refactor: extract provider pool abstractions

crates/aether-pool-core/src/lib.rs

@@ -0,0 +1,17 @@
+mod scheduler;
+mod scoring;
+
+pub use scheduler::{
+    normalize_enabled_pool_presets, run_pool_scheduler, PoolCandidateFacts, PoolCandidateInput,
+    PoolCandidateOrchestration, PoolMemberSignals, PoolRuntimeState, PoolScheduledCandidate,
+    PoolSchedulerOutcome, PoolSchedulingConfig, PoolSchedulingPreset, PoolSkippedCandidate,
+    POOL_ACCOUNT_BLOCKED_SKIP_REASON, POOL_ACCOUNT_EXHAUSTED_SKIP_REASON,
+    POOL_COOLDOWN_SKIP_REASON, POOL_COST_LIMIT_REACHED_SKIP_REASON,
+};
+pub use scoring::{
+    probe_freshness_score, probe_freshness_score_with_ttl, score_pool_member,
+    score_pool_member_with_rules, PoolMemberScoreInput, PoolMemberScoreOutput,
+    PoolMemberScoreRules, PoolMemberScoreWeights, POOL_SCORE_VERSION,
+    PROBE_FAILURE_COOLDOWN_THRESHOLD, PROBE_FAILURE_PENALTY, PROBE_FRESHNESS_TTL_SECONDS,
+    REQUEST_FAILURE_PENALTY, UNSCHEDULABLE_SCORE_CAP,
+};

crates/aether-pool-core/src/scoring.rs

