apr_cli/federation/

routing.rs

//! Router - Intelligent node selection for inference requests
//!
//! Combines catalog, health, and policy data to select the optimal
//! node for each request.

use super::catalog::ModelCatalog;
use super::health::{CircuitBreaker, HealthChecker};
use super::policy::CompositePolicy;
use super::traits::*;
use std::sync::Arc;
use std::time::Duration;

// ============================================================================
// Route Decision (exported type)
// ============================================================================

/// Final routing decision with reasoning
#[derive(Debug, Clone)]
pub struct RouteDecision {
    pub target: RouteTarget,
    pub alternatives: Vec<RouteTarget>,
    pub reasoning: String,
}

// ============================================================================
// Router Implementation
// ============================================================================

/// Configuration for the router
#[derive(Debug, Clone)]
pub struct RouterConfig {
    /// Maximum candidates to evaluate
    pub max_candidates: usize,
    /// Minimum score to be considered
    pub min_score: f64,
    /// Load balance strategy
    pub strategy: LoadBalanceStrategy,
}

impl Default for RouterConfig {
    fn default() -> Self {
        Self {
            max_candidates: 10,
            min_score: 0.1,
            strategy: LoadBalanceStrategy::LeastLatency,
        }
    }
}

/// The main router implementation
pub struct Router {
    config: RouterConfig,
    catalog: Arc<ModelCatalog>,
    health: Arc<HealthChecker>,
    circuit_breaker: Arc<CircuitBreaker>,
    policy: CompositePolicy,
}

impl Router {
    pub fn new(
        config: RouterConfig,
        catalog: Arc<ModelCatalog>,
        health: Arc<HealthChecker>,
        circuit_breaker: Arc<CircuitBreaker>,
    ) -> Self {
        Self {
            config,
            catalog,
            health,
            circuit_breaker,
            policy: CompositePolicy::enterprise_default(),
        }
    }

    /// Create router with custom policy
    #[must_use]
    pub fn with_policy(mut self, policy: CompositePolicy) -> Self {
        self.policy = policy;
        self
    }

    /// Build candidates from catalog and health data
    fn build_candidates(&self, capability: &Capability) -> Vec<RouteCandidate> {
        // This would be async in production, but for simplicity we use sync here
        // The actual routing is async via the trait

        let mut candidates = Vec::new();

        // Get nodes from catalog that support this capability
        let entries = self.catalog.all_entries();

        for entry in entries {
            // Check if model supports the capability
            let has_capability = entry.metadata.capabilities.iter().any(|c| c == capability);
            if !has_capability {
                continue;
            }

            for deployment in &entry.deployments {
                // Check circuit breaker
                if self.circuit_breaker.is_open(&deployment.node_id) {
                    continue;
                }

                // Get health status
                let health = self
                    .health
                    .get_cached_health(&deployment.node_id)
                    .unwrap_or_else(|| NodeHealth {
                        node_id: deployment.node_id.clone(),
                        status: HealthState::Unknown,
                        latency_p50: Duration::from_secs(1),
                        latency_p99: Duration::from_secs(5),
                        throughput: 0,
                        gpu_utilization: None,
                        queue_depth: 0,
                        last_check: std::time::Instant::now(),
                    });

                // Skip unhealthy nodes
                if health.status == HealthState::Unhealthy {
                    continue;
                }

                let target = RouteTarget {
                    node_id: deployment.node_id.clone(),
                    region_id: deployment.region_id.clone(),
                    endpoint: deployment.endpoint.clone(),
                    estimated_latency: health.latency_p50,
                    score: 0.0, // Will be calculated by policy
                };

                let health_score = match health.status {
                    HealthState::Healthy => 1.0,
                    HealthState::Degraded => 0.5,
                    HealthState::Unknown => 0.3,
                    HealthState::Unhealthy => 0.0,
                };

                let scores = RouteScores {
                    latency_score: 1.0 - (health.latency_p50.as_millis() as f64 / 5000.0).min(1.0),
                    throughput_score: (health.throughput as f64 / 1000.0).min(1.0),
                    cost_score: 0.5,     // Would come from region pricing
                    locality_score: 0.5, // Would be calculated based on request origin
                    health_score,
                    total: 0.0,
                };

                candidates.push(RouteCandidate {
                    target,
                    scores,
                    eligible: true,
                    rejection_reason: None,
                });
            }
        }

        candidates
    }

