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

Implement generic pool member scoring and probing

Browse Source
This commit is contained in:
fawney19
2026-05-12 01:46:22 +08:00
parent cb0ccb9cdb
commit b9e62d1667
76 changed files with 6276 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
Cargo.lock generated
View File

@@ -63,6 +63,7 @@ version = "0.1.0"
dependencies = [
"aether-ai-formats",
"aether-contracts",
"aether-data-contracts",
"aether-scheduler-core",
"async-trait",
"http",

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

@@ -43,13 +43,13 @@ pub(crate) use self::planner::{
build_local_video_sync_plan_and_reports_for_kind,
build_openai_responses_stream_plan_from_decision,
build_openai_responses_sync_plan_from_decision, build_passthrough_sync_plan_from_decision,
build_standard_family_stream_attempt_source, build_standard_family_stream_plan_and_reports,
build_standard_family_sync_attempt_source, build_standard_family_sync_plan_and_reports,
build_standard_stream_plan_from_decision, build_standard_sync_plan_from_decision,
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,
set_local_openai_chat_execution_exhausted_diagnostic,
build_provider_key_pool_score_upsert, build_standard_family_stream_attempt_source,
build_standard_family_stream_plan_and_reports, build_standard_family_sync_attempt_source,
build_standard_family_sync_plan_and_reports, build_standard_stream_plan_from_decision,
build_standard_sync_plan_from_decision, 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, set_local_openai_chat_execution_exhausted_diagnostic,
set_local_openai_image_execution_exhausted_diagnostic, CandidateFailureDiagnostic,
CandidateFailureDiagnosticKind, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot,
LocalExecutionAttemptSource, LocalResolvedOAuthRequestAuth, PlannerAppState,

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

@@ -16,6 +16,7 @@ mod materialization_policy;
mod passthrough;
mod plan_builders;
mod pool_scheduler;
pub(crate) mod pool_scores;
mod report_context;
mod route;
mod runtime_miss;
@@ -36,6 +37,7 @@ 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::route::is_matching_stream_request as planner_is_matching_stream_request;
pub(crate) use self::specialized::{
build_local_gemini_files_stream_attempt_source_for_kind,

164
apps/aether-gateway/src/ai_serving/planner/pool_scheduler.rs
View File

@@ -9,7 +9,11 @@ use aether_ai_serving::{
};
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsQuery,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
};
use aether_data_contracts::repository::pool_scores::{
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity,
PoolMemberScheduleFeedback, PoolScoreScope, StoredPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::{Map, Value};
@@ -37,6 +41,8 @@ use crate::handlers::shared::{
use crate::orchestration::LocalExecutionCandidateMetadata;
use crate::provider_key_auth::provider_key_auth_semantics;
use super::pool_scores::provider_key_pool_score_scope;
static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
const DEFAULT_POOL_KEY_PAGE_SIZE: u32 = 128;
const DEFAULT_POOL_MAX_SCANNED_KEYS: u32 = 1024;
@@ -176,6 +182,8 @@ pub(crate) struct PoolKeyCursor<'a> {
scanned_keys: u32,
page_size: u32,
max_scanned_keys: u32,
score_top_n: u32,
score_phase_loaded: bool,
skip_reason_counts: BTreeMap<&'static str, u32>,
next_pool_key_index: u32,
sticky_candidate_loaded: bool,
@@ -194,6 +202,17 @@ impl<'a> PoolKeyCursor<'a> {
request_auth_channel: Option<&str>,
) -> Self {
let pool_key_order = pool_key_candidate_order_for_group(&group);
let pool_config = pool_config_for_candidate(&group);
let score_top_n = pool_config
.as_ref()
.map(|config| config.score_top_n)
.unwrap_or(u64::from(DEFAULT_POOL_KEY_PAGE_SIZE))
.clamp(1, u64::from(u32::MAX)) as u32;
let max_scanned_keys = pool_config
.as_ref()
.map(|config| config.score_fallback_scan_limit)
.unwrap_or(u64::from(DEFAULT_POOL_MAX_SCANNED_KEYS))
.clamp(1, u64::from(u32::MAX)) as u32;
Self {
state,
group,
@@ -204,7 +223,9 @@ impl<'a> PoolKeyCursor<'a> {
next_offset: 0,
scanned_keys: 0,
page_size: DEFAULT_POOL_KEY_PAGE_SIZE,
max_scanned_keys: DEFAULT_POOL_MAX_SCANNED_KEYS,
max_scanned_keys,
score_top_n,
score_phase_loaded: false,
skip_reason_counts: BTreeMap::new(),
next_pool_key_index: 0,
sticky_candidate_loaded: false,
@@ -257,6 +278,13 @@ impl<'a> PoolKeyCursor<'a> {
}
async fn next_page_candidates(&mut self) -> Option<Vec<EligibleLocalExecutionCandidate>> {
if !self.score_phase_loaded {
self.score_phase_loaded = true;
if let Some(score_candidates) = self.next_score_candidates().await {
return Some(score_candidates);
}
}
if self.scanned_keys >= self.max_scanned_keys {
return None;
}
@@ -306,6 +334,82 @@ impl<'a> PoolKeyCursor<'a> {
Some(self.build_page_eligible_candidates(rows).await)
}
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 query = ListRankedPoolMembersQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: self.group.candidate.provider_id.clone(),
capability: scope.capability.clone(),
scope_kind: scope.scope_kind.clone(),
scope_id: scope.scope_id.clone(),
hard_states: vec![PoolMemberHardState::Available, PoolMemberHardState::Unknown],
probe_statuses: None,
offset: 0,
limit: self.score_top_n as usize,
};
let scores = match self.state.app().data.list_ranked_pool_members(&query).await {
Ok(scores) => scores,
Err(err) => {
warn!(
event_name = "pool_group_score_load_failed",
log_type = "event",
provider_id = %self.group.candidate.provider_id,
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
selected_provider_model_name = %self.group.candidate.selected_provider_model_name,
error = ?err,
"gateway pool scheduler failed to read ranked pool member scores"
);
return None;
}
};
if scores.is_empty() {
return None;
}
self.record_score_schedule_interest(&scores).await;
let key_ids = scores
.iter()
.map(|score| score.member_id.clone())
.collect::<Vec<_>>();
let rows_query = StoredPoolKeyCandidateRowsByKeyIdsQuery {
api_format: self.group.candidate.endpoint_api_format.clone(),
provider_id: self.group.candidate.provider_id.clone(),
endpoint_id: self.group.candidate.endpoint_id.clone(),
model_id: self.group.candidate.model_id.clone(),
selected_provider_model_name: self.group.candidate.selected_provider_model_name.clone(),
key_ids,
};
let rows = match self
.state
.app()
.list_pool_key_candidate_rows_for_group_key_ids(&rows_query)
.await
{
Ok(rows) => rows,
Err(err) => {
warn!(
event_name = "pool_group_score_key_load_failed",
log_type = "event",
provider_id = %self.group.candidate.provider_id,
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
selected_provider_model_name = %self.group.candidate.selected_provider_model_name,
score_count = scores.len(),
error = ?err,
"gateway pool scheduler failed to materialize ranked pool keys"
);
return None;
}
};
self.scanned_keys = self.scanned_keys.saturating_add(scores.len() as u32);
Some(self.build_page_eligible_candidates(rows).await)
}
async fn sticky_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
let pool_config = pool_config_for_candidate(&self.group)?;
if !admin_provider_pool_cache_affinity_enabled(&pool_config) {
@@ -465,6 +569,62 @@ impl<'a> PoolKeyCursor<'a> {
self.skipped_candidates.append(&mut skipped);
}
async fn record_score_schedule_interest(&self, scores: &[StoredPoolMemberScore]) {
if scores.is_empty() {
return;
}
let scheduled_at = current_unix_ms() / 1000;
let mut failed = 0usize;
for score in scores {
let identity = PoolMemberIdentity {
pool_kind: score.pool_kind.clone(),
pool_id: score.pool_id.clone(),
member_kind: score.member_kind.clone(),
member_id: score.member_id.clone(),
};
let scope = PoolScoreScope {
capability: score.capability.clone(),
scope_kind: score.scope_kind.clone(),
scope_id: score.scope_id.clone(),
};
let result = self
.state
.app()
.data
.record_pool_member_schedule_feedback(PoolMemberScheduleFeedback {
identity,
scope: Some(scope),
scheduled_at,
succeeded: None,
hard_state: None,
score_delta: None,
score_reason_patch: Some(serde_json::json!({
"last_schedule_interest": {
"provider_id": self.group.candidate.provider_id.as_str(),
"endpoint_id": self.group.candidate.endpoint_id.as_str(),
"model_id": self.group.candidate.model_id.as_str()
}
})),
})
.await;
if result.is_err() {
failed += 1;
}
}
if failed > 0 {
warn!(
event_name = "pool_group_score_interest_update_failed",
log_type = "event",
provider_id = %self.group.candidate.provider_id,
endpoint_id = %self.group.candidate.endpoint_id,
model_id = %self.group.candidate.model_id,
failed_count = failed,
score_count = scores.len(),
"gateway pool scheduler failed to record some pool score schedule interests"
);
}
}
async fn build_page_eligible_candidates(
&mut self,
rows: Vec<StoredMinimalCandidateSelectionRow>,

192
apps/aether-gateway/src/ai_serving/planner/pool_scores.rs Normal file
View File

@@ -0,0 +1,192 @@
use aether_ai_serving::{score_pool_member, PoolMemberScoreInput, POOL_SCORE_VERSION};
use aether_data_contracts::repository::pool_scores::{
PoolMemberIdentity, PoolMemberProbeStatus, PoolScoreScope, UpsertPoolMemberScore,
POOL_SCORE_SCOPE_KIND_MODEL,
};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::Value;
use crate::handlers::shared::{provider_key_health_summary, provider_key_status_snapshot_payload};
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,
) -> 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 input = provider_key_score_input(
key,
provider_type,
identity.clone(),
scope.clone(),
existing,
now_unix_secs,
);
let output = score_pool_member(&input);
UpsertPoolMemberScore {
id: provider_key_pool_score_id(&identity, &scope),
identity,
scope,
score: output.score,
hard_state: output.hard_state,
score_version: POOL_SCORE_VERSION,
score_reason: output.score_reason,
last_ranked_at: Some(now_unix_secs),
last_scheduled_at: existing.and_then(|score| score.last_scheduled_at),
last_success_at: existing.and_then(|score| score.last_success_at),
last_failure_at: existing.and_then(|score| score.last_failure_at),
failure_count: existing.map(|score| score.failure_count).unwrap_or(0),
last_probe_attempt_at: existing.and_then(|score| score.last_probe_attempt_at),
last_probe_success_at: existing.and_then(|score| score.last_probe_success_at),
last_probe_failure_at: existing.and_then(|score| score.last_probe_failure_at),
probe_failure_count: existing.map(|score| score.probe_failure_count).unwrap_or(0),
probe_status: existing
.map(|score| score.probe_status)
.unwrap_or(PoolMemberProbeStatus::Never),
updated_at: now_unix_secs,
}
}
pub(crate) fn provider_key_pool_score_scope(
api_format: &str,
model_id: Option<&str>,
) -> 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),
}
}
pub(crate) fn provider_key_pool_score_id(
identity: &PoolMemberIdentity,
scope: &PoolScoreScope,
) -> String {
let raw = format!(
"{}:{}:{}:{}:{}:{}:{}",
identity.pool_kind,
identity.pool_id,
identity.member_kind,
identity.member_id,
scope.capability,
scope.scope_kind,
scope.scope_id.as_deref().unwrap_or("*")
);
format!(
"pms-{:016x}-{:016x}",
stable_hash(raw.as_bytes()),
stable_hash(identity.member_id.as_bytes())
)
}
fn provider_key_score_input(
key: &StoredProviderCatalogKey,
provider_type: &str,
identity: PoolMemberIdentity,
scope: PoolScoreScope,
existing: Option<&aether_data_contracts::repository::pool_scores::StoredPoolMemberScore>,
now_unix_secs: u64,
) -> PoolMemberScoreInput {
let status_snapshot = provider_key_status_snapshot_payload(key, provider_type);
let quota_snapshot = status_snapshot
.as_object()
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object);
let account_snapshot = status_snapshot
.as_object()
.and_then(|snapshot| snapshot.get("account"))
.and_then(Value::as_object);
let (health_score, _, _, any_circuit_open, _) = provider_key_health_summary(key);
let health_score = key
.health_by_format
.as_ref()
.and_then(Value::as_object)
.filter(|payload| !payload.is_empty())
.map(|_| health_score);
PoolMemberScoreInput {
identity,
scope: scope.clone(),
internal_priority: key.internal_priority,
is_active: key.is_active,
health_score,
quota_usage_ratio: quota_snapshot
.and_then(|quota| quota.get("usage_ratio"))
.and_then(json_f64)
.map(|value| value.clamp(0.0, 1.0)),
quota_exhausted: quota_snapshot
.and_then(|quota| quota.get("exhausted"))
.and_then(Value::as_bool)
.unwrap_or(false),
account_blocked: account_snapshot
.and_then(|account| account.get("blocked"))
.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),
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(),
total_tokens: key.total_tokens,
total_cost_usd: key.total_cost_usd,
last_used_at: key.last_used_at_unix_secs,
last_probe_success_at: existing.and_then(|score| score.last_probe_success_at),
probe_status: existing
.map(|score| score.probe_status)
.unwrap_or(PoolMemberProbeStatus::Never),
now_unix_secs,
}
}
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
.as_str()
.map(str::trim)
.filter(|value| !value.is_empty())
.and_then(|value| value.parse::<f64>().ok())
})
}
fn api_format_lookup_keys(api_format: &str) -> Vec<String> {
let normalized = aether_ai_formats::normalize_api_format_alias(api_format);
let mut keys = aether_ai_formats::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 {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(0x100000001b3);
}
hash
}

12
apps/aether-gateway/src/control/route/admin/observability_families.rs
View File

@@ -240,6 +240,18 @@ pub(super) fn classify_admin_observability_family_route(
"admin:pool",
false,
))
} else if method == http::Method::GET
&& normalized_path_no_trailing.starts_with("/api/admin/pool/")
&& normalized_path_no_trailing.ends_with("/scores")
&& normalized_path_no_trailing.matches('/').count() == 5
{
Some(classified(
"admin_proxy",
"pool_manage",
"scores",
"admin:pool",
false,
))
} else if method == http::Method::POST
&& normalized_path_no_trailing.starts_with("/api/admin/pool/")
&& normalized_path_no_trailing.ends_with("/keys/batch-import")

8
apps/aether-gateway/src/control/tests/admin_pool.rs
View File