@@ -0,0 +1,540 @@
+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;
+pub const UNSCHEDULABLE_SCORE_CAP: f64 = 0.05;
+pub const PROBE_FAILURE_PENALTY: f64 = 0.05;
+pub const REQUEST_FAILURE_PENALTY: f64 = 0.005;
+pub const PROBE_FAILURE_COOLDOWN_THRESHOLD: u64 = 3;
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct PoolMemberScoreWeights {
+    pub manual_priority: f64,
+    pub health: f64,
+    pub probe_freshness: f64,
+    pub quota_remaining: f64,
+    pub latency: f64,
+    pub cost_lru: f64,
+}
+
+impl Default for PoolMemberScoreWeights {
+    fn default() -> Self {
+        Self {
+            manual_priority: 0.30,
+            health: 0.20,
+            probe_freshness: 0.15,
+            quota_remaining: 0.15,
+            latency: 0.10,
+            cost_lru: 0.10,
+        }
+    }
+}
+
+impl PoolMemberScoreWeights {
+    pub fn normalized(self) -> Self {
+        let sanitized = Self {
+            manual_priority: finite_non_negative(self.manual_priority),
+            health: finite_non_negative(self.health),
+            probe_freshness: finite_non_negative(self.probe_freshness),
+            quota_remaining: finite_non_negative(self.quota_remaining),
+            latency: finite_non_negative(self.latency),
+            cost_lru: finite_non_negative(self.cost_lru),
+        };
+        let total = sanitized.manual_priority
+            + sanitized.health
+            + sanitized.probe_freshness
+            + sanitized.quota_remaining
+            + sanitized.latency
+            + sanitized.cost_lru;
+        if total <= f64::EPSILON {
+            return sanitized;
+        }
+        Self {
+            manual_priority: sanitized.manual_priority / total,
+            health: sanitized.health / total,
+            probe_freshness: sanitized.probe_freshness / total,
+            quota_remaining: sanitized.quota_remaining / total,
+            latency: sanitized.latency / total,
+            cost_lru: sanitized.cost_lru / total,
+        }
+    }
+
+    fn as_reason_json(self) -> Value {
+        json!({
+            "manual_priority": self.manual_priority,
+            "health": self.health,
+            "probe_freshness": self.probe_freshness,
+            "quota_remaining": self.quota_remaining,
+            "latency": self.latency,
+            "cost_lru": self.cost_lru
+        })
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub struct PoolMemberScoreRules {
+    pub weights: PoolMemberScoreWeights,
+    pub probe_freshness_ttl_seconds: u64,
+    pub unschedulable_score_cap: f64,
+    pub probe_failure_penalty: f64,
+    pub request_failure_penalty: f64,
+    pub probe_failure_cooldown_threshold: u64,
+}
+
+impl Default for PoolMemberScoreRules {
+    fn default() -> Self {
+        Self {
+            weights: PoolMemberScoreWeights::default(),
+            probe_freshness_ttl_seconds: PROBE_FRESHNESS_TTL_SECONDS,
+            unschedulable_score_cap: UNSCHEDULABLE_SCORE_CAP,
+            probe_failure_penalty: PROBE_FAILURE_PENALTY,
+            request_failure_penalty: REQUEST_FAILURE_PENALTY,
+            probe_failure_cooldown_threshold: PROBE_FAILURE_COOLDOWN_THRESHOLD,
+        }
+    }
+}
+
+impl PoolMemberScoreRules {
+    pub fn effective(self) -> Self {
+        let defaults = Self::default();
+        Self {
+            weights: self.weights.normalized(),
+            probe_freshness_ttl_seconds: if self.probe_freshness_ttl_seconds == 0 {
+                defaults.probe_freshness_ttl_seconds
+            } else {
+                self.probe_freshness_ttl_seconds
+            },
+            unschedulable_score_cap: if self.unschedulable_score_cap.is_finite() {
+                self.unschedulable_score_cap.clamp(0.0, 1.0)
+            } else {
+                defaults.unschedulable_score_cap
+            },
+            probe_failure_penalty: if self.probe_failure_penalty.is_finite() {
+                self.probe_failure_penalty.clamp(0.0, 1.0)
+            } else {
+                defaults.probe_failure_penalty
+            },
+            request_failure_penalty: if self.request_failure_penalty.is_finite() {
+                self.request_failure_penalty.clamp(0.0, 1.0)
+            } else {
+                defaults.request_failure_penalty
+            },
+            probe_failure_cooldown_threshold: self.probe_failure_cooldown_threshold,
+        }
+    }
+}
+
+#[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_failure_count: 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 {
+    score_pool_member_with_rules(input, PoolMemberScoreRules::default())
+}
+
+pub fn score_pool_member_with_rules(
+    input: &PoolMemberScoreInput,
+    rules: PoolMemberScoreRules,
+) -> PoolMemberScoreOutput {
+    let rules = rules.effective();
+    let weights = rules.weights;
+    let hard_state = derive_hard_state(input, &rules);
+    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_with_ttl(
+        input.last_probe_success_at,
+        input.probe_status,
+        input.now_unix_secs,
+        rules.probe_freshness_ttl_seconds,
+    );
+    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 weighted_score = manual_priority * weights.manual_priority
+        + health * weights.health
+        + probe_freshness * weights.probe_freshness
+        + quota_remaining * weights.quota_remaining
+        + latency * weights.latency
+        + cost_lru * weights.cost_lru;
+    let probe_failure_penalty =
+        (input.probe_failure_count.min(10) as f64 * rules.probe_failure_penalty).min(0.5);
+    let request_failure_penalty =
+        (input.error_count.min(20) as f64 * rules.request_failure_penalty).min(0.5);
+    let total_penalty = (probe_failure_penalty + request_failure_penalty).min(1.0);
+    let mut score = weighted_score - total_penalty;
+    if !hard_state.schedulable() {
+        score = score.min(rules.unschedulable_score_cap);
+    }
+    score = score.clamp(0.0, 1.0);
+
+    PoolMemberScoreOutput {
+        score,
+        hard_state,
+        score_reason: json!({
+            "weights": weights.as_reason_json(),
+            "factors": {
+                "manual_priority": manual_priority,
+                "health": health,
+                "probe_freshness": probe_freshness,
+                "quota_remaining": quota_remaining,
+                "latency": latency,
+                "cost_lru": cost_lru
+            },
+            "rules": {
+                "probe_freshness_ttl_seconds": rules.probe_freshness_ttl_seconds,
+                "unschedulable_score_cap": rules.unschedulable_score_cap,
+                "probe_failure_penalty": rules.probe_failure_penalty,
+                "request_failure_penalty": rules.request_failure_penalty,
+                "probe_failure_cooldown_threshold": rules.probe_failure_cooldown_threshold
+            },
+            "penalties": {
+                "probe_failure": probe_failure_penalty,
+                "request_failure": request_failure_penalty,
+                "total": total_penalty
+            },
+            "hard_state": hard_state.as_database(),
+            "score_version": POOL_SCORE_VERSION
+        }),
+    }
+}
+
+fn finite_non_negative(value: f64) -> f64 {
+    if value.is_finite() {
+        value.max(0.0)
+    } else {
+        0.0
+    }
+}
+
+fn derive_hard_state(
+    input: &PoolMemberScoreInput,
+    rules: &PoolMemberScoreRules,
+) -> 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.probe_status == PoolMemberProbeStatus::Failed
+        && rules.probe_failure_cooldown_threshold > 0
+        && input.probe_failure_count >= rules.probe_failure_cooldown_threshold
+    {
+        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 {
+    probe_freshness_score_with_ttl(
+        last_probe_success_at,
+        probe_status,
+        now_unix_secs,
+        PROBE_FRESHNESS_TTL_SECONDS,
+    )
+}
+
+pub fn probe_freshness_score_with_ttl(
+    last_probe_success_at: Option<u64>,
+    probe_status: PoolMemberProbeStatus,
+    now_unix_secs: u64,
+    ttl_seconds: u64,
+) -> f64 {
+    if probe_status != PoolMemberProbeStatus::Ok {
+        return 0.0;
+    }
+    let Some(success_at) = last_probe_success_at else {
+        return 0.0;
+    };
+    let ttl_seconds = ttl_seconds.max(1);
+    let age = now_unix_secs.saturating_sub(success_at);
+    if age >= ttl_seconds {
+        0.0
+    } else {
+        1.0 - (age as f64 / 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_failure_count: 0,
+            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
+        );
+    }
+
+    #[test]
+    fn custom_rules_change_weights_and_probe_ttl() {
+        let mut input = input();
+        input.last_probe_success_at = Some(1_000);
+        input.now_unix_secs = 1_900;
+        let rules = PoolMemberScoreRules {
+            weights: PoolMemberScoreWeights {
+                manual_priority: 0.0,
+                health: 0.0,
+                probe_freshness: 1.0,
+                quota_remaining: 0.0,
+                latency: 0.0,
+                cost_lru: 0.0,
+            },
+            probe_freshness_ttl_seconds: 1_000,
+            unschedulable_score_cap: 0.05,
+            probe_failure_penalty: 0.0,
+            request_failure_penalty: 0.0,
+            probe_failure_cooldown_threshold: PROBE_FAILURE_COOLDOWN_THRESHOLD,
+        };
+
+        let output = score_pool_member_with_rules(&input, rules);
+
+        assert!((output.score - 0.1).abs() < 0.000_001);
+        assert_eq!(
+            output.score_reason["rules"]["probe_freshness_ttl_seconds"],
+            1_000
+        );
+        assert_eq!(output.score_reason["weights"]["probe_freshness"], 1.0);
+    }
+
+    #[test]
+    fn custom_rules_normalize_weights() {
+        let rules = PoolMemberScoreRules {
+            weights: PoolMemberScoreWeights {
+                manual_priority: 2.0,
+                health: 2.0,
+                probe_freshness: 0.0,
+                quota_remaining: 0.0,
+                latency: -1.0,
+                cost_lru: f64::NAN,
+            },
+            probe_freshness_ttl_seconds: 0,
+            unschedulable_score_cap: f64::INFINITY,
+            probe_failure_penalty: f64::NAN,
+            request_failure_penalty: f64::INFINITY,
+            probe_failure_cooldown_threshold: 2,
+        }
+        .effective();
+
+        assert_eq!(rules.weights.manual_priority, 0.5);
+        assert_eq!(rules.weights.health, 0.5);
+        assert_eq!(
+            rules.probe_freshness_ttl_seconds,
+            PROBE_FRESHNESS_TTL_SECONDS
+        );
+        assert_eq!(rules.unschedulable_score_cap, UNSCHEDULABLE_SCORE_CAP);
+        assert_eq!(rules.probe_failure_penalty, PROBE_FAILURE_PENALTY);
+        assert_eq!(rules.request_failure_penalty, REQUEST_FAILURE_PENALTY);
+        assert_eq!(rules.probe_failure_cooldown_threshold, 2);
+    }
+
+    #[test]
+    fn custom_rules_preserve_zero_weight_total() {
+        let rules = PoolMemberScoreRules {
+            weights: PoolMemberScoreWeights {
+                manual_priority: 0.0,
+                health: 0.0,
+                probe_freshness: 0.0,
+                quota_remaining: 0.0,
+                latency: 0.0,
+                cost_lru: 0.0,
+            },
+            ..PoolMemberScoreRules::default()
+        }
+        .effective();
+
+        assert_eq!(
+            rules.weights,
+            PoolMemberScoreWeights {
+                manual_priority: 0.0,
+                health: 0.0,
+                probe_freshness: 0.0,
+                quota_remaining: 0.0,
+                latency: 0.0,
+                cost_lru: 0.0,
+            }
+        );
+    }
+
+    #[test]
+    fn repeated_probe_failures_move_member_to_cooldown() {
+        let mut input = input();
+        input.probe_status = PoolMemberProbeStatus::Failed;
+        input.probe_failure_count = 3;
+
+        let output = score_pool_member(&input);
+
+        assert_eq!(output.hard_state, PoolMemberHardState::Cooldown);
+        assert!(output.score <= UNSCHEDULABLE_SCORE_CAP);
+        assert_eq!(
+            output.score_reason["rules"]["probe_failure_cooldown_threshold"],
+            PROBE_FAILURE_COOLDOWN_THRESHOLD
+        );
+    }
+
+    #[test]
+    fn probe_and_request_failures_penalize_schedulable_score() {
+        let mut input = input();
+        input.last_probe_success_at = None;
+        input.probe_status = PoolMemberProbeStatus::Never;
+        input.probe_failure_count = 1;
+        input.error_count = 2;
+        let rules = PoolMemberScoreRules {
+            probe_failure_penalty: 0.1,
+            request_failure_penalty: 0.01,
+            probe_failure_cooldown_threshold: 3,
+            ..PoolMemberScoreRules::default()
+        };
+
+        let output = score_pool_member_with_rules(&input, rules);
+
+        assert_eq!(output.hard_state, PoolMemberHardState::Available);
+        assert_eq!(output.score_reason["penalties"]["probe_failure"], 0.1);
+        assert_eq!(output.score_reason["penalties"]["request_failure"], 0.02);
+    }
+}