    /// Score and rank candidates
    fn rank_candidates(&self, candidates: &mut [RouteCandidate], request: &InferenceRequest) {
        for candidate in candidates.iter_mut() {
            // Check eligibility
            if !self.policy.is_eligible(candidate, request) {
                candidate.eligible = false;
                candidate.rejection_reason = Some("Policy rejected".to_string());
                continue;
            }

            // Calculate score
            let score = self.policy.score(candidate, request);
            candidate.target.score = score;
            candidate.scores.total = score;
        }

        // Sort by score descending
        candidates.sort_by(|a, b| {
            b.scores
                .total
                .partial_cmp(&a.scores.total)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
    }

    /// Select best candidate based on strategy
    fn select_best(&self, candidates: &[RouteCandidate]) -> Option<RouteCandidate> {
        let eligible: Vec<_> = candidates
            .iter()
            .filter(|c| c.eligible && c.scores.total >= self.config.min_score)
            .take(self.config.max_candidates)
            .collect();

        if eligible.is_empty() {
            return None;
        }

        match self.config.strategy {
            LoadBalanceStrategy::LeastLatency => {
                // Already sorted by score (which factors in latency)
                eligible.first().map(|c| (*c).clone())
            }
            LoadBalanceStrategy::LeastConnections => {
                // Would use queue_depth in production
                eligible.first().map(|c| (*c).clone())
            }
            LoadBalanceStrategy::RoundRobin => {
                // Would maintain counter state
                eligible.first().map(|c| (*c).clone())
            }
            LoadBalanceStrategy::WeightedRandom => {
                // Weighted random selection based on scores
                use std::collections::hash_map::DefaultHasher;
                use std::hash::{Hash, Hasher};

                let total_weight: f64 = eligible.iter().map(|c| c.scores.total).sum();
                if total_weight <= 0.0 {
                    return eligible.first().map(|c| (*c).clone());
                }

                // Simple pseudo-random using current time
                let mut hasher = DefaultHasher::new();
                std::time::SystemTime::now()
                    .duration_since(std::time::UNIX_EPOCH)
                    .unwrap_or_default()
                    .as_nanos()
                    .hash(&mut hasher);
                let random = (hasher.finish() as f64) / (u64::MAX as f64);

                let target = random * total_weight;
                let mut cumulative = 0.0;

                for candidate in &eligible {
                    cumulative += candidate.scores.total;
                    if cumulative >= target {
                        return Some((*candidate).clone());
                    }
                }

                eligible.last().map(|c| (*c).clone())
            }
            LoadBalanceStrategy::ConsistentHash => {
                // Would hash request ID to consistent bucket
                eligible.first().map(|c| (*c).clone())
            }
        }
    }
}

impl RouterTrait for Router {
    fn route(&self, request: &InferenceRequest) -> BoxFuture<'_, FederationResult<RouteTarget>> {
        // Clone request to avoid lifetime issues with async block
        let request = request.clone();

        Box::pin(async move {
            let mut candidates = self.build_candidates(&request.capability);

            if candidates.is_empty() {
                return Err(FederationError::NoCapacity(request.capability.clone()));
            }

            self.rank_candidates(&mut candidates, &request);

            self.select_best(&candidates)
                .map(|c| c.target)
                .ok_or_else(|| FederationError::AllNodesUnhealthy(request.capability.clone()))
        })
    }

    fn get_candidates(
        &self,
        request: &InferenceRequest,
    ) -> BoxFuture<'_, FederationResult<Vec<RouteCandidate>>> {
        // Clone request to avoid lifetime issues with async block
        let request = request.clone();

        Box::pin(async move {
            let mut candidates = self.build_candidates(&request.capability);
            self.rank_candidates(&mut candidates, &request);
            Ok(candidates)
        })
    }
}

// ============================================================================
// Builder for Router
// ============================================================================

/// Builder for creating routers with custom configuration
pub struct RouterBuilder {
    config: RouterConfig,
    catalog: Option<Arc<ModelCatalog>>,
    health: Option<Arc<HealthChecker>>,
    circuit_breaker: Option<Arc<CircuitBreaker>>,
    policy: Option<CompositePolicy>,
}