@@ -52,6 +52,14 @@ fn classifies_admin_pool_provider_key_routes_as_admin_proxy_route() {
assert_eq!(list.route_family.as_deref(), Some("pool_manage"));
assert_eq!(list.route_kind.as_deref(), Some("list_keys"));
let scores_uri: Uri = "/api/admin/pool/provider-1/scores?api_format=openai:responses"
.parse()
.expect("uri should parse");
let scores = classify_control_route(&http::Method::GET, &scores_uri, &headers)
.expect("route should classify");
assert_eq!(scores.route_family.as_deref(), Some("pool_manage"));
assert_eq!(scores.route_kind.as_deref(), Some("scores"));
let batch_import_uri: Uri = "/api/admin/pool/provider-1/keys/batch-import"
.parse()
.expect("uri should parse");

14
apps/aether-gateway/src/data/state/core.rs
View File

@@ -45,6 +45,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -86,6 +88,8 @@ impl GatewayDataState {
let gemini_file_mapping_writer = backends.write().gemini_file_mappings();
let provider_catalog_reader = backends.read().provider_catalog();
let provider_catalog_writer = backends.write().provider_catalog();
let pool_score_reader = backends.read().pool_scores();
let pool_score_writer = backends.write().pool_scores();
let provider_quota_reader = backends.read().provider_quotas();
let provider_quota_writer = backends.write().provider_quotas();
let usage_reader = backends.read().usage();
@@ -125,6 +129,8 @@ impl GatewayDataState {
request_candidate_writer,
provider_catalog_reader,
provider_catalog_writer,
pool_score_reader,
pool_score_writer,
provider_quota_reader,
provider_quota_writer,
usage_reader,
@@ -263,6 +269,14 @@ impl GatewayDataState {
self.provider_catalog_writer.is_some()
}
pub(crate) fn has_pool_score_reader(&self) -> bool {
self.pool_score_reader.is_some()
}
pub(crate) fn has_pool_score_writer(&self) -> bool {
self.pool_score_writer.is_some()
}
pub(crate) fn has_proxy_node_reader(&self) -> bool {
self.proxy_node_reader.is_some()
}

15
apps/aether-gateway/src/data/state/mod.rs
View File

@@ -93,7 +93,8 @@ use aether_data_contracts::repository::billing::{
};
use aether_data_contracts::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
StoredRequestedModelCandidateRowsQuery,
};
use aether_data_contracts::repository::candidates::{
PublicHealthStatusCount, PublicHealthTimelineBucket, RequestCandidateReadRepository,
@@ -107,6 +108,13 @@ use aether_data_contracts::repository::global_models::{
StoredProviderModelStats, StoredPublicCatalogModel, StoredPublicGlobalModel,
StoredPublicGlobalModelPage, UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
};
use aether_data_contracts::repository::pool_scores::{
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
UpsertPoolMemberScore,
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyListQuery, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogKeyPage,
@@ -156,6 +164,8 @@ pub(crate) struct GatewayDataState {
request_candidate_writer: Option<Arc<dyn RequestCandidateWriteRepository>>,
provider_catalog_reader: Option<Arc<dyn ProviderCatalogReadRepository>>,
provider_catalog_writer: Option<Arc<dyn ProviderCatalogWriteRepository>>,
pool_score_reader: Option<Arc<dyn PoolScoreReadRepository>>,
pool_score_writer: Option<Arc<dyn PoolMemberScoreWriteRepository>>,
provider_quota_reader: Option<Arc<dyn ProviderQuotaReadRepository>>,
provider_quota_writer: Option<Arc<dyn ProviderQuotaWriteRepository>>,
usage_reader: Option<Arc<dyn UsageReadRepository>>,
@@ -257,6 +267,8 @@ impl fmt::Debug for GatewayDataState {
"has_provider_catalog_writer",
&self.provider_catalog_writer.is_some(),
)
.field("has_pool_score_reader", &self.pool_score_reader.is_some())
.field("has_pool_score_writer", &self.pool_score_writer.is_some())
.field(
"has_provider_quota_reader",
&self.provider_quota_reader.is_some(),
@@ -287,6 +299,7 @@ mod catalog;
mod core;
mod integrations;
mod models;
mod pool_scores;
mod runtime;
#[cfg(test)]
mod testing;

13
apps/aether-gateway/src/data/state/models.rs
View File

@@ -2,7 +2,8 @@ use super::{
AdminGlobalModelListQuery, AdminProviderModelListQuery, CreateAdminGlobalModelRecord,
DataLayerError, GatewayDataState, PublicCatalogModelListQuery, PublicCatalogModelSearchQuery,
PublicGlobalModelQuery, StoredAdminGlobalModel, StoredAdminGlobalModelPage,
StoredAdminProviderModel, StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateRowsQuery,
StoredAdminProviderModel, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
StoredProviderActiveGlobalModel, StoredProviderModelStats, StoredPublicCatalogModel,
StoredPublicGlobalModel, StoredPublicGlobalModelPage, StoredRequestedModelCandidateRowsQuery,
UpdateAdminGlobalModelRecord, UpsertAdminProviderModelRecord,
@@ -73,6 +74,16 @@ impl GatewayDataState {
}
}
pub(crate) async fn list_pool_key_candidate_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
match &self.minimal_candidate_selection_reader {
Some(repository) => repository.list_pool_key_rows_for_group_key_ids(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_public_global_models(
&self,
query: &PublicGlobalModelQuery,

123
apps/aether-gateway/src/data/state/pool_scores.rs Normal file
View File

@@ -0,0 +1,123 @@
use super::{
DataLayerError, GatewayDataState, GetPoolMemberScoresByIdsQuery,
ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery, ListRankedPoolMembersQuery,
PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt, PoolMemberProbeResult,
PoolMemberScheduleFeedback, PoolScoreScope, StoredPoolMemberScore, UpsertPoolMemberScore,
};
impl GatewayDataState {
pub(crate) async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
match &self.pool_score_reader {
Some(repository) => repository.list_ranked_pool_members(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
match &self.pool_score_reader {
Some(repository) => repository.list_pool_member_probe_candidates(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
match &self.pool_score_reader {
Some(repository) => repository.list_pool_member_scores(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
match &self.pool_score_reader {
Some(repository) => repository.get_pool_member_scores_by_ids(query).await,
None => Ok(Vec::new()),
}
}
pub(crate) async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<Option<StoredPoolMemberScore>, DataLayerError> {
match &self.pool_score_writer {
Some(repository) => repository.upsert_pool_member_score(score).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
match &self.pool_score_writer {
Some(repository) => repository.record_pool_member_probe_result(result).await,
None => Ok(0),
}
}
pub(crate) async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
match &self.pool_score_writer {
Some(repository) => repository.mark_pool_member_probe_in_progress(attempt).await,
None => Ok(0),
}
}
pub(crate) async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
match &self.pool_score_writer {
Some(repository) => {
repository
.record_pool_member_schedule_feedback(feedback)
.await
}
None => Ok(0),
}
}
pub(crate) async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
match &self.pool_score_writer {
Some(repository) => {
repository
.mark_pool_member_hard_state(identity, scope, hard_state, updated_at)
.await
}
None => Ok(0),
}
}
pub(crate) async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
match &self.pool_score_writer {
Some(repository) => {
repository
.delete_pool_member_scores_for_member(identity)
.await
}
None => Ok(0),
}
}
}

4
apps/aether-gateway/src/data/state/testing/announcements.rs
View File

@@ -34,6 +34,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -86,6 +88,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,

114
apps/aether-gateway/src/data/state/testing/mod.rs
View File

@@ -2,6 +2,7 @@ use std::collections::BTreeMap;
use std::sync::{Arc, RwLock};
use aether_data_contracts::repository::candidates::RequestCandidateRepository;
use aether_data_contracts::repository::pool_scores::PoolMemberScoreRepository;
use aether_data_contracts::repository::quota::ProviderQuotaRepository;
use aether_data_contracts::repository::usage::UsageRepository;
@@ -12,12 +13,13 @@ use super::{
GeminiFileMappingWriteRepository, GlobalModelReadRepository, GlobalModelWriteRepository,
ManagementTokenReadRepository, ManagementTokenWriteRepository,
MinimalCandidateSelectionReadRepository, OAuthProviderReadRepository,
OAuthProviderWriteRepository, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, ProxyNodeReadRepository,
ProxyNodeWriteRepository, RequestCandidateReadRepository, RequestCandidateWriteRepository,
SettlementWriteRepository, StoredSystemConfigEntry, StoredUserPreferenceRecord,
UsageReadRepository, UsageWriteRepository, UserReadRepository, VideoTaskReadRepository,
VideoTaskWriteRepository, WalletReadRepository, WalletWriteRepository,
OAuthProviderWriteRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, ProviderQuotaReadRepository,
ProviderQuotaWriteRepository, ProxyNodeReadRepository, ProxyNodeWriteRepository,
RequestCandidateReadRepository, RequestCandidateWriteRepository, SettlementWriteRepository,
StoredSystemConfigEntry, StoredUserPreferenceRecord, UsageReadRepository, UsageWriteRepository,
UserReadRepository, VideoTaskReadRepository, VideoTaskWriteRepository, WalletReadRepository,
WalletWriteRepository,
};
mod announcements;
@@ -67,6 +69,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -118,6 +122,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -165,6 +171,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -300,6 +308,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -333,6 +343,18 @@ impl GatewayDataState {
self
}
#[cfg(test)]
pub(crate) fn with_pool_score_repository_for_tests<T>(mut self, repository: Arc<T>) -> Self
where
T: PoolMemberScoreRepository + 'static,
{
let pool_score_reader: Arc<dyn PoolScoreReadRepository> = repository.clone();
let pool_score_writer: Arc<dyn PoolMemberScoreWriteRepository> = repository;
self.pool_score_reader = Some(pool_score_reader);
self.pool_score_writer = Some(pool_score_writer);
self
}
#[cfg(test)]
pub(crate) fn with_provider_catalog_and_request_candidate_reader_for_tests(
provider_catalog_repository: Arc<dyn ProviderCatalogReadRepository>,
@@ -363,6 +385,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -419,6 +443,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,
@@ -484,6 +510,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,
@@ -531,6 +559,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -579,6 +609,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -638,6 +670,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -699,6 +733,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -744,6 +780,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(repository),
@@ -804,6 +842,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -857,6 +897,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -915,6 +957,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -974,6 +1018,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1032,6 +1078,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -1079,6 +1127,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1126,6 +1176,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1185,6 +1237,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1249,6 +1303,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1296,6 +1352,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1348,6 +1406,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1417,6 +1477,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1481,6 +1543,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1529,6 +1593,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1577,6 +1643,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1627,6 +1695,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_repository),
@@ -1675,6 +1745,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1723,6 +1795,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -1779,6 +1853,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,
@@ -1836,6 +1912,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,
@@ -1896,6 +1974,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,
@@ -1962,6 +2042,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -2029,6 +2111,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -2100,6 +2184,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -2178,6 +2264,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: Some(provider_catalog_writer),
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -2238,6 +2326,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,
@@ -2289,6 +2379,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -2336,6 +2428,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -2389,6 +2483,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -2446,6 +2542,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -2504,6 +2602,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: Some(usage_reader),
@@ -2555,6 +2655,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: Some(provider_quota_reader),
provider_quota_writer: Some(provider_quota_writer),
usage_reader: None,

12
apps/aether-gateway/src/data/state/testing/video_tasks.rs
View File

@@ -38,6 +38,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -92,6 +94,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -143,6 +147,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -198,6 +204,8 @@ impl GatewayDataState {
request_candidate_writer: None,
provider_catalog_reader: Some(provider_catalog_repository),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -257,6 +265,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,
@@ -325,6 +335,8 @@ impl GatewayDataState {
request_candidate_writer: Some(request_candidate_writer),
provider_catalog_reader: Some(provider_catalog_reader),
provider_catalog_writer: None,
pool_score_reader: None,
pool_score_writer: None,
provider_quota_reader: None,
provider_quota_writer: None,
usage_reader: None,

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

@@ -270,6 +270,9 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
health_policy_enabled: true,
probing_enabled: false,
probing_interval_minutes: 10,
probe_concurrency: 4,
score_top_n: 128,
score_fallback_scan_limit: 1024,
stream_timeout_threshold: 3,
stream_timeout_window_seconds: 1800,
stream_timeout_cooldown_seconds: 300,
@@ -333,6 +336,24 @@ pub(crate) fn admin_provider_pool_config_from_config_value(
.filter(|value| *value > 0)
.map(|value| value.min(1440))
.unwrap_or(10),
probe_concurrency: pool_advanced
.get("probe_concurrency")
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(64))
.unwrap_or(4),
score_top_n: pool_advanced
.get("score_top_n")
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(4096))
.unwrap_or(128),
score_fallback_scan_limit: pool_advanced
.get("score_fallback_scan_limit")
.and_then(json_u64)
.filter(|value| *value > 0)
.map(|value| value.min(50_000))
.unwrap_or(1024),
stream_timeout_threshold: pool_advanced
.get("stream_timeout_threshold")
.and_then(json_u64)
@@ -405,6 +426,9 @@ mod tests {
"health_policy_enabled": false,
"probing_enabled": true,
"probing_interval_minutes": 20,
"probe_concurrency": 6,
"score_top_n": 256,
"score_fallback_scan_limit": 2048,
"stream_timeout_threshold": 4,
"stream_timeout_window_seconds": 900,
"stream_timeout_cooldown_seconds": 180
@@ -424,6 +448,9 @@ mod tests {
assert!(!config.health_policy_enabled);
assert!(config.probing_enabled);
assert_eq!(config.probing_interval_minutes, 20);
assert_eq!(config.probe_concurrency, 6);
assert_eq!(config.score_top_n, 256);
assert_eq!(config.score_fallback_scan_limit, 2048);
assert_eq!(config.stream_timeout_threshold, 4);
assert_eq!(config.stream_timeout_window_seconds, 900);
assert_eq!(config.stream_timeout_cooldown_seconds, 180);

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

@@ -615,6 +615,9 @@ mod tests {
health_policy_enabled: true,
probing_enabled: false,
probing_interval_minutes: 10,
probe_concurrency: 4,
score_top_n: 128,
score_fallback_scan_limit: 1024,
stream_timeout_threshold: 3,
stream_timeout_window_seconds: 1800,
stream_timeout_cooldown_seconds: 300,

17
apps/aether-gateway/src/handlers/admin/provider/pool_admin/mod.rs
View File

@@ -24,6 +24,8 @@ mod read_overview;
mod read_presets;
#[path = "read_routes/resolve_selection.rs"]
mod read_resolve_selection;
#[path = "read_routes/scores.rs"]
mod read_scores;
pub(crate) mod selection;
mod support;
@@ -34,11 +36,11 @@ pub(crate) use self::batch_shared::{
AdminPoolBatchImportRequest,
};
pub(crate) use self::support::{
admin_pool_provider_id_from_path, parse_admin_pool_key_sort, parse_admin_pool_page,
parse_admin_pool_page_size, parse_admin_pool_quick_selectors, parse_admin_pool_search,
parse_admin_pool_status_filter, AdminPoolKeySort, AdminPoolKeySortDirection,
AdminPoolKeySortField, AdminPoolResolveSelectionRequest,
ADMIN_POOL_BANNED_KEY_CLEANUP_EMPTY_MESSAGE,
admin_pool_provider_id_from_path, admin_pool_provider_id_from_scores_path,
parse_admin_pool_key_sort, parse_admin_pool_page, parse_admin_pool_page_size,
parse_admin_pool_quick_selectors, parse_admin_pool_search, parse_admin_pool_status_filter,
AdminPoolKeySort, AdminPoolKeySortDirection, AdminPoolKeySortField,
AdminPoolResolveSelectionRequest, ADMIN_POOL_BANNED_KEY_CLEANUP_EMPTY_MESSAGE,
ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
ADMIN_POOL_PROVIDER_CATALOG_WRITER_UNAVAILABLE_DETAIL,
};
@@ -103,6 +105,11 @@ pub(crate) async fn maybe_build_local_admin_pool_response(
read_keys::build_admin_pool_list_keys_response(state, request_context).await?,
));
}
Some("scores") => {
return Ok(Some(
read_scores::build_admin_pool_scores_response(state, request_context).await?,
));
}
Some("resolve_selection") => {
return Ok(Some(
read_resolve_selection::build_admin_pool_resolve_selection_response(

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

@@ -0,0 +1,196 @@
use super::{
admin_pool_provider_id_from_scores_path, build_admin_pool_error_response,
parse_admin_pool_page, parse_admin_pool_page_size,
};
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
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,
};
use axum::{
body::Body,
http,
response::{IntoResponse, Response},
Json,
};
use serde_json::json;
use std::collections::BTreeMap;
pub(super) async fn build_admin_pool_scores_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
) -> Result<Response<Body>, GatewayError> {
let Some(provider_id) = admin_pool_provider_id_from_scores_path(request_context.path()) else {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
"provider_id 无效",
));
};
let query = request_context.query_string();
let page = match parse_admin_pool_page(query) {
Ok(value) => value,
Err(message) => {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
message,
));
}
};
let page_size = match parse_admin_pool_page_size(query) {
Ok(value) => value.min(500),
Err(message) => {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
message,
));
}
};
let offset = page.saturating_sub(1).saturating_mul(page_size);
let api_format = query_param_value(query, "api_format")
.map(|value| aether_ai_formats::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) => {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
message,
));
}
};
let probe_statuses = match parse_probe_status_filter(query) {
Ok(value) => value,
Err(message) => {
return Ok(build_admin_pool_error_response(
http::StatusCode::BAD_REQUEST,
message,
));
}
};
let scores = state
.app()
.data
.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(),
hard_states,
probe_statuses,
offset,
limit: page_size,
})
.await
.map_err(|err| GatewayError::Internal(format!("{err:?}")))?;
let key_ids = scores
.iter()
.map(|score| score.member_id.clone())
.collect::<Vec<_>>();
let keys = state
.app()
.read_provider_catalog_keys_by_ids(&key_ids)
.await
.unwrap_or_default()
.into_iter()
.map(|key| (key.id.clone(), key))
.collect::<BTreeMap<_, _>>();
let items = scores
.into_iter()
.map(|score| {
let key = keys.get(&score.member_id);
json!({
"id": score.id,
"pool_kind": score.pool_kind,
"pool_id": score.pool_id,
"member_kind": score.member_kind,
"member_id": score.member_id,
"capability": score.capability,
"scope_kind": score.scope_kind,
"scope_id": score.scope_id,
"score": score.score,
"hard_state": score.hard_state.as_database(),
"score_version": score.score_version,
"score_reason": score.score_reason,
"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,
"key": key.map(|key| json!({
"id": key.id,
"name": key.name,
"auth_type": key.auth_type,
"is_active": key.is_active,
"internal_priority": key.internal_priority,
"last_used_at": key.last_used_at_unix_secs,
}))
})
})
.collect::<Vec<_>>();
Ok(Json(json!({
"provider_id": provider_id,
"page": page,
"page_size": page_size,
"filters": {
"api_format": api_format,
"model_id": model_id,
"hard_state": query_param_value(query, "hard_state"),
"probe_status": query_param_value(query, "probe_status")
},
"items": items
}))
.into_response())
}
fn parse_hard_state_filter(query: Option<&str>) -> Result<Vec<PoolMemberHardState>, String> {
let Some(raw) = query_param_value(query, "hard_state") else {
return Ok(Vec::new());
};
raw.split(',')
.map(|value| match value.trim() {
"available" => Ok(PoolMemberHardState::Available),
"unknown" => Ok(PoolMemberHardState::Unknown),
"cooldown" => Ok(PoolMemberHardState::Cooldown),
"quota_exhausted" => Ok(PoolMemberHardState::QuotaExhausted),
"auth_invalid" => Ok(PoolMemberHardState::AuthInvalid),
"banned" => Ok(PoolMemberHardState::Banned),
"inactive" => Ok(PoolMemberHardState::Inactive),
_ => Err("hard_state must be one of: available, unknown, cooldown, quota_exhausted, auth_invalid, banned, inactive".to_string()),
})
.collect()
}
fn parse_probe_status_filter(
query: Option<&str>,
) -> Result<Option<Vec<PoolMemberProbeStatus>>, String> {
let Some(raw) = query_param_value(query, "probe_status") else {
return Ok(None);
};
raw.split(',')
.map(|value| match value.trim() {
"never" => Ok(PoolMemberProbeStatus::Never),
"ok" => Ok(PoolMemberProbeStatus::Ok),
"failed" => Ok(PoolMemberProbeStatus::Failed),
"stale" => Ok(PoolMemberProbeStatus::Stale),
"in_progress" => Ok(PoolMemberProbeStatus::InProgress),
_ => Err(
"probe_status must be one of: never, ok, failed, stale, in_progress".to_string(),
),
})
.collect::<Result<Vec<_>, _>>()
.map(Some)
}

16
apps/aether-gateway/src/handlers/admin/provider/pool_admin/support.rs
View File

@@ -149,6 +149,18 @@ pub(crate) fn admin_pool_provider_id_from_path(request_path: &str) -> Option<Str
}
}
pub(crate) fn admin_pool_provider_id_from_scores_path(request_path: &str) -> Option<String> {
let raw = request_path.strip_prefix("/api/admin/pool/")?;
let mut segments = raw.split('/');
let provider_id = segments.next()?.trim();
let scores_segment = segments.next()?.trim_end_matches('/').trim();
if provider_id.is_empty() || scores_segment != "scores" {
None
} else {
Some(provider_id.to_string())
}
}
pub(crate) fn is_admin_pool_route(request_context: &AdminRequestContext<'_>) -> bool {
let normalized_path = request_context.path().trim_end_matches('/');
let path = if normalized_path.is_empty() {
@@ -164,6 +176,10 @@ pub(crate) fn is_admin_pool_route(request_context: &AdminRequestContext<'_>) ->
&& path.starts_with("/api/admin/pool/")
&& path.ends_with("/keys")
&& path.matches('/').count() == 5)
|| (request_context.method() == http::Method::GET
&& path.starts_with("/api/admin/pool/")
&& path.ends_with("/scores")
&& path.matches('/').count() == 5)
|| (request_context.method() == http::Method::POST
&& path.starts_with("/api/admin/pool/")
&& path.ends_with("/keys/batch-import")

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

@@ -39,6 +39,9 @@ pub(crate) struct AdminProviderPoolConfig {
pub(crate) health_policy_enabled: bool,
pub(crate) probing_enabled: bool,
pub(crate) probing_interval_minutes: u64,
pub(crate) probe_concurrency: u64,
pub(crate) score_top_n: u64,
pub(crate) score_fallback_scan_limit: u64,
pub(crate) stream_timeout_threshold: u64,
pub(crate) stream_timeout_window_seconds: u64,
pub(crate) stream_timeout_cooldown_seconds: u64,

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

@@ -11,7 +11,7 @@ pub(crate) use runtime::{
run_admin_system_cleanup_once, skip_proxy_upgrade_rollout_node, spawn_audit_cleanup_worker,
spawn_db_maintenance_worker, spawn_gemini_file_mapping_cleanup_worker,
spawn_oauth_token_refresh_worker, spawn_pending_cleanup_worker, spawn_pool_monitor_worker,
spawn_pool_quota_probe_worker, spawn_provider_checkin_worker,
spawn_pool_quota_probe_worker, spawn_pool_score_rebuild_worker, spawn_provider_checkin_worker,
spawn_proxy_node_metrics_cleanup_worker, spawn_proxy_node_stale_cleanup_worker,
spawn_proxy_upgrade_rollout_worker, spawn_request_candidate_cleanup_worker,
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,

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

@@ -20,6 +20,8 @@ mod oauth_token_refresh;
mod pending_cleanup;
#[path = "runtime/pool_quota_probe.rs"]
mod pool_quota_probe;
#[path = "runtime/pool_score_rebuild.rs"]
mod pool_score_rebuild;
#[path = "runtime/provider_checkin.rs"]
mod provider_checkin;
#[path = "runtime/proxy_node_metrics_cleanup.rs"]
@@ -67,6 +69,10 @@ pub(crate) use pool_quota_probe::{
select_pool_quota_probe_key_ids, spawn_pool_quota_probe_worker, PoolQuotaProbeRunSummary,
PoolQuotaProbeWorkerConfig,
};
pub(crate) use pool_score_rebuild::{
perform_pool_score_rebuild_once, perform_pool_score_rebuild_once_with_config,
spawn_pool_score_rebuild_worker, PoolScoreRebuildRunSummary, PoolScoreRebuildWorkerConfig,
};
pub(crate) use provider_checkin::{perform_provider_checkin_once, ProviderCheckinRunSummary};
use proxy_node_metrics_cleanup::*;
use proxy_node_staleness::*;

371
apps/aether-gateway/src/maintenance/runtime/pool_quota_probe.rs
View File

@@ -1,10 +1,16 @@
use std::collections::BTreeMap;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::pool_scores::{
ListPoolMemberProbeCandidatesQuery, PoolMemberHardState, PoolMemberIdentity,
PoolMemberProbeAttempt, PoolMemberProbeResult, PoolMemberProbeStatus,
POOL_KIND_PROVIDER_KEY_POOL,
};
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_runtime_state::{RuntimeLockLease, RuntimeState};
use futures_util::{stream, StreamExt};
use serde_json::Value;
use tracing::{debug, info, warn};
@@ -19,6 +25,7 @@ const POOL_QUOTA_PROBE_REDIS_PREFIX: &str = "ap:quota_probe:last";
const POOL_QUOTA_PROBE_DEFAULT_SCAN_INTERVAL_SECONDS: u64 = 60;
const POOL_QUOTA_PROBE_MIN_SCAN_INTERVAL_SECONDS: u64 = 15;
const POOL_QUOTA_PROBE_DEFAULT_MAX_KEYS_PER_PROVIDER: usize = 50;
const POOL_QUOTA_PROBE_DEFAULT_GLOBAL_CONCURRENCY: usize = 16;
const POOL_QUOTA_PROBE_PROVIDER_LOCK_TTL_MS: u64 = 30_000;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -50,6 +57,7 @@ impl PoolQuotaProbeRunSummary {
pub(crate) struct PoolQuotaProbeWorkerConfig {
pub(crate) scan_interval: Duration,
pub(crate) max_keys_per_provider: usize,
pub(crate) global_concurrency: usize,
}
impl PoolQuotaProbeWorkerConfig {
@@ -63,9 +71,15 @@ impl PoolQuotaProbeWorkerConfig {
"POOL_QUOTA_PROBE_MAX_KEYS_PER_PROVIDER",
POOL_QUOTA_PROBE_DEFAULT_MAX_KEYS_PER_PROVIDER,
);
let global_concurrency = env_usize(
"POOL_QUOTA_PROBE_GLOBAL_CONCURRENCY",
POOL_QUOTA_PROBE_DEFAULT_GLOBAL_CONCURRENCY,
)
.clamp(1, 256);
Self {
scan_interval: Duration::from_secs(scan_interval_seconds),
max_keys_per_provider,
global_concurrency,
}
}
}
@@ -172,6 +186,49 @@ pub(crate) fn select_pool_quota_probe_key_ids(
stale.into_iter().map(|(_, key_id)| key_id).collect()
}
async fn select_score_probe_key_ids(
state: &AppState,
provider_id: &str,
now_ts: u64,
interval_seconds: u64,
limit: usize,
) -> Vec<String> {
if limit == 0 {
return Vec::new();
}
let stale_before_unix_secs = now_ts.saturating_sub(interval_seconds);
let query = ListPoolMemberProbeCandidatesQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: provider_id.to_string(),
capability: None,
stale_before_unix_secs,
limit: limit.saturating_mul(4).max(limit),
};
let scores = match state.data.list_pool_member_probe_candidates(&query).await {
Ok(scores) => scores,
Err(err) => {
debug!(
provider_id,
error = ?err,
"gateway pool quota probe: failed to read score probe candidates"
);
return Vec::new();
}
};
let mut selected = Vec::new();
let mut seen = std::collections::BTreeSet::new();
for score in scores {
if !seen.insert(score.member_id.clone()) {
continue;
}
selected.push(score.member_id);
if selected.len() >= limit {
break;
}
}
selected
}
fn probe_stamp_key(provider_id: &str, key_id: &str) -> String {
format!("{POOL_QUOTA_PROBE_REDIS_PREFIX}:{provider_id}:{key_id}")
}
@@ -295,7 +352,28 @@ async fn select_keys_for_provider(
let key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
let probe_stamps = load_probe_timestamps(runtime, &provider.id, &key_ids).await;
let selected_ids = select_pool_quota_probe_key_ids(
let mut selected_ids = select_score_probe_key_ids(
state,
&provider.id,
now_ts,
interval_seconds,
max_keys_per_provider,
)
.await
.into_iter()
.filter(|key_id| key_ids.iter().any(|known_id| known_id == key_id))
.filter(|key_id| {
probe_stamps.get(key_id).is_none_or(|last_probe_ts| {
now_ts.saturating_sub(*last_probe_ts) >= interval_seconds
})
})
.collect::<Vec<_>>();
let mut selected_seen = selected_ids
.iter()
.cloned()
.collect::<std::collections::BTreeSet<_>>();
let remaining = max_keys_per_provider.saturating_sub(selected_ids.len());
let fallback_selected_ids = select_pool_quota_probe_key_ids(
&keys,
provider_type,
now_ts,
@@ -303,6 +381,14 @@ async fn select_keys_for_provider(
&probe_stamps,
max_keys_per_provider,
);
for key_id in fallback_selected_ids {
if selected_ids.len() >= max_keys_per_provider || remaining == 0 {
break;
}
if selected_seen.insert(key_id.clone()) {
selected_ids.push(key_id);
}
}
if selected_ids.is_empty() {
return Ok(Vec::new());
}
@@ -381,6 +467,166 @@ fn update_summary_from_payload(
.unwrap_or(0) as usize;
}
async fn record_score_probe_results_from_payload(
state: &AppState,
provider_id: &str,
selected_key_ids: &[String],
payload: Option<&Value>,
attempted_at: u64,
) {
let mut recorded = std::collections::BTreeSet::new();
if let Some(results) = payload
.and_then(|value| value.get("results"))
.and_then(Value::as_array)
{
for item in results {
let Some(key_id) = item
.get("key_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
else {
continue;
};
recorded.insert(key_id.to_string());
record_score_probe_result_for_key(
state,
provider_id,
key_id,
attempted_at,
probe_result_succeeded(item),
probe_result_hard_state(item),
serde_json::json!({
"last_probe": {
"source": "pool_quota_probe",
"status": item.get("status").cloned().unwrap_or(Value::Null),
"status_code": item.get("status_code").cloned().unwrap_or(Value::Null),
"message": item.get("message").cloned().unwrap_or(Value::Null),
"auto_removed": item.get("auto_removed").cloned().unwrap_or(Value::Null)
}
}),
)
.await;
}
}
for key_id in selected_key_ids {
if recorded.contains(key_id) {
continue;
}
record_score_probe_result_for_key(
state,
provider_id,
key_id,
attempted_at,
false,
None,
serde_json::json!({
"last_probe": {
"source": "pool_quota_probe",
"status": "missing_result"
}
}),
)
.await;
}
}
async fn record_score_probe_result_for_key(
state: &AppState,
provider_id: &str,
key_id: &str,
attempted_at: u64,
succeeded: bool,
hard_state: Option<PoolMemberHardState>,
score_reason_patch: Value,
) {
let result = PoolMemberProbeResult {
identity: PoolMemberIdentity::provider_api_key(provider_id.to_string(), key_id.to_string()),
scope: None,
attempted_at,
succeeded,
hard_state,
probe_status: if succeeded {
PoolMemberProbeStatus::Ok
} else {
PoolMemberProbeStatus::Failed
},
score_reason_patch: Some(score_reason_patch),
};
if let Err(err) = state.data.record_pool_member_probe_result(result).await {
debug!(
provider_id,
key_id,
error = ?err,
"gateway pool quota probe: failed to record score probe result"
);
}
}
async fn record_score_probe_in_progress_for_key(
state: &AppState,
provider_id: &str,
key_id: &str,
attempted_at: u64,
) {
let attempt = PoolMemberProbeAttempt {
identity: PoolMemberIdentity::provider_api_key(provider_id.to_string(), key_id.to_string()),
scope: None,
attempted_at,
score_reason_patch: Some(serde_json::json!({
"last_probe": {
"source": "pool_quota_probe",
"status": "in_progress"
}
})),
};
if let Err(err) = state.data.mark_pool_member_probe_in_progress(attempt).await {
debug!(
provider_id,
key_id,
error = ?err,
"gateway pool quota probe: failed to mark score probe in progress"
);
}
}
fn probe_result_succeeded(item: &Value) -> bool {
item.get("status")
.and_then(Value::as_str)
.is_some_and(|status| status == "success")
}
fn probe_result_hard_state(item: &Value) -> Option<PoolMemberHardState> {
if probe_result_succeeded(item) {
return Some(PoolMemberHardState::Available);
}
if item
.get("auto_removed")
.and_then(Value::as_bool)
.unwrap_or(false)
{
return Some(PoolMemberHardState::Banned);
}
let status = item
.get("status")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
match status.as_str() {
"auth_invalid" | "forbidden" => Some(PoolMemberHardState::AuthInvalid),
"workspace_deactivated" => Some(PoolMemberHardState::Banned),
"quota_exhausted" => Some(PoolMemberHardState::QuotaExhausted),
_ => match item.get("status_code").and_then(Value::as_u64) {
Some(401 | 403) => Some(PoolMemberHardState::AuthInvalid),
Some(402) => Some(PoolMemberHardState::QuotaExhausted),
Some(429 | 500..=599) => Some(PoolMemberHardState::Cooldown),
_ => None,
},
}
}
pub(crate) async fn perform_pool_quota_probe_once_with_config(
state: &AppState,
config: PoolQuotaProbeWorkerConfig,
@@ -408,6 +654,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
provider,
provider_type,
pool_config.probing_interval_minutes,
pool_config.probe_concurrency.clamp(1, 64) as usize,
))
} else {
None
@@ -421,7 +668,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
let provider_ids = providers
.iter()
.map(|(provider, _, _)| provider.id.clone())
.map(|(provider, _, _, _)| provider.id.clone())
.collect::<Vec<_>>();
let mut endpoints_by_provider = BTreeMap::<String, Vec<StoredProviderCatalogEndpoint>>::new();
for endpoint in state
@@ -441,7 +688,7 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
..PoolQuotaProbeRunSummary::empty()
};
for (provider, provider_type, interval_minutes) in providers {
for (provider, provider_type, interval_minutes, probe_concurrency) in providers {
let endpoints = endpoints_by_provider
.remove(&provider.id)
.unwrap_or_default();
@@ -474,42 +721,98 @@ pub(crate) async fn perform_pool_quota_probe_once_with_config(
summary.providers_probed += 1;
summary.selected_keys += selected_count;
let selected_key_ids = keys.iter().map(|key| key.id.clone()).collect::<Vec<_>>();
for key_id in &selected_key_ids {
record_score_probe_in_progress_for_key(state, &provider.id, key_id, now_ts).await;
}
let provider_short_id = provider.id.chars().take(8).collect::<String>();
match refresh_provider_probe_keys(&admin_state, &provider, &endpoint, &provider_type, keys)
.await
{
Ok(payload) => {
update_summary_from_payload(&mut summary, selected_count, payload.as_ref());
let probe_success = payload
.as_ref()
.and_then(|value| value.get("success"))
.and_then(Value::as_u64)
.unwrap_or(0);
let probe_failed = payload
.as_ref()
.and_then(|value| value.get("failed"))
.and_then(Value::as_u64)
.unwrap_or(0);
info!(
provider_id = %provider_short_id,
let probe_concurrency = probe_concurrency.min(config.global_concurrency).max(1);
let probe_results = stream::iter(keys.into_iter().map(|key| {
let key_id = key.id.clone();
let admin_state = &admin_state;
let provider = &provider;
let endpoint = &endpoint;
let provider_type = provider_type.as_str();
async move {
let result = refresh_provider_probe_keys(
admin_state,
provider,
endpoint,
provider_type,
selected = selected_count,
success = probe_success,
failed = probe_failed,
"gateway pool quota probe completed"
);
vec![key],
)
.await;
(key_id, result)
}
Err(err) => {
summary.failed += selected_count;
warn!(
provider_id = %provider_short_id,
provider_type,
selected = selected_count,
error = ?err,
"gateway pool quota probe failed"
);
}))
.buffer_unordered(probe_concurrency)
.collect::<Vec<_>>()
.await;
let mut probe_success = 0usize;
let mut probe_failed = 0usize;
for (key_id, result) in probe_results {
match result {
Ok(payload) => {
update_summary_from_payload(&mut summary, 1, payload.as_ref());
probe_success += payload
.as_ref()
.and_then(|value| value.get("success"))
.and_then(Value::as_u64)
.unwrap_or(0) as usize;
probe_failed += payload
.as_ref()
.and_then(|value| value.get("failed"))
.and_then(Value::as_u64)
.unwrap_or(0) as usize;
record_score_probe_results_from_payload(
state,
&provider.id,
std::slice::from_ref(&key_id),
payload.as_ref(),
now_ts,
)
.await;
}
Err(err) => {
summary.failed += 1;
probe_failed += 1;
record_score_probe_result_for_key(
state,
&provider.id,
&key_id,
now_ts,
false,
Some(PoolMemberHardState::Cooldown),
serde_json::json!({
"last_probe": {
"source": "pool_quota_probe",
"status": "worker_error",
"message": format!("{err:?}")
}
}),
)
.await;
warn!(
provider_id = %provider_short_id,
provider_type,
key_id,
error = ?err,
"gateway pool quota probe failed"
);
}
}
}
info!(
provider_id = %provider_short_id,
provider_type,
selected = selected_count,
success = probe_success,
failed = probe_failed,
concurrency = probe_concurrency,
"gateway pool quota probe completed"
);
}
Ok(summary)

405
apps/aether-gateway/src/maintenance/runtime/pool_score_rebuild.rs Normal file
View File

@@ -0,0 +1,405 @@
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 tracing::{debug, info, warn};
use crate::admin_api::admin_provider_pool_config;
use crate::ai_serving::build_provider_key_pool_score_upsert;
use crate::{AppState, GatewayError};
const POOL_SCORE_REBUILD_DEFAULT_INTERVAL_SECONDS: u64 = 300;
const POOL_SCORE_REBUILD_MIN_INTERVAL_SECONDS: u64 = 30;
const POOL_SCORE_REBUILD_DEFAULT_MAX_UPSERTS_PER_TICK: usize = 20_000;
const POOL_SCORE_REBUILD_PROVIDER_CURSOR_KEY: &str = "ap:pool_score_rebuild:provider_cursor";
const POOL_SCORE_REBUILD_PROVIDER_OFFSET_PREFIX: &str = "ap:pool_score_rebuild:provider_offset";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct PoolScoreRebuildRunSummary {
pub(crate) providers_checked: usize,
pub(crate) providers_scored: usize,
pub(crate) keys_seen: usize,
pub(crate) scores_upserted: usize,
}
impl PoolScoreRebuildRunSummary {
const fn empty() -> Self {
Self {
providers_checked: 0,
providers_scored: 0,
keys_seen: 0,
scores_upserted: 0,
}
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct PoolScoreRebuildWorkerConfig {
pub(crate) interval: Duration,
pub(crate) max_upserts_per_tick: usize,
}
impl PoolScoreRebuildWorkerConfig {
fn from_env() -> Self {
let interval_seconds = env_u64(
"POOL_SCORE_REBUILD_INTERVAL_SECONDS",
POOL_SCORE_REBUILD_DEFAULT_INTERVAL_SECONDS,
)
.max(POOL_SCORE_REBUILD_MIN_INTERVAL_SECONDS);
let max_upserts_per_tick = env_usize(
"POOL_SCORE_REBUILD_MAX_UPSERTS_PER_TICK",
POOL_SCORE_REBUILD_DEFAULT_MAX_UPSERTS_PER_TICK,
)
.max(1);
Self {
interval: Duration::from_secs(interval_seconds),
max_upserts_per_tick,
}
}
}
fn env_u64(name: &str, default_value: u64) -> u64 {
std::env::var(name)
.ok()
.and_then(|value| value.trim().parse::<u64>().ok())
.unwrap_or(default_value)
}
fn env_usize(name: &str, default_value: usize) -> usize {
std::env::var(name)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.unwrap_or(default_value)
}
fn now_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}
async fn load_runtime_usize(state: &AppState, key: &str) -> usize {
state
.runtime_state
.kv_get(key)
.await
.ok()
.flatten()
.and_then(|value| value.trim().parse::<usize>().ok())
.unwrap_or(0)
}
async fn store_runtime_usize(state: &AppState, key: &str, value: usize) {
if let Err(err) = state
.runtime_state
.kv_set(key, value.to_string(), None)
.await
{
debug!(
key,
error = ?err,
"gateway pool score rebuild: failed to store cursor"
);
}
}
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 perform_pool_score_rebuild_once_with_config(
state: &AppState,
config: PoolScoreRebuildWorkerConfig,
) -> Result<PoolScoreRebuildRunSummary, GatewayError> {
if !state.has_provider_catalog_data_reader()
|| !state.data.has_pool_score_reader()
|| !state.data.has_pool_score_writer()
{
return Ok(PoolScoreRebuildRunSummary::empty());
}
let mut providers = state
.list_provider_catalog_providers(true)
.await?
.into_iter()
.filter(|provider| admin_provider_pool_config(provider).is_some())
.collect::<Vec<_>>();
providers.sort_by(|left, right| left.id.cmp(&right.id));
if providers.is_empty() {
return Ok(PoolScoreRebuildRunSummary::empty());
}
let provider_ids = providers
.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)
.await?
{
keys_by_provider
.entry(key.provider_id.clone())
.or_insert_with(Vec::new)
.push(key);
}
let now = now_unix_secs();
let mut summary = PoolScoreRebuildRunSummary {
providers_checked: providers.len(),
..PoolScoreRebuildRunSummary::empty()
};
let start_provider_index =
load_runtime_usize(state, POOL_SCORE_REBUILD_PROVIDER_CURSOR_KEY).await % providers.len();
let mut last_provider_index = None;
for provider_index in
(0..providers.len()).map(|offset| (start_provider_index + offset) % providers.len())
{
if summary.scores_upserted >= config.max_upserts_per_tick {
break;
}
last_provider_index = Some(provider_index);
let provider = 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?
.into_iter()
.filter(|model| model.is_available)
.collect::<Vec<_>>();
if models.is_empty() {
continue;
}
let total_combinations = endpoints
.len()
.saturating_mul(models.len())
.saturating_mul(keys.len());
if total_combinations == 0 {
continue;
}
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 remaining_budget = config
.max_upserts_per_tick
.saturating_sub(summary.scores_upserted);
let provider_budget = remaining_budget.min(total_combinations);
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 = &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,
);
build_items.push(ProviderScoreBuildItem {
endpoint_index,
model_index,
key_index,
score_id: draft.id,
});
}
if build_items.is_empty() {
store_runtime_usize(
state,
&provider_cursor_key,
(provider_cursor + provider_budget) % total_combinations,
)
.await;
continue;
}
let existing_scores = state
.data
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
ids: build_items
.iter()
.map(|item| item.score_id.clone())
.collect(),
})
.await
.unwrap_or_else(|err| {
debug!(
provider_id = %provider.id,
error = ?err,
"gateway pool score rebuild: failed to read existing scores by id"
);
Vec::new()
})
.into_iter()
.map(|score| (score.id.clone(), score))
.collect::<BTreeMap<_, _>>();
let mut provider_upserts = 0usize;
summary.keys_seen = summary.keys_seen.saturating_add(keys.len());
for item 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 upsert = build_provider_key_pool_score_upsert(
key,
provider.provider_type.as_str(),
endpoint.api_format.trim(),
Some(model.id.as_str()),
existing,
now,
);
if state
.data
.upsert_pool_member_score(upsert)
.await
.map_err(|err| GatewayError::Internal(format!("{err:?}")))?
.is_some()
{
summary.scores_upserted = summary.scores_upserted.saturating_add(1);
provider_upserts = provider_upserts.saturating_add(1);
}
}
store_runtime_usize(
state,
&provider_cursor_key,
(provider_cursor + provider_budget) % total_combinations,
)
.await;
if provider_upserts > 0 {
summary.providers_scored = summary.providers_scored.saturating_add(1);
}
}
if let Some(last_provider_index) = last_provider_index {
store_runtime_usize(
state,
POOL_SCORE_REBUILD_PROVIDER_CURSOR_KEY,
(last_provider_index + 1) % providers.len(),
)
.await;
}
Ok(summary)
}
pub(crate) async fn perform_pool_score_rebuild_once(
state: &AppState,
) -> Result<PoolScoreRebuildRunSummary, GatewayError> {
perform_pool_score_rebuild_once_with_config(state, PoolScoreRebuildWorkerConfig::from_env())
.await
}
pub(crate) fn spawn_pool_score_rebuild_worker(
state: AppState,
) -> Option<tokio::task::JoinHandle<()>> {
if !state.has_provider_catalog_data_reader()
|| !state.data.has_pool_score_reader()
|| !state.data.has_pool_score_writer()
{
return None;
}
let config = PoolScoreRebuildWorkerConfig::from_env();
Some(tokio::spawn(async move {
if let Err(err) = perform_pool_score_rebuild_once_with_config(&state, config).await {
warn!(
error = ?err,
"gateway pool score rebuild initial tick failed"
);
}
let mut interval = tokio::time::interval(config.interval);
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
loop {
interval.tick().await;
match perform_pool_score_rebuild_once_with_config(&state, config).await {
Ok(summary) if summary.scores_upserted > 0 => {
info!(
providers_checked = summary.providers_checked,
providers_scored = summary.providers_scored,
keys_seen = summary.keys_seen,
scores_upserted = summary.scores_upserted,
"gateway pool score rebuild completed"
);
}
Ok(_) => {}
Err(err) => {
warn!(
error = ?err,
"gateway pool score rebuild worker tick failed"
);
}
}
}
}))
}
#[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));
}
}

145
apps/aether-gateway/src/orchestration/effects.rs
View File

@@ -2,6 +2,9 @@ use std::collections::BTreeMap;
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
use aether_data_contracts::repository::pool_scores::{
PoolMemberHardState, PoolMemberIdentity, PoolMemberScheduleFeedback,
};
use aether_scheduler_core::{
build_scheduler_affinity_cache_key_for_api_key_id_with_client_session,
count_recent_rpm_requests_for_provider_key, ClientSessionAffinity, SchedulerAffinityTarget,
@@ -381,6 +384,19 @@ async fn record_sync_pool_success_effect(
resolve_ttfb_ms(payload.telemetry.as_ref()),
)
.await;
record_pool_score_schedule_feedback(
state,
context,
Some(true),
Some(PoolMemberHardState::Available),
Some(50),
serde_json::json!({
"last_request_feedback": {
"source": "sync_success"
}
}),
)
.await;
}
async fn record_adaptive_rate_limit_effect(
@@ -600,6 +616,19 @@ async fn record_stream_pool_success_effect(
resolve_ttfb_ms(payload.telemetry.as_ref()),
)
.await;
record_pool_score_schedule_feedback(
state,
context,
Some(true),
Some(PoolMemberHardState::Available),
Some(50),
serde_json::json!({
"last_request_feedback": {
"source": "stream_success"
}
}),
)
.await;
}
async fn record_pool_error_effect(
@@ -636,6 +665,21 @@ async fn record_pool_error_effect(
Some(effect.headers),
)
.await;
record_pool_score_schedule_feedback(
state,
context,
Some(false),
pool_score_hard_state_for_status(effect.status_code, effect.error_body),
Some(pool_score_delta_for_status(effect.status_code)),
serde_json::json!({
"last_request_feedback": {
"source": "pool_error",
"status_code": effect.status_code,
"classification": format!("{:?}", effect.classification)
}
}),
)
.await;
}
async fn open_pool_key_circuit_breaker(
@@ -730,6 +774,21 @@ async fn record_oauth_invalidation_effect(
plan.provider_id, plan.endpoint_id, plan.key_id, err
);
}
record_pool_score_schedule_feedback(
state,
context,
Some(false),
Some(PoolMemberHardState::AuthInvalid),
Some(-2_000),
serde_json::json!({
"last_request_feedback": {
"source": "oauth_invalidation",
"status_code": effect.status_code,
"reason": invalid_reason
}
}),
)
.await;
}
fn resolve_local_oauth_invalid_reason(
@@ -792,6 +851,92 @@ async fn record_pool_stream_timeout_effect(
&pool_context.pool_config,
)
.await;
record_pool_score_schedule_feedback(
state,
context,
Some(false),
Some(PoolMemberHardState::Cooldown),
Some(-250),
serde_json::json!({
"last_request_feedback": {
"source": "stream_timeout"
}
}),
)
.await;
}
async fn record_pool_score_schedule_feedback(
state: &AppState,
context: LocalExecutionEffectContext<'_>,
succeeded: Option<bool>,
hard_state: Option<PoolMemberHardState>,
score_delta: Option<i32>,
score_reason_patch: Value,
) {
if context.plan.provider_id.trim().is_empty() || context.plan.key_id.trim().is_empty() {
return;
}
let feedback = PoolMemberScheduleFeedback {
identity: PoolMemberIdentity::provider_api_key(
context.plan.provider_id.clone(),
context.plan.key_id.clone(),
),
scope: None,
scheduled_at: current_unix_secs(),
succeeded,
hard_state,
score_delta,
score_reason_patch: Some(score_reason_patch),
};
if let Err(err) = state
.data
.record_pool_member_schedule_feedback(feedback)
.await
{
warn!(
provider_id = %context.plan.provider_id,
key_id = %context.plan.key_id,
error = ?err,
"gateway orchestration effects: failed to record pool score schedule feedback"
);
}
}
fn pool_score_hard_state_for_status(
status_code: u16,
error_body: Option<&str>,
) -> Option<PoolMemberHardState> {
match status_code {
401 | 403 => Some(PoolMemberHardState::AuthInvalid),
402 => Some(PoolMemberHardState::QuotaExhausted),
429 | 500..=599 => Some(PoolMemberHardState::Cooldown),
_ => {
let body = error_body.unwrap_or_default().to_ascii_lowercase();
if body.contains("quota") && body.contains("exceed") {
Some(PoolMemberHardState::QuotaExhausted)
} else if body.contains("invalid") && body.contains("token") {
Some(PoolMemberHardState::AuthInvalid)
} else if body.contains("banned")
|| body.contains("suspended")
|| body.contains("blocked")
{
Some(PoolMemberHardState::Banned)
} else {
None
}
}
}
}
fn pool_score_delta_for_status(status_code: u16) -> i32 {
match status_code {
401 | 403 => -2_000,
402 => -1_000,
429 => -500,
500..=599 => -300,
_ => -100,
}
}
#[cfg(test)]

48
apps/aether-gateway/src/state/catalog.rs
View File

@@ -1,7 +1,8 @@
use super::{AppState, GatewayError, LocalMutationOutcome, LocalProviderDeleteTaskState};
use crate::handlers::shared::sync_provider_key_oauth_status_snapshot;
use aether_data_contracts::repository::{candidates, global_models, provider_catalog};
use aether_data_contracts::repository::{candidates, global_models, pool_scores, provider_catalog};
use std::time::{SystemTime, UNIX_EPOCH};
use tracing::warn;
impl AppState {
pub fn has_provider_catalog_data_reader(&self) -> bool {
@@ -529,6 +530,25 @@ impl AppState {
.cleanup_deleted_provider_catalog_refs(provider_id, endpoint_ids, key_ids)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
for key_id in key_ids {
if let Err(err) = self
.data
.delete_pool_member_scores_for_member(
&pool_scores::PoolMemberIdentity::provider_api_key(
provider_id.to_string(),
key_id.to_string(),
),
)
.await
{
warn!(
provider_id,
key_id,
error = ?err,
"gateway provider catalog cleanup: failed to delete pool member scores"
);
}
}
if !endpoint_ids.is_empty() || !key_ids.is_empty() {
self.clear_provider_transport_snapshot_cache();
}
@@ -620,12 +640,38 @@ impl AppState {
&self,
key_id: &str,
) -> Result<bool, GatewayError> {
let existing_key = self
.data
.list_provider_catalog_keys_by_ids(&[key_id.to_string()])
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
.into_iter()
.next();
let deleted = self
.data
.delete_provider_catalog_key(key_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if deleted {
if let Some(key) = existing_key {
if let Err(err) = self
.data
.delete_pool_member_scores_for_member(
&pool_scores::PoolMemberIdentity::provider_api_key(
key.provider_id.clone(),
key.id.clone(),
),
)
.await
{
warn!(
provider_id = %key.provider_id,
key_id = %key.id,
error = ?err,
"gateway provider catalog key delete: failed to delete pool member scores"
);
}
}
self.clear_provider_transport_snapshot_cache();
}
Ok(deleted)

5
apps/aether-gateway/src/state/core.rs
View File

@@ -47,6 +47,7 @@ use crate::maintenance::spawn_oauth_token_refresh_worker;
use crate::maintenance::spawn_pending_cleanup_worker;
use crate::maintenance::spawn_pool_monitor_worker;
use crate::maintenance::spawn_pool_quota_probe_worker;
use crate::maintenance::spawn_pool_score_rebuild_worker;
use crate::maintenance::spawn_provider_checkin_worker;
use crate::maintenance::spawn_proxy_node_metrics_cleanup_worker;
use crate::maintenance::spawn_proxy_node_stale_cleanup_worker;
@@ -1113,6 +1114,10 @@ impl AppState {
crate::task_runtime::TASK_KEY_POOL_QUOTA_PROBE,
spawn_pool_quota_probe_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_POOL_SCORE_REBUILD,
spawn_pool_score_rebuild_worker(self.clone()),
);
supervise_worker(
crate::task_runtime::TASK_KEY_STATS_HOURLY_AGG,
spawn_stats_hourly_aggregation_worker(self.data.clone()),

10
apps/aether-gateway/src/state/runtime/candidate_queries.rs
View File

@@ -57,6 +57,16 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn list_pool_key_candidate_rows_for_group_key_ids(
&self,
query: &candidate_selection::StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<candidate_selection::StoredMinimalCandidateSelectionRow>, GatewayError> {
self.data
.list_pool_key_candidate_rows_for_group_key_ids(query)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn read_provider_quota_snapshot(
&self,
provider_id: &str,

9
apps/aether-gateway/src/task_runtime/mod.rs
View File

@@ -22,6 +22,7 @@ pub(crate) const TASK_KEY_VIDEO_TASK_POLLER: &str = "video.task.poller";
pub(crate) const TASK_KEY_MODEL_FETCH_WORKER: &str = "model.fetch.worker";
pub(crate) const TASK_KEY_PROVIDER_QUOTA_RESET: &str = "provider.quota.reset.worker";
pub(crate) const TASK_KEY_POOL_QUOTA_PROBE: &str = "pool.quota.probe.worker";
pub(crate) const TASK_KEY_POOL_SCORE_REBUILD: &str = "pool.score.rebuild.worker";
pub(crate) const TASK_KEY_POOL_MONITOR: &str = "pool.monitor.worker";
pub(crate) const TASK_KEY_AUDIT_CLEANUP: &str = "maintenance.audit.cleanup";
pub(crate) const TASK_KEY_DB_MAINTENANCE: &str = "maintenance.database";
@@ -101,6 +102,14 @@ const TASK_DEFINITIONS: &[TaskDefinition] = &[
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_POOL_SCORE_REBUILD,
TaskKind::Scheduled,
"interval",
true,
true,
RETRY_ONCE,
),
TaskDefinition::new(
TASK_KEY_POOL_MONITOR,
TaskKind::Scheduled,

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

@@ -1,8 +1,13 @@
use std::sync::{Arc, Mutex};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
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,
};
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use axum::body::{to_bytes, Body, Bytes};
@@ -343,6 +348,94 @@ async fn gateway_handles_admin_pool_batch_import_locally_with_trusted_admin_prin
upstream_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_pool_scores_locally_with_trusted_admin_principal() {
let provider = sample_provider("provider-openai", "openai", 10).with_transport_fields(
true,
false,
true,
None,
None,
None,
None,
None,
Some(json!({
"pool_advanced": {
"enabled": true
}
})),
);
let mut key = sample_key("key-openai-a", "provider-openai", "openai:chat", "sk-a");
key.name = "score key".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
Vec::new(),
vec![key.clone()],
));
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()),
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 gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(Arc::clone(
&provider_catalog_repository,
))
.with_pool_score_repository_for_tests(Arc::clone(&pool_score_repository)),
),
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let response = reqwest::Client::new()
.get(format!(
"{gateway_url}/api/admin/pool/provider-openai/scores?api_format=openai:chat&model_id=model-1"
))
.header(crate::constants::GATEWAY_HEADER, "rust-phase3b")
.header(TRUSTED_ADMIN_USER_ID_HEADER, "admin-user-123")
.header(TRUSTED_ADMIN_USER_ROLE_HEADER, "admin")
.header(TRUSTED_ADMIN_SESSION_ID_HEADER, "session-123")
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(
payload["items"].as_array().map(|items| items.len()),
Some(1)
);
assert_eq!(payload["items"][0]["member_id"], json!("key-openai-a"));
assert_eq!(payload["items"][0]["key"]["name"], json!("score key"));
assert_eq!(payload["items"][0]["probe_status"], json!("ok"));
gateway_handle.abort();
}
#[tokio::test]
async fn gateway_handles_admin_pool_trailing_slash_routes_locally_with_trusted_admin_principal() {
let upstream_hits = Arc::new(Mutex::new(0usize));

1
crates/aether-ai-serving/Cargo.toml
View File

@@ -9,6 +9,7 @@ description = "AI serving application contracts and ports for Aether"
[dependencies]
aether-ai-formats.workspace = true
aether-contracts.workspace = true
aether-data-contracts.workspace = true
aether-scheduler-core.workspace = true
async-trait.workspace = true
http.workspace = true

5
crates/aether-ai-serving/src/lib.rs
View File

@@ -16,6 +16,7 @@ pub mod execution_path;
pub mod failure_diagnostic;
pub mod plan_payload;
pub mod pool_scheduler;
pub mod pool_scores;
pub mod ports;
pub mod ranking_metadata;
pub mod report_context;
@@ -99,6 +100,10 @@ pub use pool_scheduler::{
AI_POOL_ACCOUNT_EXHAUSTED_SKIP_REASON, AI_POOL_COOLDOWN_SKIP_REASON,
AI_POOL_COST_LIMIT_REACHED_SKIP_REASON,
};
pub use pool_scores::{
probe_freshness_score, score_pool_member, PoolMemberScoreInput, PoolMemberScoreOutput,
POOL_SCORE_VERSION, PROBE_FRESHNESS_TTL_SECONDS,
};
pub use ranking_metadata::append_ai_ranking_metadata_to_object;
pub use report_context::{
build_ai_execution_report_context, build_ai_report_context_original_request_echo,

237
crates/aether-ai-serving/src/pool_scores.rs Normal file
View File

@@ -0,0 +1,237 @@
use aether_data_contracts::repository::pool_scores::{
PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeStatus, PoolScoreScope,
};
use serde_json::{json, Value};
pub const POOL_SCORE_VERSION: u64 = 1;
pub const PROBE_FRESHNESS_TTL_SECONDS: u64 = 30 * 60;
#[derive(Debug, Clone)]
pub struct PoolMemberScoreInput {
pub identity: PoolMemberIdentity,
pub scope: PoolScoreScope,
pub internal_priority: i32,
pub is_active: bool,
pub health_score: Option<f64>,
pub quota_usage_ratio: Option<f64>,
pub quota_exhausted: bool,
pub account_blocked: bool,
pub oauth_invalid_reason: Option<String>,
pub circuit_open: bool,
pub success_count: u64,
pub error_count: u64,
pub total_response_time_ms: u64,
pub total_tokens: u64,
pub total_cost_usd: f64,
pub last_used_at: Option<u64>,
pub last_probe_success_at: Option<u64>,
pub probe_status: PoolMemberProbeStatus,
pub now_unix_secs: u64,
}
#[derive(Debug, Clone, PartialEq)]
pub struct PoolMemberScoreOutput {
pub score: f64,
pub hard_state: PoolMemberHardState,
pub score_reason: Value,
}
pub fn score_pool_member(input: &PoolMemberScoreInput) -> PoolMemberScoreOutput {
let hard_state = derive_hard_state(input);
let manual_priority = manual_priority_score(input.internal_priority);
let health = input.health_score.unwrap_or(0.5).clamp(0.0, 1.0);
let probe_freshness = probe_freshness_score(
input.last_probe_success_at,
input.probe_status,
input.now_unix_secs,
);
let quota_remaining = input
.quota_usage_ratio
.map(|ratio| 1.0 - ratio.clamp(0.0, 1.0))
.unwrap_or(0.5);
let latency = latency_score(input.success_count, input.total_response_time_ms);
let cost_lru = cost_lru_score(input.total_cost_usd, input.total_tokens, input.last_used_at);
let mut score = manual_priority * 0.30
+ health * 0.20
+ probe_freshness * 0.15
+ quota_remaining * 0.15
+ latency * 0.10
+ cost_lru * 0.10;
if !hard_state.schedulable() {
score = score.min(0.05);
}
score = score.clamp(0.0, 1.0);
PoolMemberScoreOutput {
score,
hard_state,
score_reason: json!({
"weights": {
"manual_priority": 0.30,
"health": 0.20,
"probe_freshness": 0.15,
"quota_remaining": 0.15,
"latency": 0.10,
"cost_lru": 0.10
},
"factors": {
"manual_priority": manual_priority,
"health": health,
"probe_freshness": probe_freshness,
"quota_remaining": quota_remaining,
"latency": latency,
"cost_lru": cost_lru
},
"hard_state": hard_state.as_database(),
"score_version": POOL_SCORE_VERSION
}),
}
}
fn derive_hard_state(input: &PoolMemberScoreInput) -> PoolMemberHardState {
if !input.is_active {
return PoolMemberHardState::Inactive;
}
if let Some(reason) = input.oauth_invalid_reason.as_deref() {
let reason = reason.to_ascii_lowercase();
if reason.contains("ban") || reason.contains("blocked") || reason.contains("suspended") {
return PoolMemberHardState::Banned;
}
return PoolMemberHardState::AuthInvalid;
}
if input.account_blocked {
return PoolMemberHardState::Banned;
}
if input.quota_exhausted {
return PoolMemberHardState::QuotaExhausted;
}
if input.circuit_open {
return PoolMemberHardState::Cooldown;
}
if input.health_score.is_some() || input.probe_status == PoolMemberProbeStatus::Ok {
PoolMemberHardState::Available
} else {
PoolMemberHardState::Unknown
}
}
fn manual_priority_score(internal_priority: i32) -> f64 {
(1.0 - (f64::from(internal_priority).clamp(0.0, 100.0) / 100.0)).clamp(0.0, 1.0)
}
pub fn probe_freshness_score(
last_probe_success_at: Option<u64>,
probe_status: PoolMemberProbeStatus,
now_unix_secs: u64,
) -> f64 {
if probe_status != PoolMemberProbeStatus::Ok {
return 0.0;
}
let Some(success_at) = last_probe_success_at else {
return 0.0;
};
let age = now_unix_secs.saturating_sub(success_at);
if age >= PROBE_FRESHNESS_TTL_SECONDS {
0.0
} else {
1.0 - (age as f64 / PROBE_FRESHNESS_TTL_SECONDS as f64)
}
}
fn latency_score(success_count: u64, total_response_time_ms: u64) -> f64 {
if success_count == 0 || total_response_time_ms == 0 {
return 0.5;
}
let avg = total_response_time_ms as f64 / success_count as f64;
if avg <= 500.0 {
1.0
} else if avg >= 60_000.0 {
0.0
} else {
1.0 - ((avg - 500.0) / 59_500.0)
}
}
fn cost_lru_score(total_cost_usd: f64, total_tokens: u64, last_used_at: Option<u64>) -> f64 {
let cost_penalty = if total_cost_usd.is_finite() {
(total_cost_usd.max(0.0) / 100.0).min(0.5)
} else {
0.0
};
let token_penalty = (total_tokens as f64 / 10_000_000.0).min(0.25);
let lru_bonus = if last_used_at.unwrap_or(0) == 0 {
0.25
} else {
0.0
};
(0.75 - cost_penalty - token_penalty + lru_bonus).clamp(0.0, 1.0)
}
#[cfg(test)]
mod tests {
use super::*;
use aether_data_contracts::repository::pool_scores::{
POOL_KIND_PROVIDER_KEY_POOL, POOL_MEMBER_KIND_PROVIDER_API_KEY, POOL_SCORE_SCOPE_KIND_MODEL,
};
fn input() -> PoolMemberScoreInput {
PoolMemberScoreInput {
identity: PoolMemberIdentity {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-1".to_string(),
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
member_id: "key-1".to_string(),
},
scope: PoolScoreScope {
capability: "openai:responses".to_string(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: Some("model-1".to_string()),
},
internal_priority: 10,
is_active: true,
health_score: Some(1.0),
quota_usage_ratio: Some(0.1),
quota_exhausted: false,
account_blocked: false,
oauth_invalid_reason: None,
circuit_open: false,
success_count: 10,
error_count: 0,
total_response_time_ms: 2_000,
total_tokens: 10,
total_cost_usd: 0.01,
last_used_at: None,
last_probe_success_at: Some(1_000),
probe_status: PoolMemberProbeStatus::Ok,
now_unix_secs: 1_000,
}
}
#[test]
fn hard_state_caps_unavailable_member_score() {
let mut input = input();
input.oauth_invalid_reason = Some("token invalid".to_string());
let output = score_pool_member(&input);
assert_eq!(output.hard_state, PoolMemberHardState::AuthInvalid);
assert!(output.score <= 0.05);
}
#[test]
fn probe_freshness_has_ttl() {
assert_eq!(
probe_freshness_score(Some(1_000), PoolMemberProbeStatus::Ok, 1_000),
1.0
);
assert_eq!(
probe_freshness_score(
Some(1_000),
PoolMemberProbeStatus::Ok,
1_000 + PROBE_FRESHNESS_TTL_SECONDS
),
0.0
);
}
}

4
crates/aether-data-contracts/src/repository/candidate_selection/mod.rs
View File

@@ -3,6 +3,6 @@ mod types;
pub use types::{
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery,
};

15
crates/aether-data-contracts/src/repository/candidate_selection/types.rs
View File

@@ -67,6 +67,16 @@ pub struct StoredPoolKeyCandidateRowsQuery {
pub limit: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredPoolKeyCandidateRowsByKeyIdsQuery {
pub api_format: String,
pub provider_id: String,
pub endpoint_id: String,
pub model_id: String,
pub selected_provider_model_name: String,
pub key_ids: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct StoredRequestedModelCandidateRowsQuery {
pub api_format: String,
@@ -124,6 +134,11 @@ pub trait MinimalCandidateSelectionReadRepository: Send + Sync {
&self,
query: &StoredPoolKeyCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, crate::DataLayerError>;
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, crate::DataLayerError>;
}
pub trait MinimalCandidateSelectionRepository:

1
crates/aether-data-contracts/src/repository/mod.rs
View File

@@ -3,6 +3,7 @@ pub mod billing;
pub mod candidate_selection;
pub mod candidates;
pub mod global_models;
pub mod pool_scores;
pub mod provider_catalog;
pub mod quota;
pub mod settlement;

11
crates/aether-data-contracts/src/repository/pool_scores/mod.rs Normal file
View File

@@ -0,0 +1,11 @@
mod types;
pub use types::{
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
PoolScoreScope, StoredPoolMemberScore, UpsertPoolMemberScore, POOL_KIND_PROVIDER_KEY_POOL,
POOL_MEMBER_KIND_PROVIDER_API_KEY, POOL_SCORE_CAPABILITY_API_FORMAT,
POOL_SCORE_SCOPE_KIND_MODEL,
};

359
crates/aether-data-contracts/src/repository/pool_scores/types.rs Normal file
View File

@@ -0,0 +1,359 @@
use async_trait::async_trait;
pub const POOL_KIND_PROVIDER_KEY_POOL: &str = "provider_key_pool";
pub const POOL_MEMBER_KIND_PROVIDER_API_KEY: &str = "provider_api_key";
pub const POOL_SCORE_CAPABILITY_API_FORMAT: &str = "api_format";
pub const POOL_SCORE_SCOPE_KIND_MODEL: &str = "model";
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
#[serde(rename_all = "snake_case")]
pub enum PoolMemberHardState {
Available,
Unknown,
Cooldown,
QuotaExhausted,
AuthInvalid,
Banned,
Inactive,
}
impl PoolMemberHardState {
pub fn as_database(self) -> &'static str {
match self {
Self::Available => "available",
Self::Unknown => "unknown",
Self::Cooldown => "cooldown",
Self::QuotaExhausted => "quota_exhausted",
Self::AuthInvalid => "auth_invalid",
Self::Banned => "banned",
Self::Inactive => "inactive",
}
}
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
match value.trim().to_ascii_lowercase().as_str() {
"available" => Ok(Self::Available),
"unknown" => Ok(Self::Unknown),
"cooldown" => Ok(Self::Cooldown),
"quota_exhausted" => Ok(Self::QuotaExhausted),
"auth_invalid" => Ok(Self::AuthInvalid),
"banned" => Ok(Self::Banned),
"inactive" => Ok(Self::Inactive),
other => Err(crate::DataLayerError::UnexpectedValue(format!(
"unknown pool member hard_state: {other}"
))),
}
}
pub fn schedulable(self) -> bool {
matches!(self, Self::Available | Self::Unknown)
}
}
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
#[serde(rename_all = "snake_case")]
pub enum PoolMemberProbeStatus {
Never,
Ok,
Failed,
Stale,
InProgress,
}
impl PoolMemberProbeStatus {
pub fn as_database(self) -> &'static str {
match self {
Self::Never => "never",
Self::Ok => "ok",
Self::Failed => "failed",
Self::Stale => "stale",
Self::InProgress => "in_progress",
}
}
pub fn from_database(value: &str) -> Result<Self, crate::DataLayerError> {
match value.trim().to_ascii_lowercase().as_str() {
"never" => Ok(Self::Never),
"ok" => Ok(Self::Ok),
"failed" => Ok(Self::Failed),
"stale" => Ok(Self::Stale),
"in_progress" => Ok(Self::InProgress),
other => Err(crate::DataLayerError::UnexpectedValue(format!(
"unknown pool member probe_status: {other}"
))),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PoolScoreScope {
pub capability: String,
pub scope_kind: String,
pub scope_id: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PoolMemberIdentity {
pub pool_kind: String,
pub pool_id: String,
pub member_kind: String,
pub member_id: String,
}
impl PoolMemberIdentity {
pub fn provider_api_key(provider_id: impl Into<String>, key_id: impl Into<String>) -> Self {
Self {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: provider_id.into(),
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
member_id: key_id.into(),
}
}
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredPoolMemberScore {
pub id: String,
pub pool_kind: String,
pub pool_id: String,
pub member_kind: String,
pub member_id: String,
pub capability: String,
pub scope_kind: String,
pub scope_id: Option<String>,
pub score: f64,
pub hard_state: PoolMemberHardState,
pub score_version: u64,
pub score_reason: serde_json::Value,
pub last_ranked_at: Option<u64>,
pub last_scheduled_at: Option<u64>,
pub last_success_at: Option<u64>,
pub last_failure_at: Option<u64>,
pub failure_count: u64,
pub last_probe_attempt_at: Option<u64>,
pub last_probe_success_at: Option<u64>,
pub last_probe_failure_at: Option<u64>,
pub probe_failure_count: u64,
pub probe_status: PoolMemberProbeStatus,
pub updated_at: u64,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct UpsertPoolMemberScore {
pub id: String,
pub identity: PoolMemberIdentity,
pub scope: PoolScoreScope,
pub score: f64,
pub hard_state: PoolMemberHardState,
pub score_version: u64,
pub score_reason: serde_json::Value,
pub last_ranked_at: Option<u64>,
pub last_scheduled_at: Option<u64>,
pub last_success_at: Option<u64>,
pub last_failure_at: Option<u64>,
pub failure_count: u64,
pub last_probe_attempt_at: Option<u64>,
pub last_probe_success_at: Option<u64>,
pub last_probe_failure_at: Option<u64>,
pub probe_failure_count: u64,
pub probe_status: PoolMemberProbeStatus,
pub updated_at: u64,
}
impl UpsertPoolMemberScore {
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
validate_non_empty(&self.id, "pool_member_scores.id")?;
validate_non_empty(&self.identity.pool_kind, "pool_member_scores.pool_kind")?;
validate_non_empty(&self.identity.pool_id, "pool_member_scores.pool_id")?;
validate_non_empty(&self.identity.member_kind, "pool_member_scores.member_kind")?;
validate_non_empty(&self.identity.member_id, "pool_member_scores.member_id")?;
validate_non_empty(&self.scope.capability, "pool_member_scores.capability")?;
validate_non_empty(&self.scope.scope_kind, "pool_member_scores.scope_kind")?;
if !self.score.is_finite() {
return Err(crate::DataLayerError::InvalidInput(
"pool_member_scores.score must be finite".to_string(),
));
}
Ok(())
}
pub fn into_stored(self) -> StoredPoolMemberScore {
StoredPoolMemberScore {
id: self.id,
pool_kind: self.identity.pool_kind,
pool_id: self.identity.pool_id,
member_kind: self.identity.member_kind,
member_id: self.identity.member_id,
capability: self.scope.capability,
scope_kind: self.scope.scope_kind,
scope_id: self.scope.scope_id,
score: self.score,
hard_state: self.hard_state,
score_version: self.score_version,
score_reason: self.score_reason,
last_ranked_at: self.last_ranked_at,
last_scheduled_at: self.last_scheduled_at,
last_success_at: self.last_success_at,
last_failure_at: self.last_failure_at,
failure_count: self.failure_count,
last_probe_attempt_at: self.last_probe_attempt_at,
last_probe_success_at: self.last_probe_success_at,
last_probe_failure_at: self.last_probe_failure_at,
probe_failure_count: self.probe_failure_count,
probe_status: self.probe_status,
updated_at: self.updated_at,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ListRankedPoolMembersQuery {
pub pool_kind: String,
pub pool_id: String,
pub capability: String,
pub scope_kind: String,
pub scope_id: Option<String>,
pub hard_states: Vec<PoolMemberHardState>,
pub probe_statuses: Option<Vec<PoolMemberProbeStatus>>,
pub offset: usize,
pub limit: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ListPoolMemberScoresQuery {
pub pool_kind: String,
pub pool_id: String,
pub capability: Option<String>,
pub scope_kind: Option<String>,
pub scope_id: Option<String>,
pub hard_states: Vec<PoolMemberHardState>,
pub probe_statuses: Option<Vec<PoolMemberProbeStatus>>,
pub offset: usize,
pub limit: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct ListPoolMemberProbeCandidatesQuery {
pub pool_kind: String,
pub pool_id: String,
pub capability: Option<String>,
pub stale_before_unix_secs: u64,
pub limit: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct GetPoolMemberScoresByIdsQuery {
pub ids: Vec<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PoolMemberProbeResult {
pub identity: PoolMemberIdentity,
pub scope: Option<PoolScoreScope>,
pub attempted_at: u64,
pub succeeded: bool,
pub hard_state: Option<PoolMemberHardState>,
pub probe_status: PoolMemberProbeStatus,
pub score_reason_patch: Option<serde_json::Value>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct PoolMemberProbeAttempt {
pub identity: PoolMemberIdentity,
pub scope: Option<PoolScoreScope>,
pub attempted_at: u64,
pub score_reason_patch: Option<serde_json::Value>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct PoolMemberScheduleFeedback {
pub identity: PoolMemberIdentity,
pub scope: Option<PoolScoreScope>,
pub scheduled_at: u64,
pub succeeded: Option<bool>,
pub hard_state: Option<PoolMemberHardState>,
pub score_delta: Option<i32>,
pub score_reason_patch: Option<serde_json::Value>,
}
#[async_trait]
pub trait PoolScoreReadRepository: Send + Sync {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, crate::DataLayerError>;
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, crate::DataLayerError>;
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, crate::DataLayerError>;
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, crate::DataLayerError>;
}
#[async_trait]
pub trait PoolMemberScoreWriteRepository: Send + Sync {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, crate::DataLayerError>;
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, crate::DataLayerError>;
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, crate::DataLayerError>;
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, crate::DataLayerError>;
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, crate::DataLayerError>;
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, crate::DataLayerError>;
}
pub trait PoolMemberScoreRepository:
PoolScoreReadRepository + PoolMemberScoreWriteRepository + Send + Sync
{
}
impl<T> PoolMemberScoreRepository for T where
T: PoolScoreReadRepository + PoolMemberScoreWriteRepository + Send + Sync
{
}
fn validate_non_empty(value: &str, field: &str) -> Result<(), crate::DataLayerError> {
if value.trim().is_empty() {
return Err(crate::DataLayerError::InvalidInput(format!(
"{field} is empty"
)));
}
Ok(())
}

30
crates/aether-data/migrations/mysql/20260403000000_baseline.sql
View File

@@ -294,6 +294,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
KEY provider_api_keys_provider_id_idx (provider_id)
);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id VARCHAR(192) PRIMARY KEY,
pool_kind VARCHAR(64) NOT NULL,
pool_id VARCHAR(64) NOT NULL,
member_kind VARCHAR(64) NOT NULL,
member_id VARCHAR(64) NOT NULL,
capability VARCHAR(64) NOT NULL,
scope_kind VARCHAR(64) NOT NULL,
scope_id VARCHAR(128),
score DOUBLE NOT NULL DEFAULT 0,
hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown',
score_version BIGINT NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at BIGINT,
last_scheduled_at BIGINT,
last_success_at BIGINT,
last_failure_at BIGINT,
failure_count BIGINT NOT NULL DEFAULT 0,
last_probe_attempt_at BIGINT,
last_probe_success_at BIGINT,
last_probe_failure_at BIGINT,
probe_failure_count BIGINT NOT NULL DEFAULT 0,
probe_status VARCHAR(64) NOT NULL DEFAULT 'never',
updated_at BIGINT NOT NULL,
KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC),
KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id),
KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at),
KEY pool_member_scores_updated_at_idx (updated_at)
);
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
id VARCHAR(64) PRIMARY KEY,
file_name VARCHAR(512) NOT NULL,

29
crates/aether-data/migrations/mysql/20260511130000_add_pool_member_scores.sql Normal file
View File

@@ -0,0 +1,29 @@
CREATE TABLE IF NOT EXISTS pool_member_scores (
id VARCHAR(192) PRIMARY KEY,
pool_kind VARCHAR(64) NOT NULL,
pool_id VARCHAR(64) NOT NULL,
member_kind VARCHAR(64) NOT NULL,
member_id VARCHAR(64) NOT NULL,
capability VARCHAR(64) NOT NULL,
scope_kind VARCHAR(64) NOT NULL,
scope_id VARCHAR(128),
score DOUBLE NOT NULL DEFAULT 0,
hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown',
score_version BIGINT NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at BIGINT,
last_scheduled_at BIGINT,
last_success_at BIGINT,
last_failure_at BIGINT,
failure_count BIGINT NOT NULL DEFAULT 0,
last_probe_attempt_at BIGINT,
last_probe_success_at BIGINT,
last_probe_failure_at BIGINT,
probe_failure_count BIGINT NOT NULL DEFAULT 0,
probe_status VARCHAR(64) NOT NULL DEFAULT 'never',
updated_at BIGINT NOT NULL,
KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC),
KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id),
KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at),
KEY pool_member_scores_updated_at_idx (updated_at)
);

79
crates/aether-data/migrations/postgres/20260403000000_baseline.sql
View File

@@ -534,6 +534,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
--
-- Name: pool_member_scores; Type: TABLE; Schema: public; Owner: -
--
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision DEFAULT 0 NOT NULL,
hard_state character varying(64) DEFAULT 'unknown'::character varying NOT NULL,
score_version bigint DEFAULT 1 NOT NULL,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint DEFAULT 0 NOT NULL,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint DEFAULT 0 NOT NULL,
probe_status character varying(64) DEFAULT 'never'::character varying NOT NULL,
updated_at bigint NOT NULL
);
--
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
--
@@ -1659,6 +1691,21 @@ END $mig$;
--
-- Name: pool_member_scores pool_member_scores_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.pool_member_scores
ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--
@@ -3401,6 +3448,38 @@ CREATE INDEX IF NOT EXISTS ix_provider_api_keys_id ON public.provider_api_keys U
--
-- Name: pool_member_scores_rank_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
--
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
--
-- Name: pool_member_scores_probe_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
--
-- Name: pool_member_scores_updated_at_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
--
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
--

40
crates/aether-data/migrations/postgres/20260511130000_add_pool_member_scores.sql Normal file
View File

@@ -0,0 +1,40 @@
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) PRIMARY KEY,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision NOT NULL DEFAULT 0,
hard_state character varying(64) NOT NULL DEFAULT 'unknown',
score_version bigint NOT NULL DEFAULT 1,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint NOT NULL DEFAULT 0,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint NOT NULL DEFAULT 0,
probe_status character varying(64) NOT NULL DEFAULT 'never',
updated_at bigint NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx
ON public.pool_member_scores USING btree
(pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx
ON public.pool_member_scores USING btree
(pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx
ON public.pool_member_scores USING btree
(pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx
ON public.pool_member_scores USING btree (updated_at);

30
crates/aether-data/migrations/sqlite/20260403000000_baseline.sql
View File

@@ -301,6 +301,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY,
pool_kind TEXT NOT NULL,
pool_id TEXT NOT NULL,
member_kind TEXT NOT NULL,
member_id TEXT NOT NULL,
capability TEXT NOT NULL,
scope_kind TEXT NOT NULL,
scope_id TEXT,
score REAL NOT NULL DEFAULT 0,
hard_state TEXT NOT NULL DEFAULT 'unknown',
score_version INTEGER NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at INTEGER,
last_scheduled_at INTEGER,
last_success_at INTEGER,
last_failure_at INTEGER,
failure_count INTEGER NOT NULL DEFAULT 0,
last_probe_attempt_at INTEGER,
last_probe_success_at INTEGER,
last_probe_failure_at INTEGER,
probe_failure_count INTEGER NOT NULL DEFAULT 0,
probe_status TEXT NOT NULL DEFAULT 'never',
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at);
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
id TEXT PRIMARY KEY,
file_name TEXT NOT NULL UNIQUE,

37
crates/aether-data/migrations/sqlite/20260511130000_add_pool_member_scores.sql Normal file
View File

@@ -0,0 +1,37 @@
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY,
pool_kind TEXT NOT NULL,
pool_id TEXT NOT NULL,
member_kind TEXT NOT NULL,
member_id TEXT NOT NULL,
capability TEXT NOT NULL,
scope_kind TEXT NOT NULL,
scope_id TEXT,
score REAL NOT NULL DEFAULT 0,
hard_state TEXT NOT NULL DEFAULT 'unknown',
score_version INTEGER NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at INTEGER,
last_scheduled_at INTEGER,
last_success_at INTEGER,
last_failure_at INTEGER,
failure_count INTEGER NOT NULL DEFAULT 0,
last_probe_attempt_at INTEGER,
last_probe_success_at INTEGER,
last_probe_failure_at INTEGER,
probe_failure_count INTEGER NOT NULL DEFAULT 0,
probe_status TEXT NOT NULL DEFAULT 'never',
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx
ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx
ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx
ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx
ON pool_member_scores (updated_at);

32
crates/aether-data/schema/bootstrap/postgres/001_types_and_tables.sql
View File

@@ -537,6 +537,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
--
-- Name: pool_member_scores; Type: TABLE; Schema: public; Owner: -
--
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision DEFAULT 0 NOT NULL,
hard_state character varying(64) DEFAULT 'unknown'::character varying NOT NULL,
score_version bigint DEFAULT 1 NOT NULL,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint DEFAULT 0 NOT NULL,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint DEFAULT 0 NOT NULL,
probe_status character varying(64) DEFAULT 'never'::character varying NOT NULL,
updated_at bigint NOT NULL
);
--
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
--

15
crates/aether-data/schema/bootstrap/postgres/003_constraints.sql
View File

@@ -252,6 +252,21 @@ END $mig$;
--
-- Name: pool_member_scores pool_member_scores_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.pool_member_scores
ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--

32
crates/aether-data/schema/bootstrap/postgres/004_indexes.sql
View File

@@ -877,6 +877,38 @@ CREATE INDEX IF NOT EXISTS ix_provider_api_keys_id ON public.provider_api_keys U
--
-- Name: pool_member_scores_rank_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
--
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
--
-- Name: pool_member_scores_probe_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
--
-- Name: pool_member_scores_updated_at_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
--
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
--

30
crates/aether-data/schema/drivers/mysql/baseline/002_provider_catalog.sql
View File

@@ -127,6 +127,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
KEY provider_api_keys_provider_id_idx (provider_id)
);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id VARCHAR(192) PRIMARY KEY,
pool_kind VARCHAR(64) NOT NULL,
pool_id VARCHAR(64) NOT NULL,
member_kind VARCHAR(64) NOT NULL,
member_id VARCHAR(64) NOT NULL,
capability VARCHAR(64) NOT NULL,
scope_kind VARCHAR(64) NOT NULL,
scope_id VARCHAR(128),
score DOUBLE NOT NULL DEFAULT 0,
hard_state VARCHAR(64) NOT NULL DEFAULT 'unknown',
score_version BIGINT NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at BIGINT,
last_scheduled_at BIGINT,
last_success_at BIGINT,
last_failure_at BIGINT,
failure_count BIGINT NOT NULL DEFAULT 0,
last_probe_attempt_at BIGINT,
last_probe_success_at BIGINT,
last_probe_failure_at BIGINT,
probe_failure_count BIGINT NOT NULL DEFAULT 0,
probe_status VARCHAR(64) NOT NULL DEFAULT 'never',
updated_at BIGINT NOT NULL,
KEY pool_member_scores_rank_idx (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC),
KEY pool_member_scores_member_idx (pool_kind, pool_id, member_kind, member_id),
KEY pool_member_scores_probe_idx (pool_kind, pool_id, probe_status, last_probe_success_at),
KEY pool_member_scores_updated_at_idx (updated_at)
);
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
id VARCHAR(64) PRIMARY KEY,
file_name VARCHAR(512) NOT NULL,

32
crates/aether-data/schema/drivers/postgres/baseline/001_types_and_tables.sql
View File

@@ -534,6 +534,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
--
-- Name: pool_member_scores; Type: TABLE; Schema: public; Owner: -
--
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision DEFAULT 0 NOT NULL,
hard_state character varying(64) DEFAULT 'unknown'::character varying NOT NULL,
score_version bigint DEFAULT 1 NOT NULL,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint DEFAULT 0 NOT NULL,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint DEFAULT 0 NOT NULL,
probe_status character varying(64) DEFAULT 'never'::character varying NOT NULL,
updated_at bigint NOT NULL
);
--
-- Name: provider_endpoints; Type: TABLE; Schema: public; Owner: -
--

15
crates/aether-data/schema/drivers/postgres/baseline/003_constraints.sql
View File

@@ -252,6 +252,21 @@ END $mig$;
--
-- Name: pool_member_scores pool_member_scores_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--
DO $mig$ BEGIN
ALTER TABLE ONLY public.pool_member_scores
ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
EXCEPTION
WHEN duplicate_object THEN NULL;
WHEN duplicate_table THEN NULL;
WHEN invalid_table_definition THEN NULL;
END $mig$;
--
-- Name: provider_endpoints provider_endpoints_pkey; Type: CONSTRAINT; Schema: public; Owner: -
--

32
crates/aether-data/schema/drivers/postgres/baseline/004_indexes.sql
View File

@@ -869,6 +869,38 @@ CREATE INDEX IF NOT EXISTS ix_provider_api_keys_id ON public.provider_api_keys U
--
-- Name: pool_member_scores_rank_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
--
-- Name: pool_member_scores_member_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
--
-- Name: pool_member_scores_probe_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
--
-- Name: pool_member_scores_updated_at_idx; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
--
-- Name: ix_provider_endpoints_id; Type: INDEX; Schema: public; Owner: -
--

30
crates/aether-data/schema/drivers/sqlite/baseline/002_provider_catalog.sql
View File

@@ -134,6 +134,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY,
pool_kind TEXT NOT NULL,
pool_id TEXT NOT NULL,
member_kind TEXT NOT NULL,
member_id TEXT NOT NULL,
capability TEXT NOT NULL,
scope_kind TEXT NOT NULL,
scope_id TEXT,
score REAL NOT NULL DEFAULT 0,
hard_state TEXT NOT NULL DEFAULT 'unknown',
score_version INTEGER NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at INTEGER,
last_scheduled_at INTEGER,
last_success_at INTEGER,
last_failure_at INTEGER,
failure_count INTEGER NOT NULL DEFAULT 0,
last_probe_attempt_at INTEGER,
last_probe_success_at INTEGER,
last_probe_failure_at INTEGER,
probe_failure_count INTEGER NOT NULL DEFAULT 0,
probe_status TEXT NOT NULL DEFAULT 'never',
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score DESC);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at);
CREATE TABLE IF NOT EXISTS gemini_file_mappings (
id TEXT PRIMARY KEY,
file_name TEXT NOT NULL UNIQUE,

31
crates/aether-data/schema/generated/mysql/baseline/002_provider_catalog.sql
View File

@@ -130,6 +130,37 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
KEY provider_api_keys_provider_id_idx (`provider_id`)
);
CREATE TABLE IF NOT EXISTS pool_member_scores (
`id` VARCHAR(192) NOT NULL,
`pool_kind` VARCHAR(64) NOT NULL,
`pool_id` VARCHAR(64) NOT NULL,
`member_kind` VARCHAR(64) NOT NULL,
`member_id` VARCHAR(64) NOT NULL,
`capability` VARCHAR(64) NOT NULL,
`scope_kind` VARCHAR(64) NOT NULL,
`scope_id` VARCHAR(128),
`score` DOUBLE NOT NULL DEFAULT 0,
`hard_state` VARCHAR(64) NOT NULL DEFAULT 'unknown',
`score_version` BIGINT NOT NULL DEFAULT 1,
`score_reason` JSON NOT NULL,
`last_ranked_at` BIGINT,
`last_scheduled_at` BIGINT,
`last_success_at` BIGINT,
`last_failure_at` BIGINT,
`failure_count` BIGINT NOT NULL DEFAULT 0,
`last_probe_attempt_at` BIGINT,
`last_probe_success_at` BIGINT,
`last_probe_failure_at` BIGINT,
`probe_failure_count` BIGINT NOT NULL DEFAULT 0,
`probe_status` VARCHAR(64) NOT NULL DEFAULT 'never',
`updated_at` BIGINT NOT NULL,
PRIMARY KEY (`id`),
KEY pool_member_scores_rank_idx (`pool_kind`, `pool_id`, `capability`, `scope_kind`, `scope_id`, `hard_state`, `score`),
KEY pool_member_scores_member_idx (`pool_kind`, `pool_id`, `member_kind`, `member_id`),
KEY pool_member_scores_probe_idx (`pool_kind`, `pool_id`, `probe_status`, `last_probe_success_at`),
KEY pool_member_scores_updated_at_idx (`updated_at`)
);
CREATE TABLE IF NOT EXISTS api_key_provider_mappings (
`id` VARCHAR(64) NOT NULL,
`api_key_id` VARCHAR(64) NOT NULL,

32
crates/aether-data/schema/generated/postgres/baseline/002_provider_catalog.sql
View File

@@ -134,6 +134,38 @@ CREATE TABLE IF NOT EXISTS public.provider_api_keys (
ALTER TABLE ONLY public.provider_api_keys ADD CONSTRAINT provider_api_keys_pkey PRIMARY KEY (id);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON public.provider_api_keys USING btree (provider_id);
CREATE TABLE IF NOT EXISTS public.pool_member_scores (
id character varying(192) NOT NULL,
pool_kind character varying(64) NOT NULL,
pool_id character varying(64) NOT NULL,
member_kind character varying(64) NOT NULL,
member_id character varying(64) NOT NULL,
capability character varying(64) NOT NULL,
scope_kind character varying(64) NOT NULL,
scope_id character varying(128),
score double precision DEFAULT 0 NOT NULL,
hard_state character varying(64) DEFAULT 'unknown' NOT NULL,
score_version bigint DEFAULT 1 NOT NULL,
score_reason jsonb NOT NULL,
last_ranked_at bigint,
last_scheduled_at bigint,
last_success_at bigint,
last_failure_at bigint,
failure_count bigint DEFAULT 0 NOT NULL,
last_probe_attempt_at bigint,
last_probe_success_at bigint,
last_probe_failure_at bigint,
probe_failure_count bigint DEFAULT 0 NOT NULL,
probe_status character varying(64) DEFAULT 'never' NOT NULL,
updated_at bigint NOT NULL
);
ALTER TABLE ONLY public.pool_member_scores ADD CONSTRAINT pool_member_scores_pkey PRIMARY KEY (id);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON public.pool_member_scores USING btree (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON public.pool_member_scores USING btree (updated_at);
CREATE TABLE IF NOT EXISTS public.api_key_provider_mappings (
id character varying(64) NOT NULL,
api_key_id character varying(64) NOT NULL,

30
crates/aether-data/schema/generated/sqlite/baseline/002_provider_catalog.sql
View File

@@ -126,6 +126,36 @@ CREATE TABLE IF NOT EXISTS provider_api_keys (
);
CREATE INDEX IF NOT EXISTS provider_api_keys_provider_id_idx ON provider_api_keys (provider_id);
CREATE TABLE IF NOT EXISTS pool_member_scores (
id TEXT PRIMARY KEY NOT NULL,
pool_kind TEXT NOT NULL,
pool_id TEXT NOT NULL,
member_kind TEXT NOT NULL,
member_id TEXT NOT NULL,
capability TEXT NOT NULL,
scope_kind TEXT NOT NULL,
scope_id TEXT,
score REAL NOT NULL DEFAULT 0,
hard_state TEXT NOT NULL DEFAULT 'unknown',
score_version INTEGER NOT NULL DEFAULT 1,
score_reason TEXT NOT NULL,
last_ranked_at INTEGER,
last_scheduled_at INTEGER,
last_success_at INTEGER,
last_failure_at INTEGER,
failure_count INTEGER NOT NULL DEFAULT 0,
last_probe_attempt_at INTEGER,
last_probe_success_at INTEGER,
last_probe_failure_at INTEGER,
probe_failure_count INTEGER NOT NULL DEFAULT 0,
probe_status TEXT NOT NULL DEFAULT 'never',
updated_at INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS pool_member_scores_rank_idx ON pool_member_scores (pool_kind, pool_id, capability, scope_kind, scope_id, hard_state, score);
CREATE INDEX IF NOT EXISTS pool_member_scores_member_idx ON pool_member_scores (pool_kind, pool_id, member_kind, member_id);
CREATE INDEX IF NOT EXISTS pool_member_scores_probe_idx ON pool_member_scores (pool_kind, pool_id, probe_status, last_probe_success_at);
CREATE INDEX IF NOT EXISTS pool_member_scores_updated_at_idx ON pool_member_scores (updated_at);
CREATE TABLE IF NOT EXISTS api_key_provider_mappings (
id TEXT PRIMARY KEY NOT NULL,
api_key_id TEXT NOT NULL,

137
crates/aether-data/schema/logical/002_provider_catalog.toml
View File

@@ -579,6 +579,143 @@ type = "unix_seconds"
name = "provider_api_keys_provider_id_idx"
columns = ["provider_id"]
[table.pool_member_scores]
domain = "provider_catalog"
order = 44
primary_key = ["id"]
[[table.pool_member_scores.columns]]
name = "id"
type = "text_id"
length = 192
[[table.pool_member_scores.columns]]
name = "pool_kind"
type = "text"
length = 64
[[table.pool_member_scores.columns]]
name = "pool_id"
type = "text_id"
length = 64
[[table.pool_member_scores.columns]]
name = "member_kind"
type = "text"
length = 64
[[table.pool_member_scores.columns]]
name = "member_id"
type = "text_id"
length = 64
[[table.pool_member_scores.columns]]
name = "capability"
type = "text"
length = 64
[[table.pool_member_scores.columns]]
name = "scope_kind"
type = "text"
length = 64
[[table.pool_member_scores.columns]]
name = "scope_id"
type = "text_id"
length = 128
nullable = true
[[table.pool_member_scores.columns]]
name = "score"
type = "float64"
default = 0
[[table.pool_member_scores.columns]]
name = "hard_state"
type = "text"
length = 64
default = "unknown"
[[table.pool_member_scores.columns]]
name = "score_version"
type = "int64"
default = 1
[[table.pool_member_scores.columns]]
name = "score_reason"
type = "json"
[[table.pool_member_scores.columns]]
name = "last_ranked_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_scheduled_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_success_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_failure_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "failure_count"
type = "int64"
default = 0
[[table.pool_member_scores.columns]]
name = "last_probe_attempt_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_probe_success_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "last_probe_failure_at"
type = "unix_seconds"
nullable = true
[[table.pool_member_scores.columns]]
name = "probe_failure_count"
type = "int64"
default = 0
[[table.pool_member_scores.columns]]
name = "probe_status"
type = "text"
length = 64
default = "never"
[[table.pool_member_scores.columns]]
name = "updated_at"
type = "unix_seconds"
[[table.pool_member_scores.indexes]]
name = "pool_member_scores_rank_idx"
columns = ["pool_kind", "pool_id", "capability", "scope_kind", "scope_id", "hard_state", "score"]
[[table.pool_member_scores.indexes]]
name = "pool_member_scores_member_idx"
columns = ["pool_kind", "pool_id", "member_kind", "member_id"]
[[table.pool_member_scores.indexes]]
name = "pool_member_scores_probe_idx"
columns = ["pool_kind", "pool_id", "probe_status", "last_probe_success_at"]
[[table.pool_member_scores.indexes]]
name = "pool_member_scores_updated_at_idx"
columns = ["updated_at"]
[table.api_key_provider_mappings]
domain = "provider_catalog"
order = 45

11
crates/aether-data/src/backend/mysql.rs
View File

@@ -37,6 +37,9 @@ use crate::repository::management_tokens::{
use crate::repository::oauth_providers::{
MysqlOAuthProviderRepository, OAuthProviderReadRepository, OAuthProviderWriteRepository,
};
use crate::repository::pool_scores::{
MysqlPoolMemberScoreRepository, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
};
use crate::repository::provider_catalog::{
MysqlProviderCatalogReadRepository, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository,
@@ -184,6 +187,14 @@ impl MysqlBackend {
Arc::new(MysqlProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone()))
}

11
crates/aether-data/src/backend/postgres.rs
View File

@@ -39,6 +39,9 @@ use crate::repository::management_tokens::{
use crate::repository::oauth_providers::{
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqlxOAuthProviderRepository,
};
use crate::repository::pool_scores::{
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PostgresPoolMemberScoreRepository,
};
use crate::repository::provider_catalog::{
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
SqlxProviderCatalogReadRepository,
@@ -195,6 +198,14 @@ impl PostgresBackend {
Arc::new(SqlxProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
Arc::new(PostgresPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
Arc::new(PostgresPoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn provider_quota_read_repository(&self) -> Arc<dyn ProviderQuotaReadRepository> {
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
}

12
crates/aether-data/src/backend/read.rs
View File

@@ -14,6 +14,7 @@ use crate::repository::gemini_file_mappings::GeminiFileMappingReadRepository;
use crate::repository::global_models::GlobalModelReadRepository;
use crate::repository::management_tokens::ManagementTokenReadRepository;
use crate::repository::oauth_providers::OAuthProviderReadRepository;
use crate::repository::pool_scores::PoolScoreReadRepository;
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
use crate::repository::proxy_nodes::ProxyNodeReadRepository;
use crate::repository::quota::ProviderQuotaReadRepository;
@@ -34,6 +35,7 @@ pub struct DataReadRepositories {
global_models: Option<Arc<dyn GlobalModelReadRepository>>,
management_tokens: Option<Arc<dyn ManagementTokenReadRepository>>,
oauth_providers: Option<Arc<dyn OAuthProviderReadRepository>>,
pool_scores: Option<Arc<dyn PoolScoreReadRepository>>,
proxy_nodes: Option<Arc<dyn ProxyNodeReadRepository>>,
minimal_candidate_selection: Option<Arc<dyn MinimalCandidateSelectionReadRepository>>,
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
@@ -61,6 +63,7 @@ impl fmt::Debug for DataReadRepositories {
.field("has_global_models", &self.global_models.is_some())
.field("has_management_tokens", &self.management_tokens.is_some())
.field("has_oauth_providers", &self.oauth_providers.is_some())
.field("has_pool_scores", &self.pool_scores.is_some())
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
.field(
"has_minimal_candidate_selection",
@@ -124,6 +127,10 @@ impl DataReadRepositories {
.map(PostgresBackend::oauth_provider_read_repository)
.or_else(|| mysql.map(MysqlBackend::oauth_provider_read_repository))
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_read_repository)),
pool_scores: postgres
.map(PostgresBackend::pool_score_read_repository)
.or_else(|| mysql.map(MysqlBackend::pool_score_read_repository))
.or_else(|| sqlite.map(SqliteBackend::pool_score_read_repository)),
proxy_nodes: postgres
.map(PostgresBackend::proxy_node_read_repository)
.or_else(|| mysql.map(MysqlBackend::proxy_node_read_repository))
@@ -208,6 +215,10 @@ impl DataReadRepositories {
self.oauth_providers.clone()
}
pub fn pool_scores(&self) -> Option<Arc<dyn PoolScoreReadRepository>> {
self.pool_scores.clone()
}
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeReadRepository>> {
self.proxy_nodes.clone()
}
@@ -257,6 +268,7 @@ impl DataReadRepositories {
|| self.global_models.is_some()
|| self.management_tokens.is_some()
|| self.oauth_providers.is_some()
|| self.pool_scores.is_some()
|| self.proxy_nodes.is_some()
|| self.minimal_candidate_selection.is_some()
|| self.request_candidates.is_some()

11
crates/aether-data/src/backend/sqlite.rs
View File

@@ -37,6 +37,9 @@ use crate::repository::management_tokens::{
use crate::repository::oauth_providers::{
OAuthProviderReadRepository, OAuthProviderWriteRepository, SqliteOAuthProviderRepository,
};
use crate::repository::pool_scores::{
PoolMemberScoreWriteRepository, PoolScoreReadRepository, SqlitePoolMemberScoreRepository,
};
use crate::repository::provider_catalog::{
ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
SqliteProviderCatalogReadRepository,
@@ -197,6 +200,14 @@ impl SqliteBackend {
Arc::new(SqliteProviderCatalogReadRepository::new(self.pool_clone()))
}
pub fn pool_score_read_repository(&self) -> Arc<dyn PoolScoreReadRepository> {
Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn pool_score_write_repository(&self) -> Arc<dyn PoolMemberScoreWriteRepository> {
Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone()))
}
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone()))
}

12
crates/aether-data/src/backend/write.rs
View File

@@ -11,6 +11,7 @@ use crate::repository::gemini_file_mappings::GeminiFileMappingWriteRepository;
use crate::repository::global_models::GlobalModelWriteRepository;
use crate::repository::management_tokens::ManagementTokenWriteRepository;
use crate::repository::oauth_providers::OAuthProviderWriteRepository;
use crate::repository::pool_scores::PoolMemberScoreWriteRepository;
use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
use crate::repository::quota::ProviderQuotaWriteRepository;
@@ -30,6 +31,7 @@ pub struct DataWriteRepositories {
global_models: Option<Arc<dyn GlobalModelWriteRepository>>,
management_tokens: Option<Arc<dyn ManagementTokenWriteRepository>>,
oauth_providers: Option<Arc<dyn OAuthProviderWriteRepository>>,
pool_scores: Option<Arc<dyn PoolMemberScoreWriteRepository>>,
proxy_nodes: Option<Arc<dyn ProxyNodeWriteRepository>>,
provider_catalog: Option<Arc<dyn ProviderCatalogWriteRepository>>,
provider_quotas: Option<Arc<dyn ProviderQuotaWriteRepository>>,
@@ -54,6 +56,7 @@ impl fmt::Debug for DataWriteRepositories {
.field("has_global_models", &self.global_models.is_some())
.field("has_management_tokens", &self.management_tokens.is_some())
.field("has_oauth_providers", &self.oauth_providers.is_some())
.field("has_pool_scores", &self.pool_scores.is_some())
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
.field("has_provider_catalog", &self.provider_catalog.is_some())
.field("has_provider_quotas", &self.provider_quotas.is_some())
@@ -108,6 +111,10 @@ impl DataWriteRepositories {
.map(PostgresBackend::oauth_provider_write_repository)
.or_else(|| mysql.map(MysqlBackend::oauth_provider_write_repository))
.or_else(|| sqlite.map(SqliteBackend::oauth_provider_write_repository)),
pool_scores: postgres
.map(PostgresBackend::pool_score_write_repository)
.or_else(|| mysql.map(MysqlBackend::pool_score_write_repository))
.or_else(|| sqlite.map(SqliteBackend::pool_score_write_repository)),
proxy_nodes: postgres
.map(PostgresBackend::proxy_node_write_repository)
.or_else(|| mysql.map(MysqlBackend::proxy_node_write_repository))
@@ -184,6 +191,10 @@ impl DataWriteRepositories {
self.oauth_providers.clone()
}
pub fn pool_scores(&self) -> Option<Arc<dyn PoolMemberScoreWriteRepository>> {
self.pool_scores.clone()
}
pub fn proxy_nodes(&self) -> Option<Arc<dyn ProxyNodeWriteRepository>> {
self.proxy_nodes.clone()
}
@@ -218,6 +229,7 @@ impl DataWriteRepositories {
|| self.global_models.is_some()
|| self.management_tokens.is_some()
|| self.oauth_providers.is_some()
|| self.pool_scores.is_some()
|| self.proxy_nodes.is_some()
|| self.provider_catalog.is_some()
|| self.provider_quotas.is_some()

37
crates/aether-data/src/repository/candidate_selection/memory.rs
View File

@@ -4,8 +4,8 @@ use async_trait::async_trait;
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery,
StoredRequestedModelCandidateRowsQuery,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredRequestedModelCandidateRowsQuery,
};
use crate::DataLayerError;
@@ -138,6 +138,39 @@ impl MinimalCandidateSelectionReadRepository for InMemoryMinimalCandidateSelecti
.take(query.limit as usize)
.collect())
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<std::collections::BTreeMap<_, _>>();
let mut rows = self
.list_for_exact_api_format(&query.api_format)
.await?
.into_iter()
.filter(|row| {
row.provider_id == query.provider_id
&& row.endpoint_id == query.endpoint_id
&& row.model_id == query.model_id
&& key_order.contains_key(row.key_id.as_str())
})
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(rows)
}
}
fn sort_pool_key_rows(

4
crates/aether-data/src/repository/candidate_selection/mod.rs
View File

@@ -7,8 +7,8 @@ mod sqlite;
pub(crate) use aether_data_contracts::repository::candidate_selection::{
MinimalCandidateSelectionReadRepository, MinimalCandidateSelectionRepository,
StoredMinimalCandidateSelectionRow, StoredPoolKeyCandidateOrder,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
StoredPoolKeyCandidateRowsByKeyIdsQuery, StoredPoolKeyCandidateRowsQuery,
StoredProviderModelMapping, StoredRequestedModelCandidateRowsQuery,
};
pub use memory::InMemoryMinimalCandidateSelectionReadRepository;
pub use mysql::MysqlMinimalCandidateSelectionReadRepository;

37
crates/aether-data/src/repository/candidate_selection/mysql.rs
View File

@@ -5,7 +5,8 @@ use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use crate::driver::mysql::MysqlPool;
@@ -196,6 +197,40 @@ impl MinimalCandidateSelectionReadRepository for MysqlMinimalCandidateSelectionR
.map(|item| item.row)
.collect())
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<BTreeMap<_, _>>();
let mut rows = self
.load_rows_for_api_format(&query.api_format)
.await?
.into_iter()
.filter(|row| {
row.row.provider_id == query.provider_id
&& row.row.endpoint_id == query.endpoint_id
&& row.row.model_id == query.model_id
&& key_order.contains_key(row.row.key_id.as_str())
})
.map(|item| item.row)
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(dedupe_candidate_selection_rows(rows))
}
}
fn select_pool_rows(rows: Vec<CandidateSelectionRow>) -> Vec<StoredMinimalCandidateSelectionRow> {

62
crates/aether-data/src/repository/candidate_selection/postgres.rs
View File

@@ -5,7 +5,8 @@ use std::collections::BTreeSet;
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use crate::{error::SqlxResultExt, DataLayerError};
@@ -535,6 +536,13 @@ fn pool_key_candidate_selection_sql(order: &StoredPoolKeyCandidateOrder) -> Stri
LIST_POOL_KEYS_FOR_GROUP_SQL.replace(default_order, &replacement)
}
fn pool_key_candidate_selection_by_key_ids_sql() -> String {
let default_order =
"ORDER BY\n pak.internal_priority ASC,\n pak.id ASC\nLIMIT $7\nOFFSET $8\n";
let replacement = "AND pak.id = ANY($7::text[])\nORDER BY\n array_position($7::text[], pak.id) ASC,\n pak.id ASC\n";
LIST_POOL_KEYS_FOR_GROUP_SQL.replace(default_order, replacement)
}
#[derive(Debug, Clone)]
pub struct SqlxMinimalCandidateSelectionReadRepository {
pool: PgPool,
@@ -704,6 +712,51 @@ impl SqlxMinimalCandidateSelectionReadRepository {
}
Ok(dedupe_candidate_selection_rows(rows))
}
pub async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let mut rows = Vec::new();
let canonical_api_format = normalize_api_format(&query.api_format);
let storage_aliases = api_format_aliases(&canonical_api_format);
let sql_match_aliases = sql_match_aliases(&storage_aliases);
let sql = pool_key_candidate_selection_by_key_ids_sql();
for api_format in storage_aliases {
rows.extend(
Self::collect_query_rows(
sqlx::query(sql.as_str())
.bind(api_format)
.bind(query.provider_id.as_str())
.bind(query.endpoint_id.as_str())
.bind(query.model_id.as_str())
.bind(sql_match_aliases.clone())
.bind(canonical_api_format.clone())
.bind(query.key_ids.clone())
.fetch(&self.pool),
map_candidate_selection_row,
)
.await?,
);
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<std::collections::BTreeMap<_, _>>();
let mut rows = dedupe_candidate_selection_rows(rows);
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(rows)
}
}
fn requested_model_selection_sql() -> String {
@@ -820,6 +873,13 @@ impl MinimalCandidateSelectionReadRepository for SqlxMinimalCandidateSelectionRe
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
Self::list_pool_key_rows_for_group(self, query).await
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
Self::list_pool_key_rows_for_group_key_ids(self, query).await
}
}
fn map_candidate_selection_row(

37
crates/aether-data/src/repository/candidate_selection/sqlite.rs
View File

@@ -5,7 +5,8 @@ use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite};
use super::{
MinimalCandidateSelectionReadRepository, StoredMinimalCandidateSelectionRow,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredPoolKeyCandidateOrder, StoredPoolKeyCandidateRowsByKeyIdsQuery,
StoredPoolKeyCandidateRowsQuery, StoredProviderModelMapping,
StoredRequestedModelCandidateRowsQuery,
};
use crate::driver::sqlite::SqlitePool;
@@ -196,6 +197,40 @@ impl MinimalCandidateSelectionReadRepository for SqliteMinimalCandidateSelection
.map(|item| item.row)
.collect())
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
if query.key_ids.is_empty() {
return Ok(Vec::new());
}
let key_order = query
.key_ids
.iter()
.enumerate()
.map(|(index, key_id)| (key_id.as_str(), index))
.collect::<BTreeMap<_, _>>();
let mut rows = self
.load_rows_for_api_format(&query.api_format)
.await?
.into_iter()
.filter(|row| {
row.row.provider_id == query.provider_id
&& row.row.endpoint_id == query.endpoint_id
&& row.row.model_id == query.model_id
&& key_order.contains_key(row.row.key_id.as_str())
})
.map(|item| item.row)
.collect::<Vec<_>>();
rows.sort_by(|left, right| {
key_order
.get(left.key_id.as_str())
.cmp(&key_order.get(right.key_id.as_str()))
.then(left.key_id.cmp(&right.key_id))
});
Ok(dedupe_candidate_selection_rows(rows))
}
}
fn select_pool_rows(rows: Vec<CandidateSelectionRow>) -> Vec<StoredMinimalCandidateSelectionRow> {

1
crates/aether-data/src/repository/mod.rs
View File

@@ -16,6 +16,7 @@ pub mod gemini_file_mappings;
pub mod global_models;
pub mod management_tokens;
pub mod oauth_providers;
pub mod pool_scores;
pub mod provider_catalog;
pub mod provider_oauth;
pub mod proxy_nodes;

512
crates/aether-data/src/repository/pool_scores/memory.rs Normal file
View File

@@ -0,0 +1,512 @@
use std::collections::BTreeMap;
use std::sync::RwLock;
use async_trait::async_trait;
use super::{
score_with_delta, GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery,
ListPoolMemberScoresQuery, ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity,
PoolMemberProbeAttempt, PoolMemberProbeResult, PoolMemberProbeStatus,
PoolMemberScheduleFeedback, PoolMemberScoreWriteRepository, PoolScoreReadRepository,
PoolScoreScope, StoredPoolMemberScore, UpsertPoolMemberScore,
};
use crate::repository::pool_scores::merge_score_reason_patch;
use crate::DataLayerError;
#[derive(Debug, Default)]
pub struct InMemoryPoolMemberScoreRepository {
scores: RwLock<BTreeMap<String, StoredPoolMemberScore>>,
}
impl InMemoryPoolMemberScoreRepository {
pub fn seed<I>(scores: I) -> Self
where
I: IntoIterator<Item = StoredPoolMemberScore>,
{
Self {
scores: RwLock::new(
scores
.into_iter()
.map(|score| (score.id.clone(), score))
.collect(),
),
}
}
fn matches_identity(score: &StoredPoolMemberScore, identity: &PoolMemberIdentity) -> bool {
score.pool_kind == identity.pool_kind
&& score.pool_id == identity.pool_id
&& score.member_kind == identity.member_kind
&& score.member_id == identity.member_id
}
fn matches_scope(score: &StoredPoolMemberScore, scope: Option<&PoolScoreScope>) -> bool {
let Some(scope) = scope else {
return true;
};
score.capability == scope.capability
&& score.scope_kind == scope.scope_kind
&& score.scope_id == scope.scope_id
}
fn sort_ranked(scores: &mut [StoredPoolMemberScore]) {
scores.sort_by(|left, right| {
right
.score
.partial_cmp(&left.score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| {
right
.last_ranked_at
.unwrap_or(0)
.cmp(&left.last_ranked_at.unwrap_or(0))
})
.then_with(|| left.member_id.cmp(&right.member_id))
.then_with(|| left.id.cmp(&right.id))
});
}
fn sort_probe(scores: &mut [StoredPoolMemberScore]) {
scores.sort_by(|left, right| {
probe_priority(left)
.cmp(&probe_priority(right))
.then_with(|| {
left.last_probe_success_at
.unwrap_or(0)
.cmp(&right.last_probe_success_at.unwrap_or(0))
})
.then_with(|| {
left.last_scheduled_at
.unwrap_or(0)
.cmp(&right.last_scheduled_at.unwrap_or(0))
.reverse()
})
.then_with(|| left.member_id.cmp(&right.member_id))
});
}
}
#[async_trait]
impl PoolScoreReadRepository for InMemoryPoolMemberScoreRepository {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let hard_states = query
.hard_states
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>();
let probe_statuses = query.probe_statuses.as_ref().map(|items| {
items
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>()
});
let mut scores = self
.scores
.read()
.expect("pool member score repository lock")
.values()
.filter(|score| {
score.pool_kind == query.pool_kind
&& score.pool_id == query.pool_id
&& score.capability == query.capability
&& score.scope_kind == query.scope_kind
&& score.scope_id == query.scope_id
&& (hard_states.is_empty() || hard_states.contains(&score.hard_state))
&& probe_statuses
.as_ref()
.is_none_or(|statuses| statuses.contains(&score.probe_status))
})
.cloned()
.collect::<Vec<_>>();
Self::sort_ranked(&mut scores);
Ok(scores
.into_iter()
.skip(query.offset)
.take(query.limit.max(1))
.collect())
}
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let hard_states = query
.hard_states
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>();
let probe_statuses = query.probe_statuses.as_ref().map(|items| {
items
.iter()
.copied()
.collect::<std::collections::BTreeSet<_>>()
});
let mut scores = self
.scores
.read()
.expect("pool member score repository lock")
.values()
.filter(|score| {
score.pool_kind == query.pool_kind
&& score.pool_id == query.pool_id
&& query
.capability
.as_ref()
.is_none_or(|capability| score.capability == *capability)
&& query
.scope_kind
.as_ref()
.is_none_or(|scope_kind| score.scope_kind == *scope_kind)
&& query
.scope_id
.as_ref()
.is_none_or(|scope_id| score.scope_id.as_ref() == Some(scope_id))
&& (hard_states.is_empty() || hard_states.contains(&score.hard_state))
&& probe_statuses
.as_ref()
.is_none_or(|statuses| statuses.contains(&score.probe_status))
})
.cloned()
.collect::<Vec<_>>();
Self::sort_ranked(&mut scores);
Ok(scores
.into_iter()
.skip(query.offset)
.take(query.limit.max(1))
.collect())
}
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut scores = self
.scores
.read()
.expect("pool member score repository lock")
.values()
.filter(|score| {
score.pool_kind == query.pool_kind
&& score.pool_id == query.pool_id
&& query
.capability
.as_ref()
.is_none_or(|capability| score.capability == *capability)
&& matches!(
score.hard_state,
PoolMemberHardState::Available
| PoolMemberHardState::Unknown
| PoolMemberHardState::Cooldown
| PoolMemberHardState::QuotaExhausted
)
&& match score.probe_status {
PoolMemberProbeStatus::Never
| PoolMemberProbeStatus::Failed
| PoolMemberProbeStatus::Stale => true,
PoolMemberProbeStatus::Ok => score
.last_probe_success_at
.is_none_or(|ts| ts <= query.stale_before_unix_secs),
PoolMemberProbeStatus::InProgress => score
.last_probe_attempt_at
.is_none_or(|ts| ts <= query.stale_before_unix_secs),
}
})
.cloned()
.collect::<Vec<_>>();
Self::sort_probe(&mut scores);
Ok(scores.into_iter().take(query.limit.max(1)).collect())
}
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let ids = query
.ids
.iter()
.cloned()
.collect::<std::collections::BTreeSet<_>>();
let scores = self
.scores
.read()
.expect("pool member score repository lock")
.values()
.filter(|score| ids.contains(&score.id))
.cloned()
.collect::<Vec<_>>();
Ok(scores)
}
}
#[async_trait]
impl PoolMemberScoreWriteRepository for InMemoryPoolMemberScoreRepository {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, DataLayerError> {
score.validate()?;
let stored = score.into_stored();
self.scores
.write()
.expect("pool member score repository lock")
.insert(stored.id.clone(), stored.clone());
Ok(stored)
}
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
let mut updated = 0;
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
for score in guard.values_mut() {
if !Self::matches_identity(score, &attempt.identity)
|| !Self::matches_scope(score, attempt.scope.as_ref())
{
continue;
}
score.last_probe_attempt_at = Some(attempt.attempted_at);
score.probe_status = PoolMemberProbeStatus::InProgress;
score.score_reason = merge_score_reason_patch(
score.score_reason.clone(),
attempt.score_reason_patch.clone(),
);
score.updated_at = attempt.attempted_at;
updated += 1;
}
Ok(updated)
}
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
let mut updated = 0;
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
for score in guard.values_mut() {
if !Self::matches_identity(score, &result.identity)
|| !Self::matches_scope(score, result.scope.as_ref())
{
continue;
}
score.last_probe_attempt_at = Some(result.attempted_at);
score.probe_status = result.probe_status;
if result.succeeded {
score.last_probe_success_at = Some(result.attempted_at);
score.probe_failure_count = 0;
} else {
score.last_probe_failure_at = Some(result.attempted_at);
score.probe_failure_count = score.probe_failure_count.saturating_add(1);
}
if let Some(hard_state) = result.hard_state {
score.hard_state = hard_state;
}
score.score_reason = merge_score_reason_patch(
score.score_reason.clone(),
result.score_reason_patch.clone(),
);
score.updated_at = result.attempted_at;
updated += 1;
}
Ok(updated)
}
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
let mut updated = 0;
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
for score in guard.values_mut() {
if !Self::matches_identity(score, &feedback.identity)
|| !Self::matches_scope(score, feedback.scope.as_ref())
{
continue;
}
score.last_scheduled_at = Some(feedback.scheduled_at);
match feedback.succeeded {
Some(true) => {
score.last_success_at = Some(feedback.scheduled_at);
}
Some(false) => {
score.last_failure_at = Some(feedback.scheduled_at);
score.failure_count = score.failure_count.saturating_add(1);
}
None => {}
}
if let Some(hard_state) = feedback.hard_state {
score.hard_state = hard_state;
}
score.score = score_with_delta(score.score, feedback.score_delta);
score.score_reason = merge_score_reason_patch(
score.score_reason.clone(),
feedback.score_reason_patch.clone(),
);
score.updated_at = feedback.scheduled_at;
updated += 1;
}
Ok(updated)
}
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
let mut updated = 0;
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
for score in guard.values_mut() {
if Self::matches_identity(score, identity) && Self::matches_scope(score, scope) {
score.hard_state = hard_state;
score.updated_at = updated_at;
updated += 1;
}
}
Ok(updated)
}
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
let mut guard = self
.scores
.write()
.expect("pool member score repository lock");
let before = guard.len();
guard.retain(|_, score| !Self::matches_identity(score, identity));
Ok(before.saturating_sub(guard.len()))
}
}
fn probe_priority(score: &StoredPoolMemberScore) -> u8 {
if score.last_scheduled_at.is_some() && score.probe_status != PoolMemberProbeStatus::Ok {
return 0;
}
match score.hard_state {
PoolMemberHardState::QuotaExhausted => 1,
PoolMemberHardState::Unknown => 2,
_ if score.probe_status == PoolMemberProbeStatus::Stale => 3,
_ => 4,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::repository::pool_scores::{
POOL_KIND_PROVIDER_KEY_POOL, POOL_MEMBER_KIND_PROVIDER_API_KEY,
POOL_SCORE_CAPABILITY_API_FORMAT, POOL_SCORE_SCOPE_KIND_MODEL,
};
fn score(id: &str, member_id: &str, value: f64) -> StoredPoolMemberScore {
StoredPoolMemberScore {
id: id.to_string(),
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-1".to_string(),
member_kind: POOL_MEMBER_KIND_PROVIDER_API_KEY.to_string(),
member_id: member_id.to_string(),
capability: POOL_SCORE_CAPABILITY_API_FORMAT.to_string(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: Some("model-1".to_string()),
score: value,
hard_state: PoolMemberHardState::Available,
score_version: 1,
score_reason: serde_json::json!({}),
last_ranked_at: Some(1),
last_scheduled_at: None,
last_success_at: None,
last_failure_at: None,
failure_count: 0,
last_probe_attempt_at: None,
last_probe_success_at: None,
last_probe_failure_at: None,
probe_failure_count: 0,
probe_status: PoolMemberProbeStatus::Never,
updated_at: 1,
}
}
#[tokio::test]
async fn lists_ranked_members_by_score() {
let repository = InMemoryPoolMemberScoreRepository::seed(vec![
score("score-1", "key-1", 0.2),
score("score-2", "key-2", 0.9),
]);
let rows = repository
.list_ranked_pool_members(&ListRankedPoolMembersQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-1".to_string(),
capability: POOL_SCORE_CAPABILITY_API_FORMAT.to_string(),
scope_kind: POOL_SCORE_SCOPE_KIND_MODEL.to_string(),
scope_id: Some("model-1".to_string()),
hard_states: vec![PoolMemberHardState::Available],
probe_statuses: None,
offset: 0,
limit: 10,
})
.await
.expect("list should succeed");
assert_eq!(
rows.into_iter()
.map(|row| row.member_id)
.collect::<Vec<_>>(),
vec!["key-2".to_string(), "key-1".to_string()]
);
}
#[tokio::test]
async fn marks_probe_in_progress_without_incrementing_failure_count() {
let repository =
InMemoryPoolMemberScoreRepository::seed(vec![score("score-1", "key-1", 0.2)]);
let updated = repository
.mark_pool_member_probe_in_progress(PoolMemberProbeAttempt {
identity: PoolMemberIdentity::provider_api_key("provider-1", "key-1"),
scope: None,
attempted_at: 100,
score_reason_patch: Some(serde_json::json!({ "last_probe": "in_progress" })),
})
.await
.expect("mark should succeed");
assert_eq!(updated, 1);
let rows = repository
.list_pool_member_scores(&ListPoolMemberScoresQuery {
pool_kind: POOL_KIND_PROVIDER_KEY_POOL.to_string(),
pool_id: "provider-1".to_string(),
capability: None,
scope_kind: None,
scope_id: None,
hard_states: Vec::new(),
probe_statuses: Some(vec![PoolMemberProbeStatus::InProgress]),
offset: 0,
limit: 10,
})
.await
.expect("list should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].last_probe_attempt_at, Some(100));
assert_eq!(rows[0].probe_failure_count, 0);
}
}

54
crates/aether-data/src/repository/pool_scores/mod.rs Normal file
View File

@@ -0,0 +1,54 @@
pub use aether_data_contracts::repository::pool_scores::*;
mod memory;
mod mysql;
mod postgres;
mod sqlite;
pub use memory::InMemoryPoolMemberScoreRepository;
pub use mysql::MysqlPoolMemberScoreRepository;
pub use postgres::PostgresPoolMemberScoreRepository;
pub use sqlite::SqlitePoolMemberScoreRepository;
fn merge_score_reason_patch(
mut current: serde_json::Value,
patch: Option<serde_json::Value>,
) -> serde_json::Value {
let Some(patch) = patch else {
return current;
};
match (current.as_object_mut(), patch) {
(Some(current), serde_json::Value::Object(patch)) => {
for (key, value) in patch {
current.insert(key, value);
}
serde_json::Value::Object(current.clone())
}
(_, patch) => patch,
}
}
fn score_with_delta(score: f64, delta_basis_points: Option<i32>) -> f64 {
let delta = delta_basis_points.unwrap_or_default() as f64 / 10_000.0;
(score + delta).clamp(0.0, 1.0)
}
fn i64_from_u64(value: u64, field: &str) -> Result<i64, crate::DataLayerError> {
i64::try_from(value).map_err(|_| {
crate::DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range"))
})
}
fn i64_opt_from_u64(value: Option<u64>, field: &str) -> Result<Option<i64>, crate::DataLayerError> {
value.map(|value| i64_from_u64(value, field)).transpose()
}
fn u64_from_i64(value: i64, field: &str) -> Result<u64, crate::DataLayerError> {
u64::try_from(value).map_err(|_| {
crate::DataLayerError::UnexpectedValue(format!("{field} is negative: {value}"))
})
}
fn u64_opt_from_i64(value: Option<i64>, field: &str) -> Result<Option<u64>, crate::DataLayerError> {
value.map(|value| u64_from_i64(value, field)).transpose()
}

585
crates/aether-data/src/repository/pool_scores/mysql.rs Normal file
View File

@@ -0,0 +1,585 @@
use async_trait::async_trait;
use sqlx::{mysql::MySqlRow, MySql, MySqlPool, QueryBuilder, Row};
use super::{
i64_from_u64, i64_opt_from_u64, score_with_delta, u64_from_i64, u64_opt_from_i64,
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
UpsertPoolMemberScore,
};
use crate::error::SqlResultExt;
use crate::repository::pool_scores::merge_score_reason_patch;
use crate::DataLayerError;
const SCORE_COLUMNS: &str = r#"
SELECT
id,
pool_kind,
pool_id,
member_kind,
member_id,
capability,
scope_kind,
scope_id,
score,
hard_state,
score_version,
score_reason,
last_ranked_at,
last_scheduled_at,
last_success_at,
last_failure_at,
failure_count,
last_probe_attempt_at,
last_probe_success_at,
last_probe_failure_at,
probe_failure_count,
probe_status,
updated_at
FROM pool_member_scores
"#;
#[derive(Debug, Clone)]
pub struct MysqlPoolMemberScoreRepository {
pool: MySqlPool,
}
impl MysqlPoolMemberScoreRepository {
pub fn new(pool: MySqlPool) -> Self {
Self { pool }
}
async fn find_scores_by_identity(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(identity.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(identity.pool_id.clone())
.push(" AND member_kind = ")
.push_bind(identity.member_kind.clone())
.push(" AND member_id = ")
.push_bind(identity.member_id.clone());
if let Some(scope) = scope {
builder
.push(" AND capability = ")
.push_bind(scope.capability.clone())
.push(" AND scope_kind = ")
.push_bind(scope.scope_kind.clone());
if let Some(scope_id) = &scope.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
}
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolScoreReadRepository for MysqlPoolMemberScoreRepository {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone())
.push(" AND capability = ")
.push_bind(query.capability.clone())
.push(" AND scope_kind = ")
.push_bind(query.scope_kind.clone());
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
if let Some(scope_kind) = &query.scope_kind {
builder
.push(" AND scope_kind = ")
.push_bind(scope_kind.clone());
}
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
builder
.push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')")
.push(" AND (probe_status IN ('never','failed','stale')")
.push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push("))")
.push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push(")))")
.push(
r#"
ORDER BY
CASE
WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0
WHEN hard_state = 'quota_exhausted' THEN 1
WHEN hard_state = 'unknown' THEN 2
WHEN probe_status = 'stale' THEN 3
ELSE 4
END ASC,
COALESCE(last_probe_success_at, 0) ASC,
COALESCE(last_scheduled_at, 0) DESC,
member_id ASC
"#,
)
.push(" LIMIT ")
.push_bind(i64_from_usize(
query.limit.max(1),
"pool probe candidate limit",
)?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
if query.ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<MySql>::new(SCORE_COLUMNS);
builder.push(" WHERE id IN (");
let mut separated = builder.separated(", ");
for id in &query.ids {
separated.push_bind(id.clone());
}
separated.push_unseparated(")");
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolMemberScoreWriteRepository for MysqlPoolMemberScoreRepository {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, DataLayerError> {
score.validate()?;
let stored = score.into_stored();
let score_reason = serde_json::to_string(&stored.score_reason)
.map_err(|err| DataLayerError::InvalidInput(err.to_string()))?;
sqlx::query(
r#"
INSERT INTO pool_member_scores (
id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id,
score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at,
last_success_at, last_failure_at, failure_count, last_probe_attempt_at,
last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at
) VALUES (
?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?
)
ON DUPLICATE KEY UPDATE
pool_kind = VALUES(pool_kind),
pool_id = VALUES(pool_id),
member_kind = VALUES(member_kind),
member_id = VALUES(member_id),
capability = VALUES(capability),
scope_kind = VALUES(scope_kind),
scope_id = VALUES(scope_id),
score = VALUES(score),
hard_state = VALUES(hard_state),
score_version = VALUES(score_version),
score_reason = VALUES(score_reason),
last_ranked_at = VALUES(last_ranked_at),
last_scheduled_at = COALESCE(VALUES(last_scheduled_at), last_scheduled_at),
last_success_at = COALESCE(VALUES(last_success_at), last_success_at),
last_failure_at = COALESCE(VALUES(last_failure_at), last_failure_at),
failure_count = VALUES(failure_count),
last_probe_attempt_at = COALESCE(VALUES(last_probe_attempt_at), last_probe_attempt_at),
last_probe_success_at = COALESCE(VALUES(last_probe_success_at), last_probe_success_at),
last_probe_failure_at = COALESCE(VALUES(last_probe_failure_at), last_probe_failure_at),
probe_failure_count = VALUES(probe_failure_count),
probe_status = VALUES(probe_status),
updated_at = VALUES(updated_at)
"#,
)
.bind(stored.id.as_str())
.bind(stored.pool_kind.as_str())
.bind(stored.pool_id.as_str())
.bind(stored.member_kind.as_str())
.bind(stored.member_id.as_str())
.bind(stored.capability.as_str())
.bind(stored.scope_kind.as_str())
.bind(stored.scope_id.as_deref())
.bind(stored.score)
.bind(stored.hard_state.as_database())
.bind(i64_from_u64(stored.score_version, "pool score version")?)
.bind(score_reason)
.bind(i64_opt_from_u64(
stored.last_ranked_at,
"pool score last_ranked_at",
)?)
.bind(i64_opt_from_u64(
stored.last_scheduled_at,
"pool score last_scheduled_at",
)?)
.bind(i64_opt_from_u64(
stored.last_success_at,
"pool score last_success_at",
)?)
.bind(i64_opt_from_u64(
stored.last_failure_at,
"pool score last_failure_at",
)?)
.bind(i64_from_u64(
stored.failure_count,
"pool score failure_count",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_attempt_at,
"pool score last_probe_attempt_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_success_at,
"pool score last_probe_success_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_failure_at,
"pool score last_probe_failure_at",
)?)
.bind(i64_from_u64(
stored.probe_failure_count,
"pool score probe_failure_count",
)?)
.bind(stored.probe_status.as_database())
.bind(i64_from_u64(stored.updated_at, "pool score updated_at")?)
.execute(&self.pool)
.await
.map_sql_err()?;
Ok(stored)
}
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&attempt.identity, attempt.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(attempt.attempted_at);
row.probe_status = PoolMemberProbeStatus::InProgress;
row.score_reason =
merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone());
row.updated_at = attempt.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&result.identity, result.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(result.attempted_at);
row.probe_status = result.probe_status;
if result.succeeded {
row.last_probe_success_at = Some(result.attempted_at);
row.probe_failure_count = 0;
} else {
row.last_probe_failure_at = Some(result.attempted_at);
row.probe_failure_count = row.probe_failure_count.saturating_add(1);
}
if let Some(hard_state) = result.hard_state {
row.hard_state = hard_state;
}
row.score_reason =
merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone());
row.updated_at = result.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&feedback.identity, feedback.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_scheduled_at = Some(feedback.scheduled_at);
match feedback.succeeded {
Some(true) => row.last_success_at = Some(feedback.scheduled_at),
Some(false) => {
row.last_failure_at = Some(feedback.scheduled_at);
row.failure_count = row.failure_count.saturating_add(1);
}
None => {}
}
if let Some(hard_state) = feedback.hard_state {
row.hard_state = hard_state;
}
row.score = score_with_delta(row.score, feedback.score_delta);
row.score_reason =
merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone());
row.updated_at = feedback.scheduled_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
let rows = self.find_scores_by_identity(identity, scope).await?;
let count = rows.len();
for mut row in rows {
row.hard_state = hard_state;
row.updated_at = updated_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
let result = sqlx::query(
r#"
DELETE FROM pool_member_scores
WHERE pool_kind = ? AND pool_id = ? AND member_kind = ? AND member_id = ?
"#,
)
.bind(identity.pool_kind.as_str())
.bind(identity.pool_id.as_str())
.bind(identity.member_kind.as_str())
.bind(identity.member_id.as_str())
.execute(&self.pool)
.await
.map_sql_err()?;
Ok(result.rows_affected() as usize)
}
}
fn map_score_row(row: &MySqlRow) -> Result<StoredPoolMemberScore, DataLayerError> {
let score_reason_raw: String = row.try_get("score_reason").map_sql_err()?;
Ok(StoredPoolMemberScore {
id: row.try_get("id").map_sql_err()?,
pool_kind: row.try_get("pool_kind").map_sql_err()?,
pool_id: row.try_get("pool_id").map_sql_err()?,
member_kind: row.try_get("member_kind").map_sql_err()?,
member_id: row.try_get("member_id").map_sql_err()?,
capability: row.try_get("capability").map_sql_err()?,
scope_kind: row.try_get("scope_kind").map_sql_err()?,
scope_id: row.try_get("scope_id").map_sql_err()?,
score: row.try_get("score").map_sql_err()?,
hard_state: PoolMemberHardState::from_database(
row.try_get::<String, _>("hard_state")
.map_sql_err()?
.as_str(),
)?,
score_version: u64_from_i64(
row.try_get("score_version").map_sql_err()?,
"pool_member_scores.score_version",
)?,
score_reason: serde_json::from_str(&score_reason_raw).unwrap_or(serde_json::Value::Null),
last_ranked_at: u64_opt_from_i64(
row.try_get("last_ranked_at").map_sql_err()?,
"pool_member_scores.last_ranked_at",
)?,
last_scheduled_at: u64_opt_from_i64(
row.try_get("last_scheduled_at").map_sql_err()?,
"pool_member_scores.last_scheduled_at",
)?,
last_success_at: u64_opt_from_i64(
row.try_get("last_success_at").map_sql_err()?,
"pool_member_scores.last_success_at",
)?,
last_failure_at: u64_opt_from_i64(
row.try_get("last_failure_at").map_sql_err()?,
"pool_member_scores.last_failure_at",
)?,
failure_count: u64_from_i64(
row.try_get("failure_count").map_sql_err()?,
"pool_member_scores.failure_count",
)?,
last_probe_attempt_at: u64_opt_from_i64(
row.try_get("last_probe_attempt_at").map_sql_err()?,
"pool_member_scores.last_probe_attempt_at",
)?,
last_probe_success_at: u64_opt_from_i64(
row.try_get("last_probe_success_at").map_sql_err()?,
"pool_member_scores.last_probe_success_at",
)?,
last_probe_failure_at: u64_opt_from_i64(
row.try_get("last_probe_failure_at").map_sql_err()?,
"pool_member_scores.last_probe_failure_at",
)?,
probe_failure_count: u64_from_i64(
row.try_get("probe_failure_count").map_sql_err()?,
"pool_member_scores.probe_failure_count",
)?,
probe_status: PoolMemberProbeStatus::from_database(
row.try_get::<String, _>("probe_status")
.map_sql_err()?
.as_str(),
)?,
updated_at: u64_from_i64(
row.try_get("updated_at").map_sql_err()?,
"pool_member_scores.updated_at",
)?,
})
}
fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore {
UpsertPoolMemberScore {
id: score.id,
identity: PoolMemberIdentity {
pool_kind: score.pool_kind,
pool_id: score.pool_id,
member_kind: score.member_kind,
member_id: score.member_id,
},
scope: PoolScoreScope {
capability: score.capability,
scope_kind: score.scope_kind,
scope_id: score.scope_id,
},
score: score.score,
hard_state: score.hard_state,
score_version: score.score_version,
score_reason: score.score_reason,
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,
updated_at: score.updated_at,
}
}
fn i64_from_usize(value: usize, field: &str) -> Result<i64, DataLayerError> {
i64::try_from(value)
.map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range")))
}

587
crates/aether-data/src/repository/pool_scores/postgres.rs Normal file
View File

@@ -0,0 +1,587 @@
use async_trait::async_trait;
use sqlx::{postgres::PgRow, PgPool, Postgres, QueryBuilder, Row};
use super::{
i64_from_u64, i64_opt_from_u64, score_with_delta, u64_from_i64, u64_opt_from_i64,
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
UpsertPoolMemberScore,
};
use crate::error::SqlxResultExt;
use crate::repository::pool_scores::merge_score_reason_patch;
use crate::DataLayerError;
const SCORE_COLUMNS: &str = r#"
SELECT
id,
pool_kind,
pool_id,
member_kind,
member_id,
capability,
scope_kind,
scope_id,
score,
hard_state,
score_version,
score_reason,
last_ranked_at,
last_scheduled_at,
last_success_at,
last_failure_at,
failure_count,
last_probe_attempt_at,
last_probe_success_at,
last_probe_failure_at,
probe_failure_count,
probe_status,
updated_at
FROM pool_member_scores
"#;
#[derive(Debug, Clone)]
pub struct PostgresPoolMemberScoreRepository {
pool: PgPool,
}
impl PostgresPoolMemberScoreRepository {
pub fn new(pool: PgPool) -> Self {
Self { pool }
}
async fn find_scores_by_identity(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(identity.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(identity.pool_id.clone())
.push(" AND member_kind = ")
.push_bind(identity.member_kind.clone())
.push(" AND member_id = ")
.push_bind(identity.member_id.clone());
if let Some(scope) = scope {
builder
.push(" AND capability = ")
.push_bind(scope.capability.clone())
.push(" AND scope_kind = ")
.push_bind(scope.scope_kind.clone());
if let Some(scope_id) = &scope.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
}
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolScoreReadRepository for PostgresPoolMemberScoreRepository {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone())
.push(" AND capability = ")
.push_bind(query.capability.clone())
.push(" AND scope_kind = ")
.push_bind(query.scope_kind.clone());
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC NULLS LAST, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
if let Some(scope_kind) = &query.scope_kind {
builder
.push(" AND scope_kind = ")
.push_bind(scope_kind.clone());
}
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC NULLS LAST, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
builder
.push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')")
.push(" AND (probe_status IN ('never','failed','stale')")
.push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push("))")
.push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push(")))")
.push(
r#"
ORDER BY
CASE
WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0
WHEN hard_state = 'quota_exhausted' THEN 1
WHEN hard_state = 'unknown' THEN 2
WHEN probe_status = 'stale' THEN 3
ELSE 4
END ASC,
COALESCE(last_probe_success_at, 0) ASC,
COALESCE(last_scheduled_at, 0) DESC,
member_id ASC
"#,
)
.push(" LIMIT ")
.push_bind(i64_from_usize(
query.limit.max(1),
"pool probe candidate limit",
)?);
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
if query.ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<Postgres>::new(SCORE_COLUMNS);
builder.push(" WHERE id IN (");
let mut separated = builder.separated(", ");
for id in &query.ids {
separated.push_bind(id.clone());
}
separated.push_unseparated(")");
let rows = builder
.build()
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolMemberScoreWriteRepository for PostgresPoolMemberScoreRepository {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, DataLayerError> {
score.validate()?;
let stored = score.clone().into_stored();
sqlx::query(
r#"
INSERT INTO pool_member_scores (
id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id,
score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at,
last_success_at, last_failure_at, failure_count, last_probe_attempt_at,
last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at
) VALUES (
$1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21,$22,$23
)
ON CONFLICT(id) DO UPDATE SET
pool_kind = EXCLUDED.pool_kind,
pool_id = EXCLUDED.pool_id,
member_kind = EXCLUDED.member_kind,
member_id = EXCLUDED.member_id,
capability = EXCLUDED.capability,
scope_kind = EXCLUDED.scope_kind,
scope_id = EXCLUDED.scope_id,
score = EXCLUDED.score,
hard_state = EXCLUDED.hard_state,
score_version = EXCLUDED.score_version,
score_reason = EXCLUDED.score_reason,
last_ranked_at = EXCLUDED.last_ranked_at,
last_scheduled_at = COALESCE(EXCLUDED.last_scheduled_at, pool_member_scores.last_scheduled_at),
last_success_at = COALESCE(EXCLUDED.last_success_at, pool_member_scores.last_success_at),
last_failure_at = COALESCE(EXCLUDED.last_failure_at, pool_member_scores.last_failure_at),
failure_count = EXCLUDED.failure_count,
last_probe_attempt_at = COALESCE(EXCLUDED.last_probe_attempt_at, pool_member_scores.last_probe_attempt_at),
last_probe_success_at = COALESCE(EXCLUDED.last_probe_success_at, pool_member_scores.last_probe_success_at),
last_probe_failure_at = COALESCE(EXCLUDED.last_probe_failure_at, pool_member_scores.last_probe_failure_at),
probe_failure_count = EXCLUDED.probe_failure_count,
probe_status = EXCLUDED.probe_status,
updated_at = EXCLUDED.updated_at
"#,
)
.bind(stored.id.as_str())
.bind(stored.pool_kind.as_str())
.bind(stored.pool_id.as_str())
.bind(stored.member_kind.as_str())
.bind(stored.member_id.as_str())
.bind(stored.capability.as_str())
.bind(stored.scope_kind.as_str())
.bind(stored.scope_id.as_deref())
.bind(stored.score)
.bind(stored.hard_state.as_database())
.bind(i64_from_u64(stored.score_version, "pool score version")?)
.bind(&stored.score_reason)
.bind(i64_opt_from_u64(stored.last_ranked_at, "pool score last_ranked_at")?)
.bind(i64_opt_from_u64(stored.last_scheduled_at, "pool score last_scheduled_at")?)
.bind(i64_opt_from_u64(stored.last_success_at, "pool score last_success_at")?)
.bind(i64_opt_from_u64(stored.last_failure_at, "pool score last_failure_at")?)
.bind(i64_from_u64(stored.failure_count, "pool score failure_count")?)
.bind(i64_opt_from_u64(
stored.last_probe_attempt_at,
"pool score last_probe_attempt_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_success_at,
"pool score last_probe_success_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_failure_at,
"pool score last_probe_failure_at",
)?)
.bind(i64_from_u64(
stored.probe_failure_count,
"pool score probe_failure_count",
)?)
.bind(stored.probe_status.as_database())
.bind(i64_from_u64(stored.updated_at, "pool score updated_at")?)
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(stored)
}
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&attempt.identity, attempt.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(attempt.attempted_at);
row.probe_status = PoolMemberProbeStatus::InProgress;
row.score_reason =
merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone());
row.updated_at = attempt.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&result.identity, result.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(result.attempted_at);
row.probe_status = result.probe_status;
if result.succeeded {
row.last_probe_success_at = Some(result.attempted_at);
row.probe_failure_count = 0;
} else {
row.last_probe_failure_at = Some(result.attempted_at);
row.probe_failure_count = row.probe_failure_count.saturating_add(1);
}
if let Some(hard_state) = result.hard_state {
row.hard_state = hard_state;
}
row.score_reason =
merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone());
row.updated_at = result.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&feedback.identity, feedback.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_scheduled_at = Some(feedback.scheduled_at);
match feedback.succeeded {
Some(true) => row.last_success_at = Some(feedback.scheduled_at),
Some(false) => {
row.last_failure_at = Some(feedback.scheduled_at);
row.failure_count = row.failure_count.saturating_add(1);
}
None => {}
}
if let Some(hard_state) = feedback.hard_state {
row.hard_state = hard_state;
}
row.score = score_with_delta(row.score, feedback.score_delta);
row.score_reason =
merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone());
row.updated_at = feedback.scheduled_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
let rows = self.find_scores_by_identity(identity, scope).await?;
let count = rows.len();
for mut row in rows {
row.hard_state = hard_state;
row.updated_at = updated_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
let result = sqlx::query(
r#"
DELETE FROM pool_member_scores
WHERE pool_kind = $1 AND pool_id = $2 AND member_kind = $3 AND member_id = $4
"#,
)
.bind(identity.pool_kind.as_str())
.bind(identity.pool_id.as_str())
.bind(identity.member_kind.as_str())
.bind(identity.member_id.as_str())
.execute(&self.pool)
.await
.map_postgres_err()?;
Ok(result.rows_affected() as usize)
}
}
fn map_score_row(row: &PgRow) -> Result<StoredPoolMemberScore, DataLayerError> {
Ok(StoredPoolMemberScore {
id: row.try_get("id").map_postgres_err()?,
pool_kind: row.try_get("pool_kind").map_postgres_err()?,
pool_id: row.try_get("pool_id").map_postgres_err()?,
member_kind: row.try_get("member_kind").map_postgres_err()?,
member_id: row.try_get("member_id").map_postgres_err()?,
capability: row.try_get("capability").map_postgres_err()?,
scope_kind: row.try_get("scope_kind").map_postgres_err()?,
scope_id: row.try_get("scope_id").map_postgres_err()?,
score: row.try_get("score").map_postgres_err()?,
hard_state: PoolMemberHardState::from_database(
row.try_get::<String, _>("hard_state")
.map_postgres_err()?
.as_str(),
)?,
score_version: u64_from_i64(
row.try_get("score_version").map_postgres_err()?,
"pool_member_scores.score_version",
)?,
score_reason: row.try_get("score_reason").map_postgres_err()?,
last_ranked_at: u64_opt_from_i64(
row.try_get("last_ranked_at").map_postgres_err()?,
"pool_member_scores.last_ranked_at",
)?,
last_scheduled_at: u64_opt_from_i64(
row.try_get("last_scheduled_at").map_postgres_err()?,
"pool_member_scores.last_scheduled_at",
)?,
last_success_at: u64_opt_from_i64(
row.try_get("last_success_at").map_postgres_err()?,
"pool_member_scores.last_success_at",
)?,
last_failure_at: u64_opt_from_i64(
row.try_get("last_failure_at").map_postgres_err()?,
"pool_member_scores.last_failure_at",
)?,
failure_count: u64_from_i64(
row.try_get("failure_count").map_postgres_err()?,
"pool_member_scores.failure_count",
)?,
last_probe_attempt_at: u64_opt_from_i64(
row.try_get("last_probe_attempt_at").map_postgres_err()?,
"pool_member_scores.last_probe_attempt_at",
)?,
last_probe_success_at: u64_opt_from_i64(
row.try_get("last_probe_success_at").map_postgres_err()?,
"pool_member_scores.last_probe_success_at",
)?,
last_probe_failure_at: u64_opt_from_i64(
row.try_get("last_probe_failure_at").map_postgres_err()?,
"pool_member_scores.last_probe_failure_at",
)?,
probe_failure_count: u64_from_i64(
row.try_get("probe_failure_count").map_postgres_err()?,
"pool_member_scores.probe_failure_count",
)?,
probe_status: PoolMemberProbeStatus::from_database(
row.try_get::<String, _>("probe_status")
.map_postgres_err()?
.as_str(),
)?,
updated_at: u64_from_i64(
row.try_get("updated_at").map_postgres_err()?,
"pool_member_scores.updated_at",
)?,
})
}
fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore {
UpsertPoolMemberScore {
id: score.id,
identity: PoolMemberIdentity {
pool_kind: score.pool_kind,
pool_id: score.pool_id,
member_kind: score.member_kind,
member_id: score.member_id,
},
scope: PoolScoreScope {
capability: score.capability,
scope_kind: score.scope_kind,
scope_id: score.scope_id,
},
score: score.score,
hard_state: score.hard_state,
score_version: score.score_version,
score_reason: score.score_reason,
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,
updated_at: score.updated_at,
}
}
fn i64_from_usize(value: usize, field: &str) -> Result<i64, DataLayerError> {
i64::try_from(value)
.map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range")))
}

570
crates/aether-data/src/repository/pool_scores/sqlite.rs Normal file
View File

@@ -0,0 +1,570 @@
use async_trait::async_trait;
use sqlx::{sqlite::SqliteRow, QueryBuilder, Row, Sqlite, SqlitePool};
use super::{
i64_from_u64, i64_opt_from_u64, score_with_delta, u64_from_i64, u64_opt_from_i64,
GetPoolMemberScoresByIdsQuery, ListPoolMemberProbeCandidatesQuery, ListPoolMemberScoresQuery,
ListRankedPoolMembersQuery, PoolMemberHardState, PoolMemberIdentity, PoolMemberProbeAttempt,
PoolMemberProbeResult, PoolMemberProbeStatus, PoolMemberScheduleFeedback,
PoolMemberScoreWriteRepository, PoolScoreReadRepository, PoolScoreScope, StoredPoolMemberScore,
UpsertPoolMemberScore,
};
use crate::error::SqlResultExt;
use crate::repository::pool_scores::merge_score_reason_patch;
use crate::DataLayerError;
const SCORE_COLUMNS: &str = r#"
SELECT
id,
pool_kind,
pool_id,
member_kind,
member_id,
capability,
scope_kind,
scope_id,
score,
hard_state,
score_version,
score_reason,
last_ranked_at,
last_scheduled_at,
last_success_at,
last_failure_at,
failure_count,
last_probe_attempt_at,
last_probe_success_at,
last_probe_failure_at,
probe_failure_count,
probe_status,
updated_at
FROM pool_member_scores
"#;
#[derive(Debug, Clone)]
pub struct SqlitePoolMemberScoreRepository {
pool: SqlitePool,
}
impl SqlitePoolMemberScoreRepository {
pub fn new(pool: SqlitePool) -> Self {
Self { pool }
}
async fn find_scores_by_identity(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(identity.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(identity.pool_id.clone())
.push(" AND member_kind = ")
.push_bind(identity.member_kind.clone())
.push(" AND member_id = ")
.push_bind(identity.member_id.clone());
if let Some(scope) = scope {
builder
.push(" AND capability = ")
.push_bind(scope.capability.clone())
.push(" AND scope_kind = ")
.push_bind(scope.scope_kind.clone());
if let Some(scope_id) = &scope.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
}
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolScoreReadRepository for SqlitePoolMemberScoreRepository {
async fn list_ranked_pool_members(
&self,
query: &ListRankedPoolMembersQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone())
.push(" AND capability = ")
.push_bind(query.capability.clone())
.push(" AND scope_kind = ")
.push_bind(query.scope_kind.clone());
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
} else {
builder.push(" AND scope_id IS NULL");
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_scores(
&self,
query: &ListPoolMemberScoresQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
if let Some(scope_kind) = &query.scope_kind {
builder
.push(" AND scope_kind = ")
.push_bind(scope_kind.clone());
}
if let Some(scope_id) = &query.scope_id {
builder.push(" AND scope_id = ").push_bind(scope_id.clone());
}
if !query.hard_states.is_empty() {
builder.push(" AND hard_state IN (");
let mut separated = builder.separated(", ");
for state in &query.hard_states {
separated.push_bind(state.as_database());
}
separated.push_unseparated(")");
}
if let Some(statuses) = &query.probe_statuses {
if !statuses.is_empty() {
builder.push(" AND probe_status IN (");
let mut separated = builder.separated(", ");
for status in statuses {
separated.push_bind(status.as_database());
}
separated.push_unseparated(")");
}
}
builder
.push(" ORDER BY score DESC, last_ranked_at DESC, member_id ASC, id ASC")
.push(" LIMIT ")
.push_bind(i64_from_usize(query.limit.max(1), "pool score limit")?)
.push(" OFFSET ")
.push_bind(i64_from_usize(query.offset, "pool score offset")?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn list_pool_member_probe_candidates(
&self,
query: &ListPoolMemberProbeCandidatesQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder
.push(" WHERE pool_kind = ")
.push_bind(query.pool_kind.clone())
.push(" AND pool_id = ")
.push_bind(query.pool_id.clone());
if let Some(capability) = &query.capability {
builder
.push(" AND capability = ")
.push_bind(capability.clone());
}
builder
.push(" AND hard_state IN ('available','unknown','cooldown','quota_exhausted')")
.push(" AND (probe_status IN ('never','failed','stale')")
.push(" OR (probe_status = 'ok' AND (last_probe_success_at IS NULL OR last_probe_success_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push("))")
.push(" OR (probe_status = 'in_progress' AND (last_probe_attempt_at IS NULL OR last_probe_attempt_at <= ")
.push_bind(i64_from_u64(
query.stale_before_unix_secs,
"pool probe stale_before_unix_secs",
)?)
.push(")))")
.push(
r#"
ORDER BY
CASE
WHEN last_scheduled_at IS NOT NULL AND probe_status <> 'ok' THEN 0
WHEN hard_state = 'quota_exhausted' THEN 1
WHEN hard_state = 'unknown' THEN 2
WHEN probe_status = 'stale' THEN 3
ELSE 4
END ASC,
COALESCE(last_probe_success_at, 0) ASC,
COALESCE(last_scheduled_at, 0) DESC,
member_id ASC
"#,
)
.push(" LIMIT ")
.push_bind(i64_from_usize(
query.limit.max(1),
"pool probe candidate limit",
)?);
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
async fn get_pool_member_scores_by_ids(
&self,
query: &GetPoolMemberScoresByIdsQuery,
) -> Result<Vec<StoredPoolMemberScore>, DataLayerError> {
if query.ids.is_empty() {
return Ok(Vec::new());
}
let mut builder = QueryBuilder::<Sqlite>::new(SCORE_COLUMNS);
builder.push(" WHERE id IN (");
let mut separated = builder.separated(", ");
for id in &query.ids {
separated.push_bind(id.clone());
}
separated.push_unseparated(")");
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
rows.iter().map(map_score_row).collect()
}
}
#[async_trait]
impl PoolMemberScoreWriteRepository for SqlitePoolMemberScoreRepository {
async fn upsert_pool_member_score(
&self,
score: UpsertPoolMemberScore,
) -> Result<StoredPoolMemberScore, DataLayerError> {
score.validate()?;
let stored = score.into_stored();
let score_reason = serde_json::to_string(&stored.score_reason)
.map_err(|err| DataLayerError::InvalidInput(err.to_string()))?;
sqlx::query(
r#"
INSERT INTO pool_member_scores (
id, pool_kind, pool_id, member_kind, member_id, capability, scope_kind, scope_id,
score, hard_state, score_version, score_reason, last_ranked_at, last_scheduled_at,
last_success_at, last_failure_at, failure_count, last_probe_attempt_at,
last_probe_success_at, last_probe_failure_at, probe_failure_count, probe_status, updated_at
) VALUES (
?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?
)
ON CONFLICT(id) DO UPDATE SET
pool_kind = excluded.pool_kind,
pool_id = excluded.pool_id,
member_kind = excluded.member_kind,
member_id = excluded.member_id,
capability = excluded.capability,
scope_kind = excluded.scope_kind,
scope_id = excluded.scope_id,
score = excluded.score,
hard_state = excluded.hard_state,
score_version = excluded.score_version,
score_reason = excluded.score_reason,
last_ranked_at = excluded.last_ranked_at,
last_scheduled_at = COALESCE(excluded.last_scheduled_at, pool_member_scores.last_scheduled_at),
last_success_at = COALESCE(excluded.last_success_at, pool_member_scores.last_success_at),
last_failure_at = COALESCE(excluded.last_failure_at, pool_member_scores.last_failure_at),
failure_count = excluded.failure_count,
last_probe_attempt_at = COALESCE(excluded.last_probe_attempt_at, pool_member_scores.last_probe_attempt_at),
last_probe_success_at = COALESCE(excluded.last_probe_success_at, pool_member_scores.last_probe_success_at),
last_probe_failure_at = COALESCE(excluded.last_probe_failure_at, pool_member_scores.last_probe_failure_at),
probe_failure_count = excluded.probe_failure_count,
probe_status = excluded.probe_status,
updated_at = excluded.updated_at
"#,
)
.bind(stored.id.as_str())
.bind(stored.pool_kind.as_str())
.bind(stored.pool_id.as_str())
.bind(stored.member_kind.as_str())
.bind(stored.member_id.as_str())
.bind(stored.capability.as_str())
.bind(stored.scope_kind.as_str())
.bind(stored.scope_id.as_deref())
.bind(stored.score)
.bind(stored.hard_state.as_database())
.bind(i64_from_u64(stored.score_version, "pool score version")?)
.bind(score_reason)
.bind(i64_opt_from_u64(stored.last_ranked_at, "pool score last_ranked_at")?)
.bind(i64_opt_from_u64(stored.last_scheduled_at, "pool score last_scheduled_at")?)
.bind(i64_opt_from_u64(stored.last_success_at, "pool score last_success_at")?)
.bind(i64_opt_from_u64(stored.last_failure_at, "pool score last_failure_at")?)
.bind(i64_from_u64(stored.failure_count, "pool score failure_count")?)
.bind(i64_opt_from_u64(
stored.last_probe_attempt_at,
"pool score last_probe_attempt_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_success_at,
"pool score last_probe_success_at",
)?)
.bind(i64_opt_from_u64(
stored.last_probe_failure_at,
"pool score last_probe_failure_at",
)?)
.bind(i64_from_u64(
stored.probe_failure_count,
"pool score probe_failure_count",
)?)
.bind(stored.probe_status.as_database())
.bind(i64_from_u64(stored.updated_at, "pool score updated_at")?)
.execute(&self.pool)
.await
.map_sql_err()?;
Ok(stored)
}
async fn mark_pool_member_probe_in_progress(
&self,
attempt: PoolMemberProbeAttempt,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&attempt.identity, attempt.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(attempt.attempted_at);
row.probe_status = PoolMemberProbeStatus::InProgress;
row.score_reason =
merge_score_reason_patch(row.score_reason, attempt.score_reason_patch.clone());
row.updated_at = attempt.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_probe_result(
&self,
result: PoolMemberProbeResult,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&result.identity, result.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_probe_attempt_at = Some(result.attempted_at);
row.probe_status = result.probe_status;
if result.succeeded {
row.last_probe_success_at = Some(result.attempted_at);
row.probe_failure_count = 0;
} else {
row.last_probe_failure_at = Some(result.attempted_at);
row.probe_failure_count = row.probe_failure_count.saturating_add(1);
}
if let Some(hard_state) = result.hard_state {
row.hard_state = hard_state;
}
row.score_reason =
merge_score_reason_patch(row.score_reason, result.score_reason_patch.clone());
row.updated_at = result.attempted_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn record_pool_member_schedule_feedback(
&self,
feedback: PoolMemberScheduleFeedback,
) -> Result<usize, DataLayerError> {
let rows = self
.find_scores_by_identity(&feedback.identity, feedback.scope.as_ref())
.await?;
let count = rows.len();
for mut row in rows {
row.last_scheduled_at = Some(feedback.scheduled_at);
match feedback.succeeded {
Some(true) => row.last_success_at = Some(feedback.scheduled_at),
Some(false) => {
row.last_failure_at = Some(feedback.scheduled_at);
row.failure_count = row.failure_count.saturating_add(1);
}
None => {}
}
if let Some(hard_state) = feedback.hard_state {
row.hard_state = hard_state;
}
row.score = score_with_delta(row.score, feedback.score_delta);
row.score_reason =
merge_score_reason_patch(row.score_reason, feedback.score_reason_patch.clone());
row.updated_at = feedback.scheduled_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn mark_pool_member_hard_state(
&self,
identity: &PoolMemberIdentity,
scope: Option<&PoolScoreScope>,
hard_state: PoolMemberHardState,
updated_at: u64,
) -> Result<usize, DataLayerError> {
let rows = self.find_scores_by_identity(identity, scope).await?;
let count = rows.len();
for mut row in rows {
row.hard_state = hard_state;
row.updated_at = updated_at;
self.upsert_pool_member_score(upsert_from_stored(row))
.await?;
}
Ok(count)
}
async fn delete_pool_member_scores_for_member(
&self,
identity: &PoolMemberIdentity,
) -> Result<usize, DataLayerError> {
let result = sqlx::query(
r#"
DELETE FROM pool_member_scores
WHERE pool_kind = ? AND pool_id = ? AND member_kind = ? AND member_id = ?
"#,
)
.bind(identity.pool_kind.as_str())
.bind(identity.pool_id.as_str())
.bind(identity.member_kind.as_str())
.bind(identity.member_id.as_str())
.execute(&self.pool)
.await
.map_sql_err()?;
Ok(result.rows_affected() as usize)
}
}
fn map_score_row(row: &SqliteRow) -> Result<StoredPoolMemberScore, DataLayerError> {
let score_reason_raw: String = row.try_get("score_reason").map_sql_err()?;
Ok(StoredPoolMemberScore {
id: row.try_get("id").map_sql_err()?,
pool_kind: row.try_get("pool_kind").map_sql_err()?,
pool_id: row.try_get("pool_id").map_sql_err()?,
member_kind: row.try_get("member_kind").map_sql_err()?,
member_id: row.try_get("member_id").map_sql_err()?,
capability: row.try_get("capability").map_sql_err()?,
scope_kind: row.try_get("scope_kind").map_sql_err()?,
scope_id: row.try_get("scope_id").map_sql_err()?,
score: row.try_get("score").map_sql_err()?,
hard_state: PoolMemberHardState::from_database(
row.try_get::<String, _>("hard_state")
.map_sql_err()?
.as_str(),
)?,
score_version: u64_from_i64(
row.try_get("score_version").map_sql_err()?,
"pool_member_scores.score_version",
)?,
score_reason: serde_json::from_str(&score_reason_raw).unwrap_or(serde_json::Value::Null),
last_ranked_at: u64_opt_from_i64(
row.try_get("last_ranked_at").map_sql_err()?,
"pool_member_scores.last_ranked_at",
)?,
last_scheduled_at: u64_opt_from_i64(
row.try_get("last_scheduled_at").map_sql_err()?,
"pool_member_scores.last_scheduled_at",
)?,
last_success_at: u64_opt_from_i64(
row.try_get("last_success_at").map_sql_err()?,
"pool_member_scores.last_success_at",
)?,
last_failure_at: u64_opt_from_i64(
row.try_get("last_failure_at").map_sql_err()?,
"pool_member_scores.last_failure_at",
)?,
failure_count: u64_from_i64(
row.try_get("failure_count").map_sql_err()?,
"pool_member_scores.failure_count",
)?,
last_probe_attempt_at: u64_opt_from_i64(
row.try_get("last_probe_attempt_at").map_sql_err()?,
"pool_member_scores.last_probe_attempt_at",
)?,
last_probe_success_at: u64_opt_from_i64(
row.try_get("last_probe_success_at").map_sql_err()?,
"pool_member_scores.last_probe_success_at",
)?,
last_probe_failure_at: u64_opt_from_i64(
row.try_get("last_probe_failure_at").map_sql_err()?,
"pool_member_scores.last_probe_failure_at",
)?,
probe_failure_count: u64_from_i64(
row.try_get("probe_failure_count").map_sql_err()?,
"pool_member_scores.probe_failure_count",
)?,
probe_status: PoolMemberProbeStatus::from_database(
row.try_get::<String, _>("probe_status")
.map_sql_err()?
.as_str(),
)?,
updated_at: u64_from_i64(
row.try_get("updated_at").map_sql_err()?,
"pool_member_scores.updated_at",
)?,
})
}
fn upsert_from_stored(score: StoredPoolMemberScore) -> UpsertPoolMemberScore {
UpsertPoolMemberScore {
id: score.id,
identity: PoolMemberIdentity {
pool_kind: score.pool_kind,
pool_id: score.pool_id,
member_kind: score.member_kind,
member_id: score.member_id,
},
scope: PoolScoreScope {
capability: score.capability,
scope_kind: score.scope_kind,
scope_id: score.scope_id,
},
score: score.score,
hard_state: score.hard_state,
score_version: score.score_version,
score_reason: score.score_reason,
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,
updated_at: score.updated_at,
}
}
fn i64_from_usize(value: usize, field: &str) -> Result<i64, DataLayerError> {
i64::try_from(value)
.map_err(|_| DataLayerError::InvalidInput(format!("{field} exceeds signed 64-bit range")))
}

92
frontend/src/api/endpoints/pool.ts
View File

@@ -192,6 +192,66 @@ export interface PoolKeysPageResponse {
keys: PoolKeyDetail[]
}
export type PoolScoreHardState =
| 'available'
| 'unknown'
| 'cooldown'
| 'quota_exhausted'
| 'auth_invalid'
| 'banned'
| 'inactive'
export type PoolScoreProbeStatus = 'never' | 'ok' | 'failed' | 'stale' | 'in_progress'
export interface PoolScoreKeySummary {
id: string
name: string
auth_type: string
is_active: boolean
internal_priority: number
last_used_at: number | null
}
export interface PoolMemberScoreItem {
id: string
pool_kind: string
pool_id: string
member_kind: string
member_id: string
capability: string
scope_kind: string
scope_id: string | null
score: number
hard_state: PoolScoreHardState
score_version: number
score_reason: Record<string, unknown> | null
last_ranked_at: number | null
last_scheduled_at: number | null
last_success_at: number | null
last_failure_at: number | null
failure_count: number
last_probe_attempt_at: number | null
last_probe_success_at: number | null
last_probe_failure_at: number | null
probe_failure_count: number
probe_status: PoolScoreProbeStatus
updated_at: number
key?: PoolScoreKeySummary | null
}
export interface PoolScoresResponse {
provider_id: string
page: number
page_size: number
filters: {
api_format?: string | null
model_id?: string | null
hard_state?: string | null
probe_status?: string | null
}
items: PoolMemberScoreItem[]
}
export interface PoolKeysQuery {
page?: number
page_size?: number
@@ -203,6 +263,15 @@ export interface PoolKeysQuery {
sort_order?: 'asc' | 'desc'
}
export interface PoolScoresQuery {
page?: number
page_size?: number
api_format?: string
model_id?: string
hard_state?: string
probe_status?: string
}
export interface PoolKeySelectionRequest {
search?: string
quick_selectors?: string[]
@@ -293,6 +362,29 @@ export async function listPoolKeys(
)
}
export async function listPoolScores(
providerId: string,
params: PoolScoresQuery = {},
options: PoolReadOptions = {},
): Promise<PoolScoresResponse> {
const normalizedParams = { ...params }
const cacheKey = buildCacheKey(
`pool:scores:${providerId}`,
normalizedParams as Record<string, unknown>,
)
return cachedRequest(
cacheKey,
async () => {
const response = await client.get<PoolScoresResponse>(
`/api/admin/pool/${providerId}/scores`,
{ params: normalizedParams },
)
return response.data
},
options.cacheTtlMs ?? 0,
)
}
export async function resolvePoolKeySelection(
providerId: string,
body: PoolKeySelectionRequest,

3
frontend/src/api/endpoints/types/provider.ts
View File

@@ -548,6 +548,9 @@ export interface PoolAdvancedConfig {
health_policy_enabled?: boolean
unschedulable_rules?: Array<Record<string, unknown>> | null
batch_concurrency?: number | null
probe_concurrency?: number | null
score_top_n?: number | null
score_fallback_scan_limit?: number | null
probing_enabled?: boolean
probing_interval_minutes?: number | null
auto_remove_banned_keys?: boolean

50
frontend/src/features/pool/components/PoolAdvancedDialog.vue
View File

@@ -230,7 +230,7 @@
</p>
</div>
<div class="rounded-xl bg-muted/30 p-4">
<div class="grid gap-3 rounded-xl bg-muted/30 p-4 sm:grid-cols-2">
<div class="space-y-1.5">
<Label>
并发数
@@ -247,6 +247,45 @@
为空时沿用默认值；数值越大，批量操作越快，但会增加瞬时请求压力。
</p>
</div>
<div class="space-y-1.5">
<Label>
探测并发
</Label>
<Input
:model-value="form.probe_concurrency ?? ''"
type="number"
min="1"
max="64"
placeholder="4"
@update:model-value="(v) => form.probe_concurrency = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>
评分 Top-N
</Label>
<Input
:model-value="form.score_top_n ?? ''"
type="number"
min="1"
max="4096"
placeholder="128"
@update:model-value="(v) => form.score_top_n = parseNum(v)"
/>
</div>
<div class="space-y-1.5">
<Label>
回退扫描
</Label>
<Input
:model-value="form.score_fallback_scan_limit ?? ''"
type="number"
min="1"
max="100000"
placeholder="1024"
@update:model-value="(v) => form.score_fallback_scan_limit = parseNum(v)"
/>
</div>
</div>
</section>
</div>
@@ -464,6 +503,9 @@ const form = ref({
cost_limit_per_key_tokens: null as number | null | undefined,
cost_soft_threshold_percent: null as number | null | undefined,
batch_concurrency: null as number | null | undefined,
probe_concurrency: null as number | null | undefined,
score_top_n: null as number | null | undefined,
score_fallback_scan_limit: null as number | null | undefined,
probing_enabled: false,
probing_interval_minutes: null as number | null | undefined,
auto_remove_banned_keys: false,
@@ -539,6 +581,9 @@ watch(() => props.modelValue, (open) => {
cost_limit_per_key_tokens: cfg?.cost_limit_per_key_tokens ?? null,
cost_soft_threshold_percent: cfg?.cost_soft_threshold_percent ?? null,
batch_concurrency: cfg?.batch_concurrency ?? null,
probe_concurrency: cfg?.probe_concurrency ?? null,
score_top_n: cfg?.score_top_n ?? null,
score_fallback_scan_limit: cfg?.score_fallback_scan_limit ?? null,
probing_enabled: cfg?.probing_enabled ?? false,
probing_interval_minutes: cfg?.probing_interval_minutes ?? null,
auto_remove_banned_keys: cfg?.auto_remove_banned_keys ?? false,
@@ -572,6 +617,9 @@ async function handleSave() {
overload_cooldown_seconds: form.value.overload_cooldown_seconds ?? undefined,
health_policy_enabled: form.value.health_policy_enabled,
batch_concurrency: form.value.batch_concurrency ?? undefined,
probe_concurrency: form.value.probe_concurrency ?? undefined,
score_top_n: form.value.score_top_n ?? undefined,
score_fallback_scan_limit: form.value.score_fallback_scan_limit ?? undefined,
probing_enabled: form.value.probing_enabled,
probing_interval_minutes: form.value.probing_enabled
? (form.value.probing_interval_minutes ?? undefined)

286
frontend/src/views/admin/PoolManagement.vue
View File

@@ -1189,6 +1189,170 @@
</template>
</Card>
<Card
v-if="selectedProviderId"
variant="default"
class="overflow-hidden"
>
<div class="flex flex-col gap-3 border-b border-border/60 px-4 py-3 sm:flex-row sm:items-center sm:justify-between sm:px-6 sm:py-3.5">
<div class="min-w-0">
<h3 class="text-base font-semibold">
评分排序
</h3>
<p class="mt-1 text-xs text-muted-foreground">
Top {{ poolScorePage.items.length }} / page {{ poolScorePage.page || 1 }}
</p>
</div>
<RefreshButton
:loading="scoresLoading"
title="刷新评分"
@click="() => loadScores()"
/>
</div>
<div class="grid gap-2 border-b border-border/60 bg-muted/20 px-4 py-3 sm:grid-cols-4 sm:px-6">
<Input
v-model="scoreApiFormatFilter"
class="h-8 text-xs"
placeholder="api_format"
/>
<Input
v-model="scoreModelIdFilter"
class="h-8 text-xs"
placeholder="model_id"
/>
<Select v-model="scoreHardStateFilter">
<SelectTrigger class="h-8 text-xs border-border/60">
<SelectValue placeholder="硬状态" />
</SelectTrigger>
<SelectContent>
<SelectItem
v-for="item in poolScoreHardStateOptions"
:key="item.value"
:value="item.value"
>
{{ item.label }}
</SelectItem>
</SelectContent>
</Select>
<Select v-model="scoreProbeStatusFilter">
<SelectTrigger class="h-8 text-xs border-border/60">
<SelectValue placeholder="探测" />
</SelectTrigger>
<SelectContent>
<SelectItem
v-for="item in poolScoreProbeStatusOptions"
:key="item.value"
:value="item.value"
>
{{ item.label }}
</SelectItem>
</SelectContent>
</Select>
</div>
<div
v-if="scoresLoading && poolScorePage.items.length === 0"
class="flex items-center justify-center py-10"
>
<div class="animate-spin rounded-full h-7 w-7 border-b-2 border-primary" />
</div>
<div
v-else-if="poolScorePage.items.length === 0"
class="flex flex-col items-center justify-center py-10 text-center"
>
<div class="mx-auto flex h-12 w-12 items-center justify-center rounded-full bg-muted">
<SlidersHorizontal class="h-6 w-6 text-muted-foreground" />
</div>
<p class="mt-3 text-sm text-muted-foreground">
{{ hasPoolScoreFilters ? '未找到匹配评分' : '暂无评分数据' }}
</p>
</div>
<div
v-else
class="overflow-x-auto"
>
<Table class="w-full min-w-[760px] table-fixed">
<TableHeader>
<TableRow class="border-b border-border/60 hover:bg-transparent">
<TableHead class="w-[24%] font-semibold">
账号
</TableHead>
<TableHead class="w-[24%] font-semibold">
范围
</TableHead>
<TableHead class="w-[12%] text-center font-semibold">
分数
</TableHead>
<TableHead class="w-[12%] text-center font-semibold">
状态
</TableHead>
<TableHead class="w-[12%] text-center font-semibold">
探测
</TableHead>
<TableHead class="w-[16%] text-center font-semibold">
更新时间
</TableHead>
</TableRow>
</TableHeader>
<TableBody>
<TableRow
v-for="score in poolScorePage.items"
:key="score.id"
class="border-b border-border/40 last:border-b-0 hover:bg-muted/30"
>
<TableCell class="py-3">
<div class="min-w-0">
<div class="truncate text-sm">
{{ score.key?.name || score.member_id }}
</div>
<div class="truncate font-mono text-[10px] text-muted-foreground">
{{ score.member_id }}
</div>
</div>
</TableCell>
<TableCell class="py-3">
<div class="min-w-0 text-xs">
<div class="truncate text-foreground/90">
{{ score.capability || '-' }}
</div>
<div class="truncate text-[10px] text-muted-foreground">
{{ score.scope_kind }}: {{ score.scope_id || '*' }}
</div>
</div>
</TableCell>
<TableCell class="py-3 text-center">
<span class="font-mono text-xs tabular-nums">
{{ formatPoolScore(score.score) }}
</span>
</TableCell>
<TableCell class="py-3 text-center">
<Badge
variant="outline"
class="text-[10px]"
>
{{ getPoolScoreHardStateLabel(score.hard_state) }}
</Badge>
</TableCell>
<TableCell class="py-3 text-center">
<Badge
variant="secondary"
class="text-[10px]"
>
{{ getPoolScoreProbeStatusLabel(score.probe_status) }}
</Badge>
</TableCell>
<TableCell class="py-3 text-center">
<span class="text-[10px] text-muted-foreground">
{{ formatUnixSeconds(score.updated_at) }}
</span>
</TableCell>
</TableRow>
</TableBody>
</Table>
</div>
</Card>
<!-- Dialogs -->
<OAuthAccountDialog
v-if="selectedProviderId"
@@ -1309,6 +1473,7 @@ import {
getPoolOverview,
getPoolSchedulingPresets,
listPoolKeys,
listPoolScores,
clearPoolCooldown,
} from '@/api/endpoints/pool'
import {
@@ -1324,7 +1489,9 @@ import type {
PoolOverviewItem,
PoolKeyDetail,
PoolKeysPageResponse,
PoolMemberScoreItem,
PoolPresetMeta,
PoolScoresResponse,
} from '@/api/endpoints/pool'
import type {
ClaudeCodeAdvancedConfig,
@@ -1420,11 +1587,14 @@ let overviewRequestId = 0
let selectProviderRequestId = 0
let providerDataRequestId = 0
let keysRequestId = 0
let scoresRequestId = 0
let keysSearchDebounceTimer: number | null = null
let scoresFilterDebounceTimer: number | null = null
let suppressFiltersWatch = false
let hasHydratedInitialProviderSelection = false
const POOL_OVERVIEW_CACHE_TTL_MS = 10 * 1000
const POOL_KEYS_CACHE_TTL_MS = 10 * 1000
const POOL_SCORES_CACHE_TTL_MS = 10 * 1000
const POOL_SCHEDULING_PRESETS_CACHE_TTL_MS = 5 * 60 * 1000
const poolKeyStatusFilterOptions: Array<{ value: PoolManagementViewState['status'], label: string }> = [
{ value: 'all', label: '全部状态' },
@@ -1432,6 +1602,24 @@ const poolKeyStatusFilterOptions: Array<{ value: PoolManagementViewState['status
{ value: 'cooldown', label: '冷却中' },
{ value: 'inactive', label: '禁用' },
]
const poolScoreHardStateOptions = [
{ value: 'all', label: '全部状态' },
{ value: 'available', label: '可用' },
{ value: 'unknown', label: '未知' },
{ value: 'cooldown', label: '冷却' },
{ value: 'quota_exhausted', label: '额度耗尽' },
{ value: 'auth_invalid', label: '授权无效' },
{ value: 'banned', label: '封禁' },
{ value: 'inactive', label: '禁用' },
]
const poolScoreProbeStatusOptions = [
{ value: 'all', label: '全部探测' },
{ value: 'never', label: '未探测' },
{ value: 'ok', label: '正常' },
{ value: 'failed', label: '失败' },
{ value: 'stale', label: '过期' },
{ value: 'in_progress', label: '探测中' },
]
async function loadOverview(options: { cacheTtlMs?: number } = {}) {
const requestId = ++overviewRequestId
@@ -1488,6 +1676,7 @@ async function loadOverview(options: { cacheTtlMs?: number } = {}) {
showAccountBatchDialog.value = false
closeProviderProxyPopovers()
resetKeyPage()
resetScorePage()
}
} catch (err) {
if (requestId !== overviewRequestId) return
@@ -1720,9 +1909,11 @@ async function selectProvider(
}
keysLoadedOnce.value = false
resetKeyPage(currentPage.value, pageSize.value)
resetScorePage()
const keysTask = loadKeys({ cacheTtlMs: options.cacheTtlMs ?? 0 })
// Provider summary is non-blocking for key list rendering.
void loadProviderData(id)
void loadScores({ cacheTtlMs: options.cacheTtlMs ? POOL_SCORES_CACHE_TTL_MS : 0 })
await keysTask
if (requestId !== selectProviderRequestId) return
}
@@ -1748,18 +1939,40 @@ function createEmptyKeyPage(page = 1, pageSizeValue = 50): PoolKeysPageResponse
return { total: 0, page, page_size: pageSizeValue, keys: [] }
}
function createEmptyScorePage(): PoolScoresResponse {
return {
provider_id: selectedProviderId.value || '',
page: 1,
page_size: 20,
filters: {},
items: [],
}
}
const keyPage = ref<PoolKeysPageResponse>(createEmptyKeyPage())
const poolScorePage = ref<PoolScoresResponse>(createEmptyScorePage())
const keysLoading = ref(false)
const scoresLoading = ref(false)
const keysLoadedOnce = ref(false)
const refreshingCurrentPageQuota = ref(false)
const searchQuery = ref(restoredViewState.search)
const statusFilter = ref(restoredViewState.status)
const scoreApiFormatFilter = ref('')
const scoreModelIdFilter = ref('')
const scoreHardStateFilter = ref('all')
const scoreProbeStatusFilter = ref('all')
const currentPage = ref(restoredViewState.page)
const pageSize = ref(restoredViewState.pageSize)
const sortBy = ref<PoolManagementSortBy | null>(restoredViewState.sortBy)
const sortOrder = ref<PoolManagementSortOrder>(restoredViewState.sortOrder)
const poolStatsMode = ref<PoolManagementStatsMode>(restoredViewState.statsMode)
const hasPoolKeyFilters = computed(() => searchQuery.value.trim().length > 0 || statusFilter.value !== 'all')
const hasPoolScoreFilters = computed(() => {
return scoreApiFormatFilter.value.trim().length > 0
|| scoreModelIdFilter.value.trim().length > 0
|| scoreHardStateFilter.value !== 'all'
|| scoreProbeStatusFilter.value !== 'all'
})
const MANUAL_QUOTA_REFRESH_COOLDOWN_SECONDS = 5 * 60
const refreshingOAuthKeyId = ref<string | null>(null)
const resettingCycleKeyId = ref<string | null>(null)
@@ -2060,6 +2273,10 @@ function resetKeyPage(page = currentPage.value, pageSizeValue = pageSize.value):
keyPage.value = createEmptyKeyPage(page, pageSizeValue)
}
function resetScorePage(): void {
poolScorePage.value = createEmptyScorePage()
}
function refreshOverviewInBackground(): void {
void loadOverview()
}
@@ -2199,6 +2416,36 @@ async function refreshCurrentPage() {
if (!quotaDidReload) {
await refresh()
}
void loadScores()
}
async function loadScores(options: { cacheTtlMs?: number } = {}) {
if (!selectedProviderId.value) return
const requestId = ++scoresRequestId
const providerId = selectedProviderId.value
scoresLoading.value = true
try {
const nextPage = await listPoolScores(providerId, {
page: 1,
page_size: 20,
api_format: scoreApiFormatFilter.value.trim() || undefined,
model_id: scoreModelIdFilter.value.trim() || undefined,
hard_state: scoreHardStateFilter.value === 'all' ? undefined : scoreHardStateFilter.value,
probe_status: scoreProbeStatusFilter.value === 'all' ? undefined : scoreProbeStatusFilter.value,
}, {
cacheTtlMs: options.cacheTtlMs ?? 0,
})
if (requestId !== scoresRequestId || selectedProviderId.value !== providerId) return
poolScorePage.value = nextPage
} catch (err) {
if (requestId !== scoresRequestId || selectedProviderId.value !== providerId) return
resetScorePage()
showError(parseApiError(err, '读取评分失败'))
} finally {
if (requestId === scoresRequestId) {
scoresLoading.value = false
}
}
}
async function loadKeys(options: { cacheTtlMs?: number } = {}) {
@@ -2276,6 +2523,20 @@ watch(searchQuery, () => {
}, 300)
})
watch([scoreApiFormatFilter, scoreModelIdFilter], () => {
if (scoresFilterDebounceTimer !== null) {
clearTimeout(scoresFilterDebounceTimer)
}
scoresFilterDebounceTimer = window.setTimeout(() => {
scoresFilterDebounceTimer = null
void loadScores({ cacheTtlMs: POOL_SCORES_CACHE_TTL_MS })
}, 300)
})
watch([scoreHardStateFilter, scoreProbeStatusFilter], () => {
void loadScores({ cacheTtlMs: POOL_SCORES_CACHE_TTL_MS })
})
function normalizeAuthTypeForEdit(key: PoolKeyDetail): EndpointAPIKey['auth_type'] {
if (isOAuthManagedCredential(key)) return 'oauth'
if (isServiceAccountCredential(key)) return 'service_account'
@@ -3584,6 +3845,26 @@ function formatStatUsd(value: number | string | null | undefined): string {
return `$${n.toLocaleString('en-US', { minimumFractionDigits: 2, maximumFractionDigits: 2 })}`
}
function formatPoolScore(value: number): string {
const n = Number(value)
if (!Number.isFinite(n)) return '0.000'
return n.toFixed(3)
}
function getPoolScoreHardStateLabel(value: PoolMemberScoreItem['hard_state']): string {
return poolScoreHardStateOptions.find(item => item.value === value)?.label || value
}
function getPoolScoreProbeStatusLabel(value: PoolMemberScoreItem['probe_status']): string {
return poolScoreProbeStatusOptions.find(item => item.value === value)?.label || value
}
function formatUnixSeconds(seconds: number | null | undefined): string {
const raw = Number(seconds ?? 0)
if (!Number.isFinite(raw) || raw <= 0) return '-'
return formatRelativeTime(new Date(raw * 1000).toISOString())
}
function formatRelativeTime(isoStr: string): string {
const date = new Date(isoStr)
const pad = (n: number) => String(n).padStart(2, '0')
@@ -3611,9 +3892,14 @@ onBeforeUnmount(() => {
clearTimeout(keysSearchDebounceTimer)
keysSearchDebounceTimer = null
}
if (scoresFilterDebounceTimer !== null) {
clearTimeout(scoresFilterDebounceTimer)
scoresFilterDebounceTimer = null
}
overviewRequestId += 1
selectProviderRequestId += 1
providerDataRequestId += 1
keysRequestId += 1
scoresRequestId += 1
})
</script>