litellm-rs 0.9.0

//! Deployment selection logic
//!
//! This module contains the core routing logic for selecting
//! the best deployment for a given model.

use super::admission::{AdmissionBackend, AdmissionHold, AdmissionReserve};
use super::config::RoutingStrategy;
use super::deployment::{Deployment, DeploymentId, current_timestamp};
use super::error::RouterError;
use super::execution::router_error_to_provider_error;
use super::strategy_impl::{self, RoutingContext};
use super::unified::Router;
use super::{RoutingSnapshot, RuntimeHandle};
use crate::core::types::model::ProviderCapability;
use dashmap::DashMap;
use std::sync::Arc;
use std::sync::atomic::{
    AtomicUsize,
    Ordering::{AcqRel, Acquire, Relaxed},
};

/// Active-request reservation for a selected deployment.
///
/// Dropping the lease releases the selected deployment unless ownership was
/// explicitly converted back into the legacy deployment-id API.
#[derive(Debug)]
pub struct DeploymentLease {
    deployment: Arc<Deployment>,
    release_on_drop: bool,
    admission: AdmissionBackend,
    hold: Option<AdmissionHold>,
    async_cleanup: bool,
}

impl DeploymentLease {
    fn new(
        deployment: Arc<Deployment>,
        admission: AdmissionBackend,
        hold: Option<AdmissionHold>,
    ) -> Self {
        Self {
            deployment,
            release_on_drop: true,
            admission,
            hold,
            async_cleanup: true,
        }
    }

    #[cfg(all(test, feature = "gateway"))]
    pub(crate) fn commit_admission(&mut self, actual_tokens: u64) {
        if let Some(hold) = self.hold.take() {
            self.admission.settle(&hold, actual_tokens);
        }
    }

    pub(crate) fn preserve_admission_usage(&self, actual_tokens: u64) {
        if let Some(hold) = &self.hold {
            hold.prepare_settlement(actual_tokens);
        }
    }

    pub(crate) fn record_success(&self, tokens: u64, latency_us: u64) {
        self.deployment
            .record_success_with_admission(tokens, latency_us, self.hold.as_ref());
    }

    pub(crate) fn record_interrupted_usage(&self, tokens: u64) {
        self.deployment
            .record_interrupted_usage_with_admission(tokens, self.hold.as_ref());
    }

    pub(crate) fn preserve_admission_reservation(&self, minimum_tokens: u64) {
        if let Some(hold) = &self.hold {
            hold.prepare_retention(minimum_tokens);
        }
    }

    pub(crate) async fn retain_admission_async(&mut self, minimum_tokens: u64) {
        if let Some(hold) = self.hold.take() {
            self.admission.retain_async(&hold, minimum_tokens).await;
        }
    }

    #[cfg(feature = "gateway")]
    pub(crate) fn clone_admission_hold(&self) -> Option<AdmissionHold> {
        self.hold.clone()
    }

    pub(crate) async fn commit_admission_async(&mut self, actual_tokens: u64) {
        if let Some(hold) = self.hold.take() {
            self.admission.settle_async(&hold, actual_tokens).await;
        }
    }

    pub(crate) async fn cancel_admission_async(&mut self) {
        if let Some(hold) = self.hold.take() {
            self.admission.cancel_async(&hold).await;
        }
    }

    fn synchronous_cleanup(mut self) -> Self {
        self.async_cleanup = false;
        self
    }

    #[cfg(feature = "gateway")]
    pub(crate) fn take_admission(&mut self) -> (AdmissionBackend, Option<AdmissionHold>) {
        (self.admission.clone(), self.hold.take())
    }

    pub fn deployment_id(&self) -> &str {
        &self.deployment.id
    }

    pub fn deployment(&self) -> &Deployment {
        &self.deployment
    }

    pub fn clone_deployment_id(&self) -> DeploymentId {
        self.deployment_id().to_string()
    }

    pub fn clone_deployment(&self) -> Arc<Deployment> {
        self.deployment.clone()
    }