impl RouterBuilder {
    pub fn new() -> Self {
        Self {
            config: RouterConfig::default(),
            catalog: None,
            health: None,
            circuit_breaker: None,
            policy: None,
        }
    }

    #[must_use]
    pub fn config(mut self, config: RouterConfig) -> Self {
        self.config = config;
        self
    }

    #[must_use]
    pub fn catalog(mut self, catalog: Arc<ModelCatalog>) -> Self {
        self.catalog = Some(catalog);
        self
    }

    #[must_use]
    pub fn health(mut self, health: Arc<HealthChecker>) -> Self {
        self.health = Some(health);
        self
    }

    #[must_use]
    pub fn circuit_breaker(mut self, cb: Arc<CircuitBreaker>) -> Self {
        self.circuit_breaker = Some(cb);
        self
    }

    #[must_use]
    pub fn policy(mut self, policy: CompositePolicy) -> Self {
        self.policy = Some(policy);
        self
    }

    pub fn build(self) -> Router {
        let catalog = self
            .catalog
            .unwrap_or_else(|| Arc::new(ModelCatalog::new()));
        let health = self
            .health
            .unwrap_or_else(|| Arc::new(HealthChecker::default()));
        let circuit_breaker = self
            .circuit_breaker
            .unwrap_or_else(|| Arc::new(CircuitBreaker::default()));

        let router = Router::new(self.config, catalog, health, circuit_breaker);

        if let Some(policy) = self.policy {
            router.with_policy(policy)
        } else {
            router
        }
    }
}

