use super::deployment::{Deployment, DeploymentId};
use super::error::{CooldownReason, RouterError};
use super::execution::{
build_execution_result, infer_cooldown_reason, provider_error_to_router_error,
retryable_budget_scope, router_error_to_provider_error,
};
use super::fallback::{ExecutionResult, FallbackType};
use super::retry_policy::{RetryContext, RetryPolicy};
use super::selection::DeploymentLease;
use super::unified::Router;
use super::{RoutingSnapshot, RuntimeHandle};
use crate::core::providers::unified_provider::ProviderError;
use crate::core::types::model::ProviderCapability;
use std::collections::HashSet;
use std::sync::Arc;
impl Router {
async fn record_execution_failure(&self, deployment: &Deployment, reason: CooldownReason) {
self.record_local_failure(deployment, reason);
self.record_failure_circuit_for_deployment_async(deployment, reason)
.await;
}
pub async fn execute_with_selected_deployment_retry<T, F, Fut>(
&self,
model_name: &str,
operation: F,
) -> Result<(T, DeploymentId, u32, u64), (ProviderError, u32)>
where
F: Fn(Arc<Deployment>) -> Fut + Clone,
Fut: std::future::Future<Output = Result<(T, u64), ProviderError>>,
{
let snapshot = self.load_routing_snapshot();
self.execute_with_retry_inner(snapshot.as_ref(), model_name, None, 0, move |deployment| {
let result = operation(deployment);
async move {
result
.await
.map(|(value, tokens)| (value, Some(tokens)))
.map_err(|error| (error, None))
}
})
.await
.map(
|(value, deployment_id, _model_used, attempts, latency_us)| {
(value, deployment_id, attempts, latency_us)
},
)
}
#[deprecated(
since = "0.5.0",
note = "Use execute_with_selected_deployment_retry so callbacks receive the selected snapshot deployment"
)]
pub async fn execute_with_retry<T, F, Fut>(
&self,
model_name: &str,
operation: F,
) -> Result<(T, DeploymentId, u32, u64), (ProviderError, u32)>
where
F: Fn(DeploymentId) -> Fut + Clone,
Fut: std::future::Future<Output = Result<(T, u64), ProviderError>>,
{
self.execute_with_selected_deployment_retry(model_name, move |deployment| {
operation(deployment.id.clone())
})
.await
}
async fn execute_with_retry_inner<T, F, Fut>(
&self,
snapshot: &RoutingSnapshot,
model_name: &str,
capability: Option<&ProviderCapability>,
estimated_tokens: u64,
operation: F,
) -> Result<(T, DeploymentId, String, u32, u64), (ProviderError, u32)>
where
F: Fn(Arc<Deployment>) -> Fut + Clone,
Fut: std::future::Future<
Output = Result<(T, Option<u64>), (ProviderError, Option<Option<u64>>)>,
>,
{
let max_attempts = self.config.num_retries + 1;
let mut attempt = 1;
let mut last_error = None;
let mut excluded_budget_deployments = HashSet::new();
let mut tried_deployments = HashSet::new();
while attempt <= max_attempts {
let start = std::time::Instant::now();
let mut deployment_lease = match self
.select_retry_candidate(
snapshot,
model_name,
capability,
&excluded_budget_deployments,
&mut tried_deployments,
estimated_tokens,
)
.await
{
Ok(lease) => lease,
Err(router_err) => {
if matches!(
&router_err,
RouterError::UnsupportedCapability { .. }
| RouterError::NoAvailableDeployment(_)
) && let Some(err) = last_error.clone()
{
return Err((err, attempt));
}
if !excluded_budget_deployments.is_empty()
&& let Some(err) = last_error.clone()
{
return Err((err, attempt));
}
let provider_err = router_error_to_provider_error(router_err);
let retry_decision = RetryPolicy.decide(
&self.config,
&provider_err,
RetryContext::unary(attempt, max_attempts),
);
if retry_decision.should_retry {
last_error = Some(provider_err);
attempt += 1;
if let Some(delay) = retry_decision.delay {
tokio::time::sleep(delay).await;
}
continue;
} else {
return Err((provider_err, attempt));
}
}
};
let selected_deployment = deployment_lease.clone_deployment();
let deployment_id = selected_deployment.id.clone();
let result = operation(selected_deployment.clone()).await;
let latency_us = start.elapsed().as_micros() as u64;
match result {
Ok((value, tokens_used)) => {
let model_used = selected_deployment.model.clone();
if let Some(tokens) = tokens_used {
deployment_lease.preserve_admission_usage(tokens);
} else {
deployment_lease.preserve_admission_reservation(0);
}
deployment_lease.record_success(tokens_used.unwrap_or(0), latency_us);
self.record_success_circuit_for_deployment_async(deployment_lease.deployment())
.await;
if let Some(tokens) = tokens_used {
deployment_lease.commit_admission_async(tokens).await;
} else {
deployment_lease.retain_admission_async(0).await;
}
drop(deployment_lease);
return Ok((value, deployment_id, model_used, attempt, latency_us));
}
Err((err, completed_usage)) => {
if let Some(tokens) = completed_usage {
if let Some(tokens) = tokens {
deployment_lease.preserve_admission_usage(tokens);
} else {
deployment_lease.preserve_admission_reservation(0);
}
deployment_lease.record_interrupted_usage(tokens.unwrap_or(0));
}
if retryable_budget_scope(&err).is_some() {
excluded_budget_deployments.insert(deployment_id);
match completed_usage {
Some(Some(tokens)) => {
deployment_lease.commit_admission_async(tokens).await
}
Some(None) => deployment_lease.retain_admission_async(0).await,
None => deployment_lease.cancel_admission_async().await,
}
drop(deployment_lease);
last_error = Some(err);
continue;
}
let retry_decision = RetryPolicy.decide_for_deployment(
&self.config,
&selected_deployment.config,
&err,
RetryContext::unary(attempt, max_attempts),
);
if retry_decision.should_retry {
self.record_execution_failure(
deployment_lease.deployment(),
CooldownReason::ConsecutiveFailures,
)
.await;
tried_deployments.insert(deployment_id);
match completed_usage {
Some(Some(tokens)) => {
deployment_lease.commit_admission_async(tokens).await
}
Some(None) => deployment_lease.retain_admission_async(0).await,
None => deployment_lease.cancel_admission_async().await,
}
drop(deployment_lease);
last_error = Some(err);
attempt += 1;
if let Some(delay) = retry_decision.delay {
tokio::time::sleep(delay).await;
}
continue;
} else {
let cooldown_reason = infer_cooldown_reason(&err);
self.record_execution_failure(
deployment_lease.deployment(),
cooldown_reason,
)
.await;
match completed_usage {
Some(Some(tokens)) => {
deployment_lease.commit_admission_async(tokens).await
}
Some(None) => deployment_lease.retain_admission_async(0).await,
None => deployment_lease.cancel_admission_async().await,
}
drop(deployment_lease);
return Err((err, attempt));
}
}
}
}
Err((
last_error.unwrap_or_else(|| ProviderError::Other {
provider: "router",
message: if capability.is_some() {
"Unknown error during capability retry".to_string()
} else {
"Unknown error during retry".to_string()
},
}),
max_attempts,
))
}
async fn select_retry_candidate(
&self,
snapshot: &RoutingSnapshot,
model_name: &str,
capability: Option<&ProviderCapability>,
excluded_budget_deployments: &HashSet<String>,
tried_deployments: &mut HashSet<String>,
estimated_tokens: u64,
) -> Result<DeploymentLease, RouterError> {
let select = async |prefer_untried: bool| {
let is_candidate = |deployment: &Deployment| {
!excluded_budget_deployments.contains(deployment.id.as_str())
&& (!prefer_untried || !tried_deployments.contains(deployment.id.as_str()))
};
match capability {
Some(capability) => {
self.select_deployment_lease_for_capability_matching_with_estimate(
snapshot,
model_name,
capability,
is_candidate,
estimated_tokens,
)
.await
}
None => {
self.select_deployment_lease_matching_in_snapshot(
snapshot,
model_name,
is_candidate,
)
.await
}
}
};
match select(true).await {
Ok(lease) => Ok(lease),
Err(err)
if !tried_deployments.is_empty()
&& matches!(
err,
RouterError::UnsupportedCapability { .. }
| RouterError::NoAvailableDeployment(_)
) =>
{
match select(false).await {
Ok(lease) => {
tried_deployments.clear();
Ok(lease)
}
Err(fallback_err) => Err(fallback_err),
}
}
other => other,
}
}
async fn execute_stream_with_retry_inner<T, F, Fut>(
&self,
snapshot: &RoutingSnapshot,
model_name: &str,
capability: &ProviderCapability,
estimated_tokens: u64,
operation: F,
) -> Result<(T, DeploymentLease), ProviderError>
where
F: Fn(Arc<Deployment>) -> Fut + Clone,
Fut: std::future::Future<Output = Result<T, ProviderError>>,
{
let max_attempts = self.config.num_retries + 1;
let mut attempt = 1;
let mut last_operation_error = None;
let mut excluded_budget_deployments = HashSet::new();
let mut tried_deployments = HashSet::new();
while attempt <= max_attempts {
let mut deployment_lease = match self
.select_retry_candidate(
snapshot,
model_name,
Some(capability),
&excluded_budget_deployments,
&mut tried_deployments,
estimated_tokens,
)
.await
{
Ok(lease) => lease,
Err(router_err) => {
if matches!(
&router_err,
RouterError::UnsupportedCapability { .. }
| RouterError::NoAvailableDeployment(_)
) && let Some(err) = last_operation_error.clone()
{
return Err(err);
}
let provider_err = router_error_to_provider_error(router_err);
let retry_decision = RetryPolicy.decide(
&self.config,
&provider_err,
RetryContext::stream_pre_output(attempt, max_attempts),
);
if retry_decision.should_retry {
attempt += 1;
if let Some(delay) = retry_decision.delay {
tokio::time::sleep(delay).await;
}
continue;
}
return Err(last_operation_error.unwrap_or(provider_err));
}
};
let selected_deployment = deployment_lease.clone_deployment();
let deployment_id = selected_deployment.id.clone();
match operation.clone()(selected_deployment.clone()).await {
Ok(stream) => return Ok((stream, deployment_lease)),
Err(err) => {
if retryable_budget_scope(&err).is_some() {
excluded_budget_deployments.insert(deployment_id);
deployment_lease.cancel_admission_async().await;
drop(deployment_lease);
last_operation_error = Some(err);
continue;
}
let retry_decision = RetryPolicy.decide_for_deployment(
&self.config,
&selected_deployment.config,
&err,
RetryContext::stream_pre_output(attempt, max_attempts),
);
if retry_decision.should_retry {
self.record_execution_failure(
deployment_lease.deployment(),
CooldownReason::ConsecutiveFailures,
)
.await;
tried_deployments.insert(deployment_id);
deployment_lease.cancel_admission_async().await;
drop(deployment_lease);
last_operation_error = Some(err);
attempt += 1;
if let Some(delay) = retry_decision.delay {
tokio::time::sleep(delay).await;
}
continue;
}
let cooldown_reason = infer_cooldown_reason(&err);
self.record_execution_failure(deployment_lease.deployment(), cooldown_reason)
.await;
deployment_lease.cancel_admission_async().await;
drop(deployment_lease);
return Err(err);
}
}
}
Err(
last_operation_error.unwrap_or_else(|| ProviderError::Other {
provider: "router",
message: "Unknown error during streaming retry".to_string(),
}),
)
}
pub async fn execute_with_selected_deployment_capability_retry<T, F, Fut>(
&self,
model_name: &str,
capability: &ProviderCapability,
operation: F,
) -> Result<(T, DeploymentId, u32, u64), (ProviderError, u32)>
where
F: Fn(Arc<Deployment>) -> Fut + Clone,
Fut: std::future::Future<Output = Result<(T, u64), ProviderError>>,
{
let snapshot = self.load_routing_snapshot();
self.execute_with_retry_inner(
snapshot.as_ref(),
model_name,
Some(capability),
0,
move |deployment| {
let result = operation(deployment);
async move {
result
.await
.map(|(value, tokens)| (value, Some(tokens)))
.map_err(|error| (error, None))
}
},
)
.await
.map(
|(value, deployment_id, _model_used, attempts, latency_us)| {
(value, deployment_id, attempts, latency_us)
},
)
}
#[deprecated(
since = "0.5.0",
note = "Use execute_with_selected_deployment_capability_retry so callbacks receive the selected snapshot deployment"
)]
pub async fn execute_with_capability_retry<T, F, Fut>(
&self,
model_name: &str,
capability: &ProviderCapability,
operation: F,
) -> Result<(T, DeploymentId, u32, u64), (ProviderError, u32)>
where
F: Fn(DeploymentId) -> Fut + Clone,
Fut: std::future::Future<Output = Result<(T, u64), ProviderError>>,
{
self.execute_with_selected_deployment_capability_retry(
model_name,
capability,
move |deployment| operation(deployment.id.clone()),
)
.await
}
pub async fn execute_with_selected_deployment<T, F, Fut>(
&self,
model_name: &str,
operation: F,
) -> Result<ExecutionResult<T>, RouterError>
where
F: Fn(Arc<Deployment>) -> Fut + Clone,
Fut: std::future::Future<Output = Result<(T, u64), ProviderError>>,
{
let snapshot = self.load_routing_snapshot();
self.execute_with_selected_deployment_in_snapshot_typed(
snapshot.as_ref(),
model_name,
None,
0,
move |deployment| {
let result = operation(deployment);
async move {
result
.await
.map(|(value, tokens)| (value, Some(tokens)))
.map_err(|error| (error, None))
}
},
)
.await
.map_err(|error| provider_error_to_router_error(error, model_name))
}
async fn execute_with_selected_deployment_in_snapshot_typed<T, F, Fut>(
&self,
snapshot: &RoutingSnapshot,
model_name: &str,
capability: Option<&ProviderCapability>,
estimated_tokens: u64,
operation: F,
) -> Result<ExecutionResult<T>, ProviderError>
where
F: Fn(Arc<Deployment>) -> Fut + Clone,
Fut: std::future::Future<
Output = Result<(T, Option<u64>), (ProviderError, Option<Option<u64>>)>,
>,
{
let start = std::time::Instant::now();
let models_to_try = self.get_models_with_fallbacks_for_snapshot(
snapshot,
model_name,
FallbackType::General,
);
let max_models = 1 + self.config.max_fallbacks as usize;
let mut seen = std::collections::HashSet::new();
let models_to_try: Vec<_> = models_to_try
.into_iter()
.filter(|m| seen.insert(m.clone()))
.take(max_models)
.collect();
let mut last_error: Option<ProviderError> = None;
let mut total_attempts = 0;
for (model_idx, model) in models_to_try.iter().enumerate() {
let is_fallback = model_idx > 0;
if is_fallback {
self.fallback_triggered_count
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
tracing::info!(
original_model = %model_name,
fallback_model = %model,
fallback_index = model_idx,
error_type = %last_error.as_ref().map_or("unknown".to_string(), |e| e.redacted().to_string()),
"fallback triggered, trying next model"
);
}
match self
.execute_with_retry_inner(
snapshot,
model,
capability,
estimated_tokens,
operation.clone(),
)
.await
{
Ok((result, deployment_id, model_used, attempts, _latency_us)) => {
total_attempts += attempts;
let total_latency_us = start.elapsed().as_micros() as u64;
return Ok(build_execution_result(
result,
deployment_id,
total_attempts,
model_used,
is_fallback,
total_latency_us,
));
}
Err((err, attempts)) => {
total_attempts += attempts;
last_error = Some(err);
}
}
}
Err(last_error
.unwrap_or_else(|| ProviderError::model_not_found("router", model_name.to_string())))
}
#[deprecated(
since = "0.5.0",
note = "Use execute_with_selected_deployment so callbacks receive the selected snapshot deployment"
)]
pub async fn execute<T, F, Fut>(
&self,
model_name: &str,
operation: F,
) -> Result<ExecutionResult<T>, RouterError>
where
F: Fn(DeploymentId) -> Fut + Clone,
Fut: std::future::Future<Output = Result<(T, u64), ProviderError>>,
{
self.execute_with_selected_deployment(model_name, move |deployment| {
operation(deployment.id.clone())
})
.await
}
pub async fn execute_once_with_selected_deployment<T, F, Fut>(
&self,
model_name: &str,
operation: F,
) -> Result<ExecutionResult<T>, RouterError>
where
F: FnOnce(Arc<Deployment>) -> Fut,
Fut: std::future::Future<Output = Result<(T, u64), ProviderError>>,
{
let start = std::time::Instant::now();
let mut deployment_lease = self.select_deployment_lease_async(model_name).await?;
let selected_deployment = deployment_lease.clone_deployment();
let deployment_id = selected_deployment.id.clone();
let result = operation(selected_deployment.clone()).await;
let latency_us = start.elapsed().as_micros() as u64;
match result {
Ok((value, tokens_used)) => {
let model_used = selected_deployment.model.clone();
deployment_lease.preserve_admission_usage(tokens_used);
deployment_lease.record_success(tokens_used, latency_us);
self.record_success_circuit_for_deployment_async(deployment_lease.deployment())
.await;
deployment_lease.commit_admission_async(tokens_used).await;
drop(deployment_lease);
Ok(build_execution_result(
value,
deployment_id,
1,
model_used,
false,
latency_us,
))
}
Err(err) => {
let cooldown_reason = infer_cooldown_reason(&err);
self.record_execution_failure(deployment_lease.deployment(), cooldown_reason)
.await;
deployment_lease.cancel_admission_async().await;
drop(deployment_lease);
Err(provider_error_to_router_error(err, model_name))
}
}
}
#[deprecated(
since = "0.5.0",
note = "Use execute_once_with_selected_deployment so callbacks receive the selected snapshot deployment"
)]
pub async fn execute_once<T, F, Fut>(
&self,
model_name: &str,
operation: F,
) -> Result<ExecutionResult<T>, RouterError>
where
F: FnOnce(DeploymentId) -> Fut,
Fut: std::future::Future<Output = Result<(T, u64), ProviderError>>,
{
self.execute_once_with_selected_deployment(model_name, move |deployment| {
operation(deployment.id.clone())
})
.await
}
}
impl RuntimeHandle {
pub(crate) async fn execute_with_selected_deployment_typed<T, F, Fut>(
&self,
model_name: &str,
operation: F,
) -> Result<ExecutionResult<T>, ProviderError>
where
F: Fn(Arc<Deployment>) -> Fut + Clone,
Fut: std::future::Future<Output = Result<(T, u64), ProviderError>>,
{
self.binding
.router
.execute_with_selected_deployment_in_snapshot_typed(
self.snapshot.as_ref(),
model_name,
None,
0,
move |deployment| {
let result = operation(deployment);
async move {
result
.await
.map(|(value, tokens)| (value, Some(tokens)))
.map_err(|error| (error, None))
}
},
)
.await
}
pub(crate) async fn execute_with_selected_deployment_capability_typed<T, F, Fut>(
&self,
model_name: &str,
capability: &ProviderCapability,
estimated_tokens: u64,
operation: F,
) -> Result<ExecutionResult<T>, ProviderError>
where
F: Fn(Arc<Deployment>) -> Fut + Clone,
Fut: std::future::Future<
Output = Result<(T, Option<u64>), (ProviderError, Option<Option<u64>>)>,
>,
{
self.binding
.router
.execute_with_selected_deployment_in_snapshot_typed(
self.snapshot.as_ref(),
model_name,
Some(capability),
estimated_tokens,
operation,
)
.await
}
pub(crate) async fn execute_stream_with_selected_deployment_capability_typed<T, F, Fut>(
&self,
model_name: &str,
capability: &ProviderCapability,
estimated_tokens: u64,
operation: F,
) -> Result<(T, DeploymentLease), ProviderError>
where
F: Fn(Arc<Deployment>) -> Fut + Clone,
Fut: std::future::Future<Output = Result<T, ProviderError>>,
{
self.binding
.router
.execute_stream_with_retry_inner(
self.snapshot.as_ref(),
model_name,
capability,
estimated_tokens,
operation,
)
.await
}
pub async fn execute_with_selected_deployment<T, F, Fut>(
&self,
model_name: &str,
operation: F,
) -> Result<ExecutionResult<T>, RouterError>
where
F: Fn(Arc<Deployment>) -> Fut + Clone,
Fut: std::future::Future<Output = Result<(T, u64), ProviderError>>,
{
self.execute_with_selected_deployment_typed(model_name, operation)
.await
.map_err(|error| provider_error_to_router_error(error, model_name))
}
}