    /// Convert this lease into the legacy deployment-id API.
    ///
    /// Prefer keeping the lease alive so drop releases the exact snapshot
    /// deployment. Deprecated ID selectors retain only the local active count
    /// until `release_deployment`; shared admission requires an owned lease.
    pub fn into_deployment_id(mut self) -> DeploymentId {
        self.release_on_drop = false;
        if let Some(hold) = self.hold.take() {
            // The ID cannot carry a shared reservation's cleanup obligation.
            // Failed cancellation leaves the hold's Drop retry intact.
            self.admission.cancel(&hold);
        }
        self.deployment_id().to_string()
    }
}

impl Drop for DeploymentLease {
    fn drop(&mut self) {
        if self.release_on_drop {
            Router::release_selected_deployment(self.deployment());
            if !self.async_cleanup
                && let Some(hold) = self.hold.take()
            {
                self.admission.cancel(&hold);
            }
        }
    }
}

impl Router {
    /// Select the best deployment for a given model (core routing method)
    ///
    /// # Flow
    ///
    /// 1. Resolve model_name (handle aliases)
    /// 2. Get all deployment IDs for this model
    /// 3. Filter: healthy + not in cooldown + not rate limited
    /// 4. Select based on routing strategy
    /// 5. Increment active_requests counter
    #[deprecated(
        since = "0.5.0",
        note = "Use select_deployment_lease so the selected snapshot deployment is released by RAII"
    )]
    pub fn select_deployment(&self, model_name: &str) -> Result<DeploymentId, RouterError> {
        self.select_deployment_lease(model_name)
            .map(DeploymentLease::into_deployment_id)
    }

    /// Select and reserve the best deployment for a given model.
    ///
    /// The returned lease releases the deployment on drop, which protects
    /// router-owned execution paths from cancellation leaks.
    pub fn select_deployment_lease(
        &self,
        model_name: &str,
    ) -> Result<DeploymentLease, RouterError> {
        super::sync_compat::wait(self.select_deployment_lease_async(model_name))
            .map(DeploymentLease::synchronous_cleanup)
    }

    /// Select without blocking the caller's executor on shared Redis state.
    pub async fn select_deployment_lease_async(
        &self,
        model_name: &str,
    ) -> Result<DeploymentLease, RouterError> {
        let snapshot = self.load_routing_snapshot();
        self.select_deployment_matching(snapshot.as_ref(), model_name, |_| true, None, 0)
            .await
    }

    pub(crate) async fn select_deployment_lease_matching_in_snapshot<F>(
        &self,
        snapshot: &RoutingSnapshot,
        model_name: &str,
        is_candidate: F,
    ) -> Result<DeploymentLease, RouterError>
    where
        F: Fn(&Deployment) -> bool,
    {
        self.select_deployment_matching(snapshot, model_name, is_candidate, None, 0)
            .await
    }

    /// Select the best deployment for a model that supports `capability`.
    ///
    /// This uses the same health, cooldown, concurrency, rate-limit, and
    /// strategy filters as `select_deployment`, then reserves the selected
    /// deployment by incrementing its active request count.
    #[deprecated(
        since = "0.5.0",
        note = "Use select_deployment_lease_for_capability so the selected snapshot deployment is released by RAII"
    )]
    pub fn select_deployment_for_capability(
        &self,
        model_name: &str,
        capability: &ProviderCapability,
    ) -> Result<DeploymentId, RouterError> {
        self.select_deployment_lease_for_capability(model_name, capability)
            .map(DeploymentLease::into_deployment_id)
    }

    /// Select and reserve the best deployment for a model that supports
    /// `capability`.
    pub fn select_deployment_lease_for_capability(
        &self,
        model_name: &str,
        capability: &ProviderCapability,
    ) -> Result<DeploymentLease, RouterError> {
        self.select_deployment_lease_for_capability_matching(model_name, capability, |_| true)
    }

    /// Select and reserve the best deployment for a model that supports
    /// `capability` and passes an additional route-level predicate.
    pub fn select_deployment_lease_for_capability_matching<F>(
        &self,
        model_name: &str,
        capability: &ProviderCapability,
        is_candidate: F,
    ) -> Result<DeploymentLease, RouterError>
    where
        F: Fn(&Deployment) -> bool,
    {
        super::sync_compat::wait(self.select_deployment_lease_for_capability_matching_async(
            model_name,
            capability,
            is_candidate,
        ))
        .map(DeploymentLease::synchronous_cleanup)
    }

    /// Capability selection with asynchronous shared admission and circuit checks.
    pub async fn select_deployment_lease_for_capability_matching_async<F>(
        &self,
        model_name: &str,
        capability: &ProviderCapability,
        is_candidate: F,
    ) -> Result<DeploymentLease, RouterError>
    where
        F: Fn(&Deployment) -> bool,
    {
        let snapshot = self.load_routing_snapshot();
        self.select_deployment_lease_for_capability_matching_in_snapshot(
            snapshot.as_ref(),
            model_name,
            capability,
            is_candidate,
        )
        .await
    }

    pub(crate) async fn select_deployment_lease_for_capability_matching_in_snapshot<F>(
        &self,
        snapshot: &RoutingSnapshot,
        model_name: &str,
        capability: &ProviderCapability,
        is_candidate: F,
    ) -> Result<DeploymentLease, RouterError>
    where
        F: Fn(&Deployment) -> bool,
    {
        self.select_deployment_lease_for_capability_matching_with_estimate(
            snapshot,
            model_name,
            capability,
            is_candidate,
            0,
        )
        .await
    }

    pub(crate) async fn select_deployment_lease_for_capability_matching_with_estimate<F>(
        &self,
        snapshot: &RoutingSnapshot,
        model_name: &str,
        capability: &ProviderCapability,
        is_candidate: F,
        estimated_tokens: u64,
    ) -> Result<DeploymentLease, RouterError>
    where
        F: Fn(&Deployment) -> bool,
    {
        let resolved_name = snapshot.resolve_model_name(model_name);
        let candidates = snapshot
            .model_index
            .get(&resolved_name)
            .into_iter()
            .flatten()
            .filter_map(|id| snapshot.deployments.get(id));
        let (candidate_count, missing_mapping_count) =
            candidates.fold((0usize, 0usize), |(count, missing), deployment| {
                let missing_mapping = deployment
                    .provider
                    .deployment_model_identity()
                    .is_some_and(|identity| identity.capability_catalog_model().is_none());
                (count + 1, missing + usize::from(missing_mapping))
            });
        let no_matching_candidate_error = if candidate_count > 0
            && candidate_count == missing_mapping_count
        {
            RouterError::InvalidConfiguration(format!(
                "model '{model_name}' requires settings.model_identity_mappings.{resolved_name}.capability_catalog_model before capability routing"
            ))
        } else {
            RouterError::UnsupportedCapability {
                model: model_name.to_string(),
                capability: format!("{capability:?}"),
            }
        };

        self.select_deployment_matching(
            snapshot,
            model_name,
            |deployment| {
                deployment
                    .provider
                    .supports_capability_for_model(&deployment.model, capability)
                    && is_candidate(deployment)
            },
            Some(no_matching_candidate_error),
            estimated_tokens,
        )
        .await
    }

    /// Select a deployment ID from pre-built routing contexts.
    ///
    /// This is the shared routing-policy entry point used by the full
    /// `UnifiedRouter` and lightweight adapters such as the SDK client.
    pub fn select_from_routing_contexts<'id>(
        strategy: RoutingStrategy,
        model_name: &str,
        routing_contexts: &[RoutingContext<'id>],
        round_robin_counters: &DashMap<String, AtomicUsize>,
    ) -> Option<&'id DeploymentId> {
        match strategy {
            RoutingStrategy::SimpleShuffle => {
                strategy_impl::weighted_random_from_context(routing_contexts)
            }
            RoutingStrategy::LeastBusy => strategy_impl::least_busy_from_context(routing_contexts),
            RoutingStrategy::UsageBased => {
                strategy_impl::lowest_usage_from_context(routing_contexts)
            }
            RoutingStrategy::LatencyBased => {
                strategy_impl::lowest_latency_from_context(routing_contexts)
            }
            RoutingStrategy::PriorityBased => {
                strategy_impl::lowest_priority_from_context(routing_contexts)
            }
            RoutingStrategy::RateLimitAware => {
                strategy_impl::rate_limit_aware_from_context(routing_contexts)
            }
            RoutingStrategy::RoundRobin => strategy_impl::round_robin_from_context(
                model_name,
                routing_contexts,
                round_robin_counters,
            ),
        }
    }

    async fn select_deployment_matching<F>(
        &self,
        snapshot: &RoutingSnapshot,
        model_name: &str,
        is_candidate: F,
        no_matching_candidate_error: Option<RouterError>,
        estimated_tokens: u64,
    ) -> Result<DeploymentLease, RouterError>
    where
        F: Fn(&Deployment) -> bool,
    {
        // 1. Resolve model aliases and deployment indexes from one immutable
        // routing generation.
        let resolved_name = snapshot.resolve_model_name(model_name);

        // 2. Get all deployment IDs for this model.
        let deployment_ids_ref = snapshot
            .model_index
            .get(&resolved_name)
            .ok_or_else(|| RouterError::ModelNotFound(model_name.to_string()))?;

        if deployment_ids_ref.is_empty() {
            return Err(RouterError::ModelNotFound(model_name.to_string()));
        }

        // 3. Filter and build routing contexts in one pass to avoid cloning a
        // temporary candidate ID vector and then looking everything up again.
        let total_deployments = deployment_ids_ref.len();
        let mut existing_deployments = 0;
        let mut matching_deployments = 0;
        let mut routing_contexts = Vec::with_capacity(total_deployments);
        let now = current_timestamp();

        for id in deployment_ids_ref.iter() {
            let Some(deployment) = snapshot.deployments.get(id.as_str()) else {
                continue;
            };

            if deployment.model_name != resolved_name {
                tracing::debug!(
                    deployment_id = id.as_str(),
                    requested_model = %model_name,
                    resolved_model = %resolved_name,
                    deployment_model = %deployment.model_name,
                    "deployment index entry points at a different model"
                );
                continue;
            }
            existing_deployments += 1;

            if !is_candidate(deployment) {
                tracing::trace!(
                    deployment_id = id.as_str(),
                    model = %resolved_name,
                    reason = "capability_mismatch",
                    "deployment excluded from routing candidates"
                );
                continue;
            }
            matching_deployments += 1;

            // Shared circuit (or local cooldown/health) owns selection eligibility.
            if !self.deployment_is_selectable_async(deployment).await {
                tracing::trace!(
                    deployment_id = id.as_str(),
                    model = %resolved_name,
                    reason = "circuit_or_unhealthy",
                    "deployment excluded from routing candidates"
                );
                continue;
            }

            let active_requests = deployment.state.active_requests.load(Relaxed);
            if let Some(limit) = deployment.config.max_parallel_requests
                && active_requests >= limit
            {
                tracing::trace!(
                    deployment_id = id.as_str(),
                    model = %resolved_name,
                    reason = "parallel_limit_reached",
                    "deployment excluded from routing candidates"
                );
                continue;
            }

            let minute = deployment.state.minute_counters(now);
            let rpm_current = minute.rpm;
            if let Some(limit) = deployment.config.rpm_limit
                && rpm_current >= limit
            {
                tracing::trace!(
                    deployment_id = id.as_str(),
                    model = %resolved_name,
                    reason = "rate_limited",
                    "deployment excluded from routing candidates"
                );
                continue;
            }

            let tpm_current = deployment.state.admission_tpm(now);
            if let Some(limit) = deployment.config.tpm_limit
                && tpm_current >= limit
            {
                tracing::trace!(
                    deployment_id = id.as_str(),
                    model = %resolved_name,
                    reason = "rate_limited",
                    "deployment excluded from routing candidates"
                );
                continue;
            }

            routing_contexts.push(strategy_impl::RoutingContext {
                deployment_id: id,
                weight: deployment.config.weight,
                priority: deployment.config.priority,
                active_requests,
                tpm_current,
                tpm_limit: deployment.config.tpm_limit,
                rpm_current,
                rpm_limit: deployment.config.rpm_limit,
                avg_latency_us: deployment.state.avg_latency_us.load(Relaxed),
            });
        }

        if routing_contexts.is_empty() {
            if existing_deployments > 0
                && matching_deployments == 0
                && let Some(err) = no_matching_candidate_error
            {
                return Err(err);
            }

            tracing::warn!(
                model = %model_name,
                total_deployments = total_deployments,
                "no available deployments after filtering"
            );
            return Err(RouterError::NoAvailableDeployment(model_name.to_string()));
        }

        // 4. Select based on immutable routing contexts, then atomically reserve
        // the selected deployment. If another caller wins the race for the last
        // slot, remove that candidate and retry within this selection attempt.
        while !routing_contexts.is_empty() {
            let selected_id = Self::select_from_routing_contexts(
                self.config.routing_strategy,
                &resolved_name,
                &routing_contexts,
                &self.round_robin_counters,
            )
            .ok_or_else(|| RouterError::NoAvailableDeployment(model_name.to_string()))?
            .clone();

            let Some(deployment) = snapshot.deployments.get(&selected_id).cloned() else {
                continue;
            };

            if let Ok(hold) = self
                .try_reserve_deployment(&deployment, &resolved_name, estimated_tokens)
                .await
            {
                self.provider_selected_count.fetch_add(1, Relaxed);
                self.strategy_used_count.fetch_add(1, Relaxed);

                tracing::debug!(
                    model = %model_name,
                    strategy = ?self.config.routing_strategy,
                    candidate_count = routing_contexts.len(),
                    selected_id = %selected_id,
                    "deployment selected for routing"
                );

                return Ok(DeploymentLease::new(
                    deployment,
                    self.admission.clone(),
                    hold,
                ));
            }

            if let Some(pos) = routing_contexts
                .iter()
                .position(|ctx| ctx.deployment_id == &selected_id)
            {
                routing_contexts.swap_remove(pos);
            } else {
                break;
            }
        }

        tracing::warn!(
            model = %model_name,
            "no available deployments after reservation"
        );
        Err(RouterError::NoAvailableDeployment(model_name.to_string()))
    }

    async fn try_reserve_deployment(
        &self,
        deployment: &Deployment,
        expected_model: &str,
        estimated_tokens: u64,
    ) -> Result<Option<AdmissionHold>, RouterError> {
        if deployment.model_name != expected_model {
            return Err(RouterError::DeploymentNotFound(deployment.id.clone()));
        }

        let hold = match self
            .admission
            .reserve_async(deployment, estimated_tokens)
            .await
        {
            AdmissionReserve::Skipped => None,
            AdmissionReserve::Denied => {
                return Err(RouterError::RateLimitExceeded(expected_model.into()));
            }
            #[cfg(any(feature = "gateway", test))]
            AdmissionReserve::Unavailable => {
                return Err(RouterError::NoAvailableDeployment(expected_model.into()));
            }
            AdmissionReserve::Granted(hold) => Some(hold),
        };

        let local_ok = match deployment.config.max_parallel_requests {
            Some(limit) => {
                let mut current = deployment.state.active_requests.load(Acquire);
                loop {
                    if current >= limit {
                        break false;
                    }

                    match deployment.state.active_requests.compare_exchange_weak(
                        current,
                        current + 1,
                        AcqRel,
                        Acquire,
                    ) {
                        Ok(_) => break true,
                        Err(next) => current = next,
                    }
                }
            }
            None => {
                deployment.state.active_requests.fetch_add(1, Relaxed);
                true
            }
        };

        if local_ok {
            Ok(hold)
        } else {
            if let Some(hold) = hold {
                self.admission.cancel_async(&hold).await;
            }
            Err(RouterError::RateLimitExceeded(expected_model.into()))
        }
    }

    #[cfg(test)]
    pub(crate) fn select_deployment_lease_with_tokens(
        &self,
        model_name: &str,
        estimated_tokens: u64,
    ) -> Result<DeploymentLease, RouterError> {
        let snapshot = self.load_routing_snapshot();
        super::sync_compat::wait(self.select_deployment_matching(
            snapshot.as_ref(),
            model_name,
            |_| true,
            None,
            estimated_tokens,
        ))
        .map(DeploymentLease::synchronous_cleanup)
    }

    #[cfg(all(test, feature = "gateway", feature = "websockets"))]
    pub(crate) fn select_pinned_response_lease(
        &self,
        deployment: &Deployment,
        estimated_tokens: u64,
    ) -> Result<DeploymentLease, RouterError> {
        super::sync_compat::wait(
            self.select_pinned_response_lease_async(deployment, estimated_tokens),
        )
        .map(DeploymentLease::synchronous_cleanup)
    }

    #[cfg(all(feature = "gateway", feature = "websockets"))]
    pub(crate) async fn select_pinned_response_lease_async(
        &self,
        deployment: &Deployment,
        estimated_tokens: u64,
    ) -> Result<DeploymentLease, RouterError> {
        let snapshot = self.load_routing_snapshot();
        let current = snapshot
            .deployments
            .get(&deployment.id)
            .filter(|current| std::ptr::eq(current.as_ref(), deployment))
            .ok_or_else(|| RouterError::DeploymentNotFound(deployment.id.clone()))?;
        if !self.deployment_is_selectable_async(current).await {
            return Err(RouterError::NoAvailableDeployment(
                deployment.model_name.clone(),
            ));
        }
        let minute = current.state.minute_counters(current_timestamp());
        if current
            .config
            .rpm_limit
            .is_some_and(|limit| minute.rpm >= limit)
            || current
                .config
                .tpm_limit
                .is_some_and(|limit| current.state.admission_tpm(current_timestamp()) >= limit)
        {
            return Err(RouterError::RateLimitExceeded(
                deployment.model_name.clone(),
            ));
        }
        let hold = self
            .try_reserve_deployment(current, &deployment.model_name, estimated_tokens)
            .await?;
        self.provider_selected_count.fetch_add(1, Relaxed);
        self.strategy_used_count.fetch_add(1, Relaxed);
        Ok(DeploymentLease::new(
            current.clone(),
            self.admission.clone(),
            hold,
        ))
    }

    /// Release a deployment after request completion
    ///
    /// Decrements the active_requests counter for the deployment.
    #[deprecated(
        since = "0.5.0",
        note = "Use select_deployment_lease so release targets the selected snapshot deployment"
    )]
    pub fn release_deployment(&self, deployment_id: &str) {
        let snapshot = self.load_routing_snapshot();
        if let Some(deployment) = snapshot.deployments.get(deployment_id) {
            Self::release_selected_deployment(deployment);
        }
    }

    pub(crate) fn release_selected_deployment(deployment: &Deployment) {
        let result = deployment
            .state
            .active_requests
            .fetch_update(Relaxed, Relaxed, |v| Some(v.saturating_sub(1)));
        debug_assert!(result.is_ok());
    }
}

impl RuntimeHandle {
    /// Select and reserve a capable deployment from this handle's pinned generation.
    #[allow(dead_code)]
    pub(crate) fn select_deployment_lease_for_capability_typed(
        &self,
        model_name: &str,
        capability: &ProviderCapability,
    ) -> Result<DeploymentLease, crate::core::providers::ProviderError> {
        super::sync_compat::wait(
            self.binding
                .router
                .select_deployment_lease_for_capability_matching_in_snapshot(
                    self.snapshot.as_ref(),
                    model_name,
                    capability,
                    |_| true,
                ),
        )
        .map(DeploymentLease::synchronous_cleanup)
        .map_err(router_error_to_provider_error)
    }
}

impl RuntimeHandle {
    pub fn select_deployment_lease(
        &self,
        model_name: &str,
    ) -> Result<DeploymentLease, RouterError> {
        super::sync_compat::wait(self.select_deployment_lease_async(model_name))
            .map(DeploymentLease::synchronous_cleanup)
    }

    pub async fn select_deployment_lease_async(
        &self,
        model_name: &str,
    ) -> Result<DeploymentLease, RouterError> {
        self.binding
            .router
            .select_deployment_matching(self.snapshot.as_ref(), model_name, |_| true, None, 0)
            .await
    }
}