impl Default for RouterBuilder {
    fn default() -> Self {
        Self::new()
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    fn setup_test_router() -> (Router, Arc<ModelCatalog>, Arc<HealthChecker>) {
        let catalog = Arc::new(ModelCatalog::new());
        let health = Arc::new(HealthChecker::default());
        let circuit_breaker = Arc::new(CircuitBreaker::default());

        let router = Router::new(
            RouterConfig::default(),
            Arc::clone(&catalog),
            Arc::clone(&health),
            circuit_breaker,
        );

        (router, catalog, health)
    }

    #[tokio::test]
    async fn test_route_no_nodes() {
        let (router, _, _) = setup_test_router();

        let request = InferenceRequest {
            capability: Capability::Transcribe,
            input: vec![],
            qos: QoSRequirements::default(),
            request_id: "test-1".to_string(),
            tenant_id: None,
        };

        let result = router.route(&request).await;
        assert!(matches!(result, Err(FederationError::NoCapacity(_))));
    }

    #[tokio::test]
    async fn test_route_single_node() {
        let (router, catalog, health) = setup_test_router();

        // Register a node
        catalog
            .register(
                ModelId("whisper".to_string()),
                NodeId("node-1".to_string()),
                RegionId("us-west".to_string()),
                vec![Capability::Transcribe],
            )
            .await
            .expect("registration failed");

        health.register_node(NodeId("node-1".to_string()));
        health.report_success(&NodeId("node-1".to_string()), Duration::from_millis(50));

        let request = InferenceRequest {
            capability: Capability::Transcribe,
            input: vec![],
            qos: QoSRequirements::default(),
            request_id: "test-2".to_string(),
            tenant_id: None,
        };

        let result = router.route(&request).await;
        assert!(result.is_ok());

        let target = result.expect("routing failed");
        assert_eq!(target.node_id, NodeId("node-1".to_string()));
    }

    #[tokio::test]
    async fn test_route_prefers_healthy() {
        let (router, catalog, health) = setup_test_router();

        // Register two nodes
        catalog
            .register(
                ModelId("llama".to_string()),
                NodeId("healthy-node".to_string()),
                RegionId("us-west".to_string()),
                vec![Capability::Generate],
            )
            .await
            .expect("registration failed");

        catalog
            .register(
                ModelId("llama".to_string()),
                NodeId("degraded-node".to_string()),
                RegionId("us-east".to_string()),
                vec![Capability::Generate],
            )
            .await
            .expect("registration failed");

        // Make one healthy, one degraded
        health.register_node(NodeId("healthy-node".to_string()));
        health.register_node(NodeId("degraded-node".to_string()));

        for _ in 0..5 {
            health.report_success(
                &NodeId("healthy-node".to_string()),
                Duration::from_millis(20),
            );
            health.report_failure(&NodeId("degraded-node".to_string()));
        }

        let request = InferenceRequest {
            capability: Capability::Generate,
            input: vec![],
            qos: QoSRequirements::default(),
            request_id: "test-3".to_string(),
            tenant_id: None,
        };

        let result = router.route(&request).await;
        assert!(result.is_ok());

        let target = result.expect("routing failed");
        assert_eq!(target.node_id, NodeId("healthy-node".to_string()));
    }

    #[tokio::test]
    async fn test_get_candidates_returns_all() {
        let (router, catalog, health) = setup_test_router();

        // Register multiple nodes
        for i in 0..3 {
            catalog
                .register(
                    ModelId("embed".to_string()),
                    NodeId(format!("node-{}", i)),
                    RegionId("us-west".to_string()),
                    vec![Capability::Embed],
                )
                .await
                .expect("registration failed");

            health.register_node(NodeId(format!("node-{}", i)));
            health.report_success(&NodeId(format!("node-{}", i)), Duration::from_millis(50));
        }

        let request = InferenceRequest {
            capability: Capability::Embed,
            input: vec![],
            qos: QoSRequirements::default(),
            request_id: "test-4".to_string(),
            tenant_id: None,
        };

        let candidates = router
            .get_candidates(&request)
            .await
            .expect("get_candidates failed");

        assert_eq!(candidates.len(), 3);
    }

    #[test]
    fn test_router_builder() {
        let router = RouterBuilder::new()
            .config(RouterConfig {
                max_candidates: 5,
                min_score: 0.2,
                strategy: LoadBalanceStrategy::RoundRobin,
            })
            .build();

        assert_eq!(router.config.max_candidates, 5);
        assert_eq!(router.config.min_score, 0.2);
    }

    // =========================================================================
    // RouterConfig tests
    // =========================================================================

    #[test]
    fn test_router_config_default() {
        let config = RouterConfig::default();
        assert_eq!(config.max_candidates, 10);
        assert_eq!(config.min_score, 0.1);
        assert!(matches!(config.strategy, LoadBalanceStrategy::LeastLatency));
    }

    #[test]
    fn test_router_config_clone() {
        let config = RouterConfig {
            max_candidates: 20,
            min_score: 0.5,
            strategy: LoadBalanceStrategy::ConsistentHash,
        };
        let cloned = config.clone();
        assert_eq!(cloned.max_candidates, 20);
        assert_eq!(cloned.min_score, 0.5);
    }

    // =========================================================================
    // RouteDecision tests
    // =========================================================================

    #[test]
    fn test_route_decision_construction() {
        let decision = RouteDecision {
            target: RouteTarget {
                node_id: NodeId("n1".to_string()),
                region_id: RegionId("r1".to_string()),
                endpoint: "http://n1:8080".to_string(),
                estimated_latency: Duration::from_millis(50),
                score: 0.95,
            },
            alternatives: vec![],
            reasoning: "best latency score".to_string(),
        };
        assert_eq!(decision.target.score, 0.95);
        assert!(decision.alternatives.is_empty());
    }

    #[test]
    fn test_route_decision_with_alternatives() {
        let primary = RouteTarget {
            node_id: NodeId("n1".to_string()),
            region_id: RegionId("us-west".to_string()),
            endpoint: "http://n1:8080".to_string(),
            estimated_latency: Duration::from_millis(50),
            score: 0.95,
        };
        let alt = RouteTarget {
            node_id: NodeId("n2".to_string()),
            region_id: RegionId("eu-west".to_string()),
            endpoint: "http://n2:8080".to_string(),
            estimated_latency: Duration::from_millis(120),
            score: 0.7,
        };
        let decision = RouteDecision {
            target: primary,
            alternatives: vec![alt],
            reasoning: "latency-based".to_string(),
        };
        assert_eq!(decision.alternatives.len(), 1);
    }

    // =========================================================================
    // Router with_policy tests
    // =========================================================================

    #[test]
    fn test_router_with_custom_policy() {
        let catalog = Arc::new(ModelCatalog::new());
        let health = Arc::new(HealthChecker::default());
        let circuit_breaker = Arc::new(CircuitBreaker::default());

        let custom_policy =
            CompositePolicy::new().with_policy(super::super::policy::LatencyPolicy::default());

        let router = Router::new(RouterConfig::default(), catalog, health, circuit_breaker)
            .with_policy(custom_policy);

        assert_eq!(router.config.max_candidates, 10);
    }

    // =========================================================================
    // RouterBuilder chaining tests
    // =========================================================================

    #[test]
    fn test_router_builder_default() {
        let builder = RouterBuilder::default();
        assert!(builder.catalog.is_none());
        assert!(builder.health.is_none());
        assert!(builder.circuit_breaker.is_none());
        assert!(builder.policy.is_none());
    }

    #[test]
    fn test_router_builder_with_catalog() {
        let catalog = Arc::new(ModelCatalog::new());
        let builder = RouterBuilder::new().catalog(Arc::clone(&catalog));
        assert!(builder.catalog.is_some());
    }

    #[test]
    fn test_router_builder_with_health() {
        let health = Arc::new(HealthChecker::default());
        let builder = RouterBuilder::new().health(health);
        assert!(builder.health.is_some());
    }

    #[test]
    fn test_router_builder_with_circuit_breaker() {
        let cb = Arc::new(CircuitBreaker::default());
        let builder = RouterBuilder::new().circuit_breaker(cb);
        assert!(builder.circuit_breaker.is_some());
    }

    #[test]
    fn test_router_builder_with_policy() {
        let policy = CompositePolicy::new();
        let builder = RouterBuilder::new().policy(policy);
        assert!(builder.policy.is_some());
    }

    #[test]
    fn test_router_builder_full_chain() {
        let catalog = Arc::new(ModelCatalog::new());
        let health = Arc::new(HealthChecker::default());
        let cb = Arc::new(CircuitBreaker::default());

        let router = RouterBuilder::new()
            .config(RouterConfig {
                max_candidates: 20,
                min_score: 0.05,
                strategy: LoadBalanceStrategy::WeightedRandom,
            })
            .catalog(catalog)
            .health(health)
            .circuit_breaker(cb)
            .policy(CompositePolicy::enterprise_default())
            .build();

        assert_eq!(router.config.max_candidates, 20);
        assert_eq!(router.config.min_score, 0.05);
    }

    #[test]
    fn test_router_builder_build_without_policy() {
        let router = RouterBuilder::new().build();
        // Should use enterprise default policy
        assert_eq!(router.config.max_candidates, 10);
    }

    // =========================================================================
    // build_candidates tests
    // =========================================================================

    #[tokio::test]
    async fn test_build_candidates_skips_circuit_open() {
        let catalog = Arc::new(ModelCatalog::new());
        let health = Arc::new(HealthChecker::default());
        let circuit_breaker = Arc::new(CircuitBreaker::default());

        // Register node
        catalog
            .register(
                ModelId("model".to_string()),
                NodeId("n1".to_string()),
                RegionId("us-west".to_string()),
                vec![Capability::Generate],
            )
            .await
            .expect("registration failed");

        health.register_node(NodeId("n1".to_string()));
        for _ in 0..3 {
            health.report_success(&NodeId("n1".to_string()), Duration::from_millis(50));
        }

        // Open circuit breaker for the node
        for _ in 0..5 {
            circuit_breaker.record_failure(&NodeId("n1".to_string()));
        }

        let router = Router::new(RouterConfig::default(), catalog, health, circuit_breaker);

        let request = InferenceRequest {
            capability: Capability::Generate,
            input: vec![],
            qos: QoSRequirements::default(),
            request_id: "test".to_string(),
            tenant_id: None,
        };

        // Should have no candidates since circuit is open
        let result = router.route(&request).await;
        assert!(matches!(result, Err(FederationError::NoCapacity(_))));
    }

    #[tokio::test]
    async fn test_build_candidates_skips_unhealthy() {
        let catalog = Arc::new(ModelCatalog::new());
        let health = Arc::new(HealthChecker::new(super::super::health::HealthConfig {
            failure_threshold: 2,
            ..Default::default()
        }));
        let circuit_breaker = Arc::new(CircuitBreaker::default());

        catalog
            .register(
                ModelId("model".to_string()),
                NodeId("unhealthy-node".to_string()),
                RegionId("us-west".to_string()),
                vec![Capability::Generate],
            )
            .await
            .expect("registration failed");

        health.register_node(NodeId("unhealthy-node".to_string()));
        // Make unhealthy
        for _ in 0..3 {
            health.report_failure(&NodeId("unhealthy-node".to_string()));
        }

        let router = Router::new(RouterConfig::default(), catalog, health, circuit_breaker);

        let request = InferenceRequest {
            capability: Capability::Generate,
            input: vec![],
            qos: QoSRequirements::default(),
            request_id: "test".to_string(),
            tenant_id: None,
        };

        // Unhealthy nodes are skipped
        let result = router.route(&request).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_build_candidates_wrong_capability() {
        let catalog = Arc::new(ModelCatalog::new());
        let health = Arc::new(HealthChecker::default());
        let circuit_breaker = Arc::new(CircuitBreaker::default());

        catalog
            .register(
                ModelId("whisper".to_string()),
                NodeId("n1".to_string()),
                RegionId("us-west".to_string()),
                vec![Capability::Transcribe],
            )
            .await
            .expect("registration failed");

        let router = Router::new(RouterConfig::default(), catalog, health, circuit_breaker);

        let request = InferenceRequest {
            capability: Capability::Generate, // Different capability
            input: vec![],
            qos: QoSRequirements::default(),
            request_id: "test".to_string(),
            tenant_id: None,
        };

        let result = router.route(&request).await;
        assert!(matches!(result, Err(FederationError::NoCapacity(_))));
    }

    // =========================================================================
    // select_best strategy tests
    // =========================================================================

    #[tokio::test]
    async fn test_route_with_round_robin_strategy() {
        let catalog = Arc::new(ModelCatalog::new());
        let health = Arc::new(HealthChecker::default());
        let circuit_breaker = Arc::new(CircuitBreaker::default());

        catalog
            .register(
                ModelId("model".to_string()),
                NodeId("n1".to_string()),
                RegionId("us-west".to_string()),
                vec![Capability::Generate],
            )
            .await
            .expect("registration failed");

        health.register_node(NodeId("n1".to_string()));
        for _ in 0..3 {
            health.report_success(&NodeId("n1".to_string()), Duration::from_millis(50));
        }

        let router = Router::new(
            RouterConfig {
                strategy: LoadBalanceStrategy::RoundRobin,
                ..Default::default()
            },
            catalog,
            health,
            circuit_breaker,
        );

        let request = InferenceRequest {
            capability: Capability::Generate,
            input: vec![],
            qos: QoSRequirements::default(),
            request_id: "test".to_string(),
            tenant_id: None,
        };

        let result = router.route(&request).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_route_with_consistent_hash_strategy() {
        let catalog = Arc::new(ModelCatalog::new());
        let health = Arc::new(HealthChecker::default());
        let circuit_breaker = Arc::new(CircuitBreaker::default());

        catalog
            .register(
                ModelId("model".to_string()),
                NodeId("n1".to_string()),
                RegionId("us-west".to_string()),
                vec![Capability::Embed],
            )
            .await
            .expect("registration failed");

        health.register_node(NodeId("n1".to_string()));
        for _ in 0..3 {
            health.report_success(&NodeId("n1".to_string()), Duration::from_millis(50));
        }

        let router = Router::new(
            RouterConfig {
                strategy: LoadBalanceStrategy::ConsistentHash,
                ..Default::default()
            },
            catalog,
            health,
            circuit_breaker,
        );

        let request = InferenceRequest {
            capability: Capability::Embed,
            input: vec![],
            qos: QoSRequirements::default(),
            request_id: "test".to_string(),
            tenant_id: None,
        };

        let result = router.route(&request).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_route_with_weighted_random_strategy() {
        let catalog = Arc::new(ModelCatalog::new());
        let health = Arc::new(HealthChecker::default());
        let circuit_breaker = Arc::new(CircuitBreaker::default());

        catalog
            .register(
                ModelId("model".to_string()),
                NodeId("n1".to_string()),
                RegionId("us-west".to_string()),
                vec![Capability::Generate],
            )
            .await
            .expect("registration failed");

        health.register_node(NodeId("n1".to_string()));
        for _ in 0..3 {
            health.report_success(&NodeId("n1".to_string()), Duration::from_millis(50));
        }

        let router = Router::new(
            RouterConfig {
                strategy: LoadBalanceStrategy::WeightedRandom,
                ..Default::default()
            },
            catalog,
            health,
            circuit_breaker,
        );

        let request = InferenceRequest {
            capability: Capability::Generate,
            input: vec![],
            qos: QoSRequirements::default(),
            request_id: "test".to_string(),
            tenant_id: None,
        };

        let result = router.route(&request).await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_route_with_least_connections_strategy() {
        let catalog = Arc::new(ModelCatalog::new());
        let health = Arc::new(HealthChecker::default());
        let circuit_breaker = Arc::new(CircuitBreaker::default());

        catalog
            .register(
                ModelId("model".to_string()),
                NodeId("n1".to_string()),
                RegionId("us-west".to_string()),
                vec![Capability::Generate],
            )
            .await
            .expect("registration failed");

        health.register_node(NodeId("n1".to_string()));
        for _ in 0..3 {
            health.report_success(&NodeId("n1".to_string()), Duration::from_millis(50));
        }

        let router = Router::new(
            RouterConfig {
                strategy: LoadBalanceStrategy::LeastConnections,
                ..Default::default()
            },
            catalog,
            health,
            circuit_breaker,
        );

        let request = InferenceRequest {
            capability: Capability::Generate,
            input: vec![],
            qos: QoSRequirements::default(),
            request_id: "test".to_string(),
            tenant_id: None,
        };

        let result = router.route(&request).await;
        assert!(result.is_ok());
    }

    // =========================================================================
    // rank_candidates with policy rejection
    // =========================================================================

    #[tokio::test]
    async fn test_get_candidates_with_rejected() {
        let catalog = Arc::new(ModelCatalog::new());
        let health = Arc::new(HealthChecker::default());
        let circuit_breaker = Arc::new(CircuitBreaker::default());

        // Register node
        catalog
            .register(
                ModelId("model".to_string()),
                NodeId("n1".to_string()),
                RegionId("us-west".to_string()),
                vec![Capability::Generate],
            )
            .await
            .expect("registration failed");

        health.register_node(NodeId("n1".to_string()));
        for _ in 0..3 {
            health.report_success(&NodeId("n1".to_string()), Duration::from_millis(50));
        }

        // Use a policy that rejects based on privacy
        let policy = CompositePolicy::new().with_policy(
            super::super::policy::PrivacyPolicy::default()
                .with_region(RegionId("us-west".to_string()), PrivacyLevel::Public),
        );

        let router = Router::new(RouterConfig::default(), catalog, health, circuit_breaker)
            .with_policy(policy);

        // Request requires Confidential, but us-west is Public
        let request = InferenceRequest {
            capability: Capability::Generate,
            input: vec![],
            qos: QoSRequirements {
                privacy: PrivacyLevel::Confidential,
                ..Default::default()
            },
            request_id: "test".to_string(),
            tenant_id: None,
        };

        let candidates = router
            .get_candidates(&request)
            .await
            .expect("get_candidates failed");
        assert!(!candidates.is_empty());
        // Candidate should be marked as not eligible
        assert!(!candidates[0].eligible);
        assert!(candidates[0].rejection_reason.is_some());
    }

    // =========================================================================
    // min_score filtering
    // =========================================================================

    #[tokio::test]
    async fn test_route_min_score_filters_candidates() {
        let catalog = Arc::new(ModelCatalog::new());
        let health = Arc::new(HealthChecker::default());
        let circuit_breaker = Arc::new(CircuitBreaker::default());

        catalog
            .register(
                ModelId("model".to_string()),
                NodeId("n1".to_string()),
                RegionId("us-west".to_string()),
                vec![Capability::Generate],
            )
            .await
            .expect("registration failed");

        health.register_node(NodeId("n1".to_string()));
        // Leave at Unknown state (low health score)

        let router = Router::new(
            RouterConfig {
                min_score: 0.99, // Very high minimum
                ..Default::default()
            },
            catalog,
            health,
            circuit_breaker,
        );

        let request = InferenceRequest {
            capability: Capability::Generate,
            input: vec![],
            qos: QoSRequirements::default(),
            request_id: "test".to_string(),
            tenant_id: None,
        };

        let result = router.route(&request).await;
        // Might fail if score is below min_score
        // Unknown health = 0.3 score, which is below 0.99
        assert!(result.is_err());
    }
}