feat: add routing profile scheduling policies

This commit is contained in:
fawney19
2026-05-18 11:03:49 +08:00
parent a2f91b4108
commit 92813e6122
124 changed files with 11681 additions and 578 deletions

View File

@@ -0,0 +1,91 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RoutingRulePhase {
#[default]
ClientRequest,
ProviderRequest,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RoutingSetPriorityMode {
#[default]
Provider,
GlobalKey,
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RoutingSchedulingMode {
#[default]
CacheAffinity,
LoadBalance,
FixedOrder,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "op")]
pub enum RoutingJsonPatchOperation {
Add { path: String, value: Value },
Replace { path: String, value: Value },
Remove { path: String },
}
impl RoutingJsonPatchOperation {
pub fn path(&self) -> &str {
match self {
Self::Add { path, .. } | Self::Replace { path, .. } | Self::Remove { path } => path,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "op")]
pub enum RoutingHeaderPatch {
Set { name: String, value: String },
Remove { name: String },
}
impl RoutingHeaderPatch {
pub fn name(&self) -> &str {
match self {
Self::Set { name, .. } | Self::Remove { name } => name,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case", tag = "type")]
pub enum RoutingAction {
RestrictModels {
models: Vec<String>,
},
RestrictProviders {
provider_ids: Vec<String>,
},
RestrictKeys {
key_ids: Vec<String>,
},
SetScheduling {
priority_mode: Option<RoutingSetPriorityMode>,
scheduling_mode: Option<RoutingSchedulingMode>,
keep_priority_on_conversion: Option<bool>,
},
SetProviderPriority {
provider_id: String,
priority: i32,
},
SetKeyPriority {
key_id: String,
priority: i32,
},
JsonPatchBody {
patch: Vec<RoutingJsonPatchOperation>,
},
PatchHeaders {
patch: Vec<RoutingHeaderPatch>,
},
}

View File

@@ -0,0 +1,232 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RoutingConditionOp {
Eq,
Ne,
In,
Contains,
Exists,
Prefix,
Suffix,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum RoutingCondition {
All {
all: Vec<RoutingCondition>,
},
Any {
any: Vec<RoutingCondition>,
},
Not {
not: Box<RoutingCondition>,
},
Predicate {
field: String,
op: RoutingConditionOp,
#[serde(default, skip_serializing_if = "Option::is_none")]
value: Option<Value>,
},
Empty {},
}
impl Default for RoutingCondition {
fn default() -> Self {
Self::Empty {}
}
}
#[derive(Debug, Clone)]
pub struct RoutingConditionContext<'a> {
pub model: &'a str,
pub api_format: &'a str,
pub user_id: Option<&'a str>,
pub api_key_id: Option<&'a str>,
pub headers: &'a Value,
pub body: &'a Value,
}
impl RoutingCondition {
pub fn matches(&self, context: &RoutingConditionContext<'_>) -> bool {
match self {
Self::All { all } => all.iter().all(|condition| condition.matches(context)),
Self::Any { any } => any.iter().any(|condition| condition.matches(context)),
Self::Not { not } => !not.matches(context),
Self::Predicate { field, op, value } => {
let actual = resolve_field(context, field);
compare_condition(actual, *op, value.as_ref())
}
Self::Empty {} => true,
}
}
}
fn resolve_field(context: &RoutingConditionContext<'_>, field: &str) -> Option<Value> {
let normalized = field.trim();
match normalized {
"model" => return Some(Value::String(context.model.to_string())),
"api_format" | "client_api_format" => {
return Some(Value::String(context.api_format.to_string()))
}
"user_id" => {
return context
.user_id
.map(|value| Value::String(value.to_string()))
}
"api_key_id" => {
return context
.api_key_id
.map(|value| Value::String(value.to_string()))
}
_ => {}
}
if let Some(path) = normalized.strip_prefix("headers.") {
return lookup_dotted_path(context.headers, path).cloned();
}
if let Some(path) = normalized.strip_prefix("body.") {
return lookup_dotted_path(context.body, path).cloned();
}
None
}
fn lookup_dotted_path<'a>(root: &'a Value, path: &str) -> Option<&'a Value> {
let mut current = root;
for part in path.split('.').filter(|part| !part.is_empty()) {
match current {
Value::Object(map) => current = map.get(part)?,
Value::Array(items) => {
let index = part.parse::<usize>().ok()?;
current = items.get(index)?;
}
_ => return None,
}
}
Some(current)
}
fn compare_condition(
actual: Option<Value>,
op: RoutingConditionOp,
expected: Option<&Value>,
) -> bool {
match op {
RoutingConditionOp::Exists => actual.is_some(),
RoutingConditionOp::Eq => actual
.as_ref()
.zip(expected)
.is_some_and(|(actual, expected)| values_equal(actual, expected)),
RoutingConditionOp::Ne => actual
.as_ref()
.zip(expected)
.is_none_or(|(actual, expected)| !values_equal(actual, expected)),
RoutingConditionOp::In => {
actual
.as_ref()
.zip(expected)
.is_some_and(|(actual, expected)| {
expected
.as_array()
.is_some_and(|items| items.iter().any(|item| values_equal(actual, item)))
})
}
RoutingConditionOp::Contains => {
actual
.as_ref()
.zip(expected)
.is_some_and(|(actual, expected)| {
let Some(expected) = expected.as_str() else {
return false;
};
value_as_string(actual).is_some_and(|actual| actual.contains(expected))
})
}
RoutingConditionOp::Prefix => {
actual
.as_ref()
.zip(expected)
.is_some_and(|(actual, expected)| {
let Some(expected) = expected.as_str() else {
return false;
};
value_as_string(actual).is_some_and(|actual| actual.starts_with(expected))
})
}
RoutingConditionOp::Suffix => {
actual
.as_ref()
.zip(expected)
.is_some_and(|(actual, expected)| {
let Some(expected) = expected.as_str() else {
return false;
};
value_as_string(actual).is_some_and(|actual| actual.ends_with(expected))
})
}
}
}
fn values_equal(left: &Value, right: &Value) -> bool {
match (value_as_string(left), value_as_string(right)) {
(Some(left), Some(right)) => left == right,
_ => left == right,
}
}
fn value_as_string(value: &Value) -> Option<String> {
match value {
Value::String(value) => Some(value.clone()),
Value::Number(value) => Some(value.to_string()),
Value::Bool(value) => Some(value.to_string()),
_ => None,
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
fn context<'a>(headers: &'a Value, body: &'a Value) -> RoutingConditionContext<'a> {
RoutingConditionContext {
model: "gpt-5",
api_format: "openai:chat",
user_id: Some("user-1"),
api_key_id: Some("key-1"),
headers,
body,
}
}
#[test]
fn matches_all_body_and_header_predicates() {
let condition = RoutingCondition::All {
all: vec![
RoutingCondition::Predicate {
field: "model".to_string(),
op: RoutingConditionOp::Eq,
value: Some(json!("gpt-5")),
},
RoutingCondition::Predicate {
field: "headers.x-app".to_string(),
op: RoutingConditionOp::Eq,
value: Some(json!("coding")),
},
RoutingCondition::Predicate {
field: "body.reasoning_effort".to_string(),
op: RoutingConditionOp::In,
value: Some(json!(["high", "xhigh"])),
},
],
};
let headers = json!({"x-app":"coding"});
assert!(condition.matches(&context(&headers, &json!({"reasoning_effort":"high"}))));
assert!(!condition.matches(&context(&headers, &json!({"reasoning_effort":"low"}))));
}
}

View File

@@ -0,0 +1,36 @@
mod actions;
mod conditions;
mod model;
mod mutations;
mod policy;
mod ranking;
mod trace;
mod validation;
pub use actions::{
RoutingAction, RoutingHeaderPatch, RoutingJsonPatchOperation, RoutingRulePhase,
RoutingSchedulingMode, RoutingSetPriorityMode,
};
pub use conditions::{RoutingCondition, RoutingConditionContext, RoutingConditionOp};
pub use model::{
RoutingGroupBinding, RoutingGroupBindingSubject, RoutingGroupConfig, RoutingGroupRecord,
RoutingGroupVersionRecord, RoutingModelPolicy, RoutingPoolPolicyOverride, RoutingRule,
RoutingSchedulingPreset,
};
pub use mutations::{
apply_json_patch_operations, validate_header_patch, validate_json_patch_operations,
HeaderMutation, MutationError, MutationPlan,
};
pub use policy::{
resolve_routing_policy, MatchedRoutingRule, ResolvedRoutingPolicy, RoutingPolicyError,
RoutingPolicyInput,
};
pub use ranking::{
rank_vector_for_candidate, CandidateKind, RankingOverlay, RoutingCandidateFacts,
RoutingCandidateRankVector, ROUTING_PRIORITY_UNSPECIFIED,
};
pub use trace::{
RoutingCandidateTrace, RoutingDecisionTrace, RoutingPatchSummary, RoutingPoolExpansionTrace,
RoutingRuntimeFacts,
};
pub use validation::{validate_routing_group_config, RoutingValidationError};

View File

@@ -0,0 +1,131 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::actions::{
RoutingAction, RoutingRulePhase, RoutingSchedulingMode, RoutingSetPriorityMode,
};
use crate::conditions::RoutingCondition;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingSchedulingPreset {
pub preset: String,
#[serde(default = "default_true")]
pub enabled: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub mode: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct RoutingPoolPolicyOverride {
#[serde(default)]
pub scheduling_presets: Vec<RoutingSchedulingPreset>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct RoutingDefaultPolicy {
#[serde(default)]
pub priority_mode: RoutingSetPriorityMode,
#[serde(default)]
pub scheduling_mode: RoutingSchedulingMode,
#[serde(default)]
pub keep_priority_on_conversion: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct RoutingModelPolicy {
pub model: String,
#[serde(default)]
pub allowed_providers: Vec<String>,
#[serde(default)]
pub allowed_keys: Vec<String>,
#[serde(default)]
pub provider_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub key_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub pool_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub pool_policy_overrides: BTreeMap<String, RoutingPoolPolicyOverride>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingRule {
pub id: String,
#[serde(default)]
pub priority: i32,
#[serde(default = "default_true")]
pub enabled: bool,
#[serde(default)]
pub phase: RoutingRulePhase,
#[serde(default)]
pub conditions: RoutingCondition,
#[serde(default)]
pub actions: Vec<RoutingAction>,
#[serde(default)]
pub stop_processing: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
pub struct RoutingGroupConfig {
#[serde(default)]
pub allowed_models: Vec<String>,
#[serde(default)]
pub default_policy: RoutingDefaultPolicy,
#[serde(default)]
pub model_policies: Vec<RoutingModelPolicy>,
#[serde(default)]
pub rules: Vec<RoutingRule>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingGroupRecord {
pub id: String,
pub name: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
pub enabled: bool,
pub is_system_default: bool,
pub config_json: Value,
pub version: i64,
pub created_at: i64,
pub updated_at: i64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub published_at: Option<i64>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RoutingGroupBindingSubject {
User,
ApiKey,
UserGroup,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingGroupBinding {
pub id: String,
pub group_id: String,
pub subject_type: RoutingGroupBindingSubject,
pub subject_id: String,
pub is_default: bool,
pub allow_explicit_select: bool,
pub created_at: i64,
pub updated_at: i64,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingGroupVersionRecord {
pub id: String,
pub group_id: String,
pub version: i64,
pub config_json: Value,
pub created_at: i64,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub created_by: Option<String>,
}
fn default_true() -> bool {
true
}

View File

@@ -0,0 +1,236 @@
use std::collections::BTreeSet;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use thiserror::Error;
use crate::actions::{RoutingHeaderPatch, RoutingJsonPatchOperation};
const RESERVED_HEADERS: &[&str] = &[
"authorization",
"x-api-key",
"api-key",
"cookie",
"set-cookie",
"x-aether-trace-id",
"x-aether-internal",
"x-aether-scheduler-group",
];
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum MutationError {
#[error("json patch path must be an absolute JSON pointer: {0}")]
InvalidJsonPointer(String),
#[error("json patch cannot target reserved path: {0}")]
ReservedJsonPath(String),
#[error("json patch target does not exist: {0}")]
MissingTarget(String),
#[error("json patch parent is not an object: {0}")]
InvalidParent(String),
#[error("header patch targets reserved header: {0}")]
ReservedHeader(String),
#[error("header patch has invalid header name: {0}")]
InvalidHeaderName(String),
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct HeaderMutation {
pub set: Vec<(String, String)>,
pub remove: Vec<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct MutationPlan {
pub body_patch: Vec<RoutingJsonPatchOperation>,
pub header_patch: Vec<RoutingHeaderPatch>,
}
impl MutationPlan {
pub fn is_empty(&self) -> bool {
self.body_patch.is_empty() && self.header_patch.is_empty()
}
}
pub fn validate_json_patch_operations(
operations: &[RoutingJsonPatchOperation],
) -> Result<(), MutationError> {
for operation in operations {
let path = operation.path();
validate_json_pointer(path)?;
if is_reserved_json_path(path) {
return Err(MutationError::ReservedJsonPath(path.to_string()));
}
}
Ok(())
}
pub fn apply_json_patch_operations(
body: &mut Value,
operations: &[RoutingJsonPatchOperation],
) -> Result<(), MutationError> {
validate_json_patch_operations(operations)?;
for operation in operations {
match operation {
RoutingJsonPatchOperation::Add { path, value } => {
set_json_pointer(body, path, value.clone(), true)?;
}
RoutingJsonPatchOperation::Replace { path, value } => {
set_json_pointer(body, path, value.clone(), false)?;
}
RoutingJsonPatchOperation::Remove { path } => {
remove_json_pointer(body, path)?;
}
}
}
Ok(())
}
pub fn validate_header_patch(patch: &[RoutingHeaderPatch]) -> Result<(), MutationError> {
let reserved = RESERVED_HEADERS.iter().copied().collect::<BTreeSet<_>>();
for item in patch {
let name = item.name().trim().to_ascii_lowercase();
if name.is_empty()
|| name
.chars()
.any(|ch| !(ch.is_ascii_alphanumeric() || ch == '-'))
{
return Err(MutationError::InvalidHeaderName(item.name().to_string()));
}
if reserved.contains(name.as_str()) {
return Err(MutationError::ReservedHeader(item.name().to_string()));
}
}
Ok(())
}
fn validate_json_pointer(path: &str) -> Result<(), MutationError> {
if !path.starts_with('/') {
return Err(MutationError::InvalidJsonPointer(path.to_string()));
}
Ok(())
}
fn is_reserved_json_path(path: &str) -> bool {
matches!(
path,
"/authorization"
| "/api_key"
| "/provider_secret"
| "/upstream_url"
| "/upstream_base_url"
| "/auth"
)
}
fn set_json_pointer(
root: &mut Value,
pointer: &str,
value: Value,
allow_create: bool,
) -> Result<(), MutationError> {
let tokens = pointer_tokens(pointer);
if tokens.is_empty() {
*root = value;
return Ok(());
}
let (parents, leaf) = tokens.split_at(tokens.len() - 1);
let parent = descend_mut(root, parents, pointer)?;
match parent {
Value::Object(map) => {
if !allow_create && !map.contains_key(&leaf[0]) {
return Err(MutationError::MissingTarget(pointer.to_string()));
}
map.insert(leaf[0].clone(), value);
Ok(())
}
_ => Err(MutationError::InvalidParent(pointer.to_string())),
}
}
fn remove_json_pointer(root: &mut Value, pointer: &str) -> Result<(), MutationError> {
let tokens = pointer_tokens(pointer);
if tokens.is_empty() {
*root = Value::Null;
return Ok(());
}
let (parents, leaf) = tokens.split_at(tokens.len() - 1);
let parent = descend_mut(root, parents, pointer)?;
match parent {
Value::Object(map) => map
.remove(&leaf[0])
.map(|_| ())
.ok_or_else(|| MutationError::MissingTarget(pointer.to_string())),
_ => Err(MutationError::InvalidParent(pointer.to_string())),
}
}
fn descend_mut<'a>(
root: &'a mut Value,
tokens: &[String],
pointer: &str,
) -> Result<&'a mut Value, MutationError> {
let mut current = root;
for token in tokens {
match current {
Value::Object(map) => {
current = map
.get_mut(token)
.ok_or_else(|| MutationError::MissingTarget(pointer.to_string()))?;
}
Value::Null => {
*current = Value::Object(Map::new());
if let Value::Object(map) = current {
current = map
.entry(token.clone())
.or_insert_with(|| Value::Object(Map::new()));
}
}
_ => return Err(MutationError::InvalidParent(pointer.to_string())),
}
}
Ok(current)
}
fn pointer_tokens(pointer: &str) -> Vec<String> {
pointer
.trim_start_matches('/')
.split('/')
.filter(|part| !part.is_empty())
.map(|part| part.replace("~1", "/").replace("~0", "~"))
.collect()
}
#[cfg(test)]
mod tests {
use serde_json::json;
use crate::actions::RoutingJsonPatchOperation;
use super::*;
#[test]
fn applies_body_patch() {
let mut body = json!({"metadata":{}});
apply_json_patch_operations(
&mut body,
&[RoutingJsonPatchOperation::Add {
path: "/metadata/routing".to_string(),
value: json!("high"),
}],
)
.expect("patch should apply");
assert_eq!(body["metadata"]["routing"], json!("high"));
}
#[test]
fn rejects_reserved_headers() {
assert_eq!(
validate_header_patch(&[RoutingHeaderPatch::Set {
name: "authorization".to_string(),
value: "secret".to_string()
}]),
Err(MutationError::ReservedHeader("authorization".to_string()))
);
}
}

View File

@@ -0,0 +1,423 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
use crate::actions::{
RoutingAction, RoutingRulePhase, RoutingSchedulingMode, RoutingSetPriorityMode,
};
use crate::conditions::RoutingConditionContext;
use crate::model::{RoutingGroupConfig, RoutingModelPolicy, RoutingPoolPolicyOverride};
use crate::mutations::{validate_header_patch, validate_json_patch_operations, MutationPlan};
use crate::ranking::RankingOverlay;
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum RoutingPolicyError {
#[error("routing group config is invalid: {0}")]
InvalidConfig(String),
#[error("model is not allowed by routing group: {0}")]
ModelNotAllowed(String),
#[error("mutation action is invalid: {0}")]
InvalidMutation(String),
}
#[derive(Debug, Clone)]
pub struct RoutingPolicyInput<'a> {
pub group_id: Option<&'a str>,
pub group_version: Option<i64>,
pub selection_source: &'a str,
pub requested_model: &'a str,
pub resolved_model: &'a str,
pub api_format: &'a str,
pub user_id: Option<&'a str>,
pub api_key_id: Option<&'a str>,
pub headers: &'a Value,
pub body: &'a Value,
pub phase: RoutingRulePhase,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MatchedRoutingRule {
pub id: String,
pub priority: i32,
pub phase: RoutingRulePhase,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ResolvedRoutingPolicy {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_version: Option<i64>,
pub selection_source: String,
pub requested_model: String,
pub resolved_model: String,
pub priority_mode: RoutingSetPriorityMode,
pub scheduling_mode: RoutingSchedulingMode,
pub keep_priority_on_conversion: bool,
pub ranking_overlay: RankingOverlay,
pub mutation_plan: MutationPlan,
#[serde(default)]
pub pool_policy_overrides: BTreeMap<String, RoutingPoolPolicyOverride>,
#[serde(default)]
pub matched_rules: Vec<MatchedRoutingRule>,
}
pub fn resolve_routing_policy(
config: &RoutingGroupConfig,
input: RoutingPolicyInput<'_>,
) -> Result<ResolvedRoutingPolicy, RoutingPolicyError> {
if !model_allowed(&config.allowed_models, input.requested_model)
&& !model_allowed(&config.allowed_models, input.resolved_model)
{
return Err(RoutingPolicyError::ModelNotAllowed(
input.requested_model.to_string(),
));
}
let mut policy = ResolvedRoutingPolicy {
group_id: input.group_id.map(str::to_string),
group_version: input.group_version,
selection_source: input.selection_source.to_string(),
requested_model: input.requested_model.to_string(),
resolved_model: input.resolved_model.to_string(),
priority_mode: config.default_policy.priority_mode,
scheduling_mode: config.default_policy.scheduling_mode,
keep_priority_on_conversion: config.default_policy.keep_priority_on_conversion,
ranking_overlay: RankingOverlay::default(),
mutation_plan: MutationPlan::default(),
pool_policy_overrides: BTreeMap::new(),
matched_rules: Vec::new(),
};
for model_policy in matching_model_policies(config, input.requested_model, input.resolved_model)
{
apply_model_policy(&mut policy, model_policy);
}
let condition_context = RoutingConditionContext {
model: input.requested_model,
api_format: input.api_format,
user_id: input.user_id,
api_key_id: input.api_key_id,
headers: input.headers,
body: input.body,
};
let mut rules = config
.rules
.iter()
.filter(|rule| rule.enabled && rule.phase == input.phase)
.collect::<Vec<_>>();
rules.sort_by(|left, right| {
left.priority
.cmp(&right.priority)
.then(left.id.cmp(&right.id))
});
for rule in rules {
if !rule.conditions.matches(&condition_context) {
continue;
}
for action in &rule.actions {
apply_action(
&mut policy,
action,
input.requested_model,
input.resolved_model,
)?;
}
policy.matched_rules.push(MatchedRoutingRule {
id: rule.id.clone(),
priority: rule.priority,
phase: rule.phase,
});
if rule.stop_processing {
break;
}
}
Ok(policy)
}
fn apply_model_policy(policy: &mut ResolvedRoutingPolicy, model_policy: &RoutingModelPolicy) {
if !model_policy.allowed_providers.is_empty() {
policy.ranking_overlay.allowed_providers = model_policy.allowed_providers.clone();
}
if !model_policy.allowed_keys.is_empty() {
policy.ranking_overlay.allowed_keys = model_policy.allowed_keys.clone();
}
policy.ranking_overlay.provider_priority_overrides.extend(
model_policy
.provider_priority_overrides
.iter()
.map(|(key, value)| (key.clone(), *value)),
);
policy.ranking_overlay.key_priority_overrides.extend(
model_policy
.key_priority_overrides
.iter()
.map(|(key, value)| (key.clone(), *value)),
);
policy.ranking_overlay.pool_priority_overrides.extend(
model_policy
.pool_priority_overrides
.iter()
.map(|(key, value)| (key.clone(), *value)),
);
policy
.pool_policy_overrides
.extend(model_policy.pool_policy_overrides.clone());
}
fn apply_action(
policy: &mut ResolvedRoutingPolicy,
action: &RoutingAction,
requested_model: &str,
resolved_model: &str,
) -> Result<(), RoutingPolicyError> {
match action {
RoutingAction::RestrictModels { models } => {
if !model_allowed(models, requested_model) && !model_allowed(models, resolved_model) {
return Err(RoutingPolicyError::ModelNotAllowed(
requested_model.to_string(),
));
}
}
RoutingAction::RestrictProviders { provider_ids } => {
policy.ranking_overlay.allowed_providers = provider_ids.clone();
}
RoutingAction::RestrictKeys { key_ids } => {
policy.ranking_overlay.allowed_keys = key_ids.clone();
}
RoutingAction::SetScheduling {
priority_mode,
scheduling_mode,
keep_priority_on_conversion,
} => {
if let Some(priority_mode) = priority_mode {
policy.priority_mode = *priority_mode;
}
if let Some(scheduling_mode) = scheduling_mode {
policy.scheduling_mode = *scheduling_mode;
}
if let Some(keep_priority_on_conversion) = keep_priority_on_conversion {
policy.keep_priority_on_conversion = *keep_priority_on_conversion;
}
}
RoutingAction::SetProviderPriority {
provider_id,
priority,
} => {
policy
.ranking_overlay
.provider_priority_overrides
.insert(provider_id.clone(), *priority);
}
RoutingAction::SetKeyPriority { key_id, priority } => {
policy
.ranking_overlay
.key_priority_overrides
.insert(key_id.clone(), *priority);
}
RoutingAction::JsonPatchBody { patch } => {
validate_json_patch_operations(patch)
.map_err(|error| RoutingPolicyError::InvalidMutation(error.to_string()))?;
policy.mutation_plan.body_patch.extend(patch.clone());
}
RoutingAction::PatchHeaders { patch } => {
validate_header_patch(patch)
.map_err(|error| RoutingPolicyError::InvalidMutation(error.to_string()))?;
policy.mutation_plan.header_patch.extend(patch.clone());
}
}
Ok(())
}
fn matching_model_policies<'a>(
config: &'a RoutingGroupConfig,
requested_model: &str,
resolved_model: &str,
) -> Vec<&'a RoutingModelPolicy> {
config
.model_policies
.iter()
.filter(|policy| {
model_pattern_matches(&policy.model, requested_model)
|| model_pattern_matches(&policy.model, resolved_model)
})
.collect()
}
fn model_allowed(patterns: &[String], requested_model: &str) -> bool {
patterns.is_empty()
|| patterns
.iter()
.any(|pattern| model_pattern_matches(pattern, requested_model))
}
fn model_pattern_matches(pattern: &str, value: &str) -> bool {
let pattern = pattern.trim();
if pattern == "*" {
return true;
}
if let Some(prefix) = pattern.strip_suffix('*') {
return value.starts_with(prefix);
}
pattern == value
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use serde_json::json;
use crate::actions::{RoutingJsonPatchOperation, RoutingRulePhase};
use crate::conditions::{RoutingCondition, RoutingConditionOp};
use crate::model::{RoutingDefaultPolicy, RoutingRule};
use super::*;
#[test]
fn resolves_model_policy_and_matching_rule() {
let config = RoutingGroupConfig {
allowed_models: vec!["gpt-*".to_string()],
default_policy: RoutingDefaultPolicy::default(),
model_policies: vec![RoutingModelPolicy {
model: "gpt-5".to_string(),
allowed_providers: vec!["provider-a".to_string()],
provider_priority_overrides: BTreeMap::from([("provider-a".to_string(), 0)]),
pool_priority_overrides: BTreeMap::from([("provider-a".to_string(), 3)]),
..RoutingModelPolicy::default()
}],
rules: vec![RoutingRule {
id: "high".to_string(),
priority: 10,
enabled: true,
phase: RoutingRulePhase::ClientRequest,
conditions: RoutingCondition::Predicate {
field: "body.reasoning_effort".to_string(),
op: RoutingConditionOp::Eq,
value: Some(json!("high")),
},
actions: vec![RoutingAction::JsonPatchBody {
patch: vec![RoutingJsonPatchOperation::Add {
path: "/metadata/routing".to_string(),
value: json!("high"),
}],
}],
stop_processing: false,
}],
};
let policy = resolve_routing_policy(
&config,
RoutingPolicyInput {
group_id: Some("group-1"),
group_version: Some(1),
selection_source: "explicit",
requested_model: "gpt-5",
resolved_model: "gpt-5",
api_format: "openai:chat",
user_id: Some("user-1"),
api_key_id: Some("api-key-1"),
headers: &json!({}),
body: &json!({"reasoning_effort":"high"}),
phase: RoutingRulePhase::ClientRequest,
},
)
.expect("policy should resolve");
assert_eq!(policy.ranking_overlay.allowed_providers, vec!["provider-a"]);
assert_eq!(
policy
.ranking_overlay
.provider_priority_overrides
.get("provider-a"),
Some(&0)
);
assert_eq!(
policy
.ranking_overlay
.pool_priority_overrides
.get("provider-a"),
Some(&3)
);
assert_eq!(policy.matched_rules.len(), 1);
assert_eq!(policy.mutation_plan.body_patch.len(), 1);
}
#[test]
fn rejects_disallowed_model() {
let config = RoutingGroupConfig {
allowed_models: vec!["gpt-5".to_string()],
..RoutingGroupConfig::default()
};
let err = resolve_routing_policy(
&config,
RoutingPolicyInput {
group_id: None,
group_version: None,
selection_source: "test",
requested_model: "claude",
resolved_model: "claude",
api_format: "openai:chat",
user_id: None,
api_key_id: None,
headers: &json!({}),
body: &json!({}),
phase: RoutingRulePhase::ClientRequest,
},
)
.unwrap_err();
assert_eq!(
err,
RoutingPolicyError::ModelNotAllowed("claude".to_string())
);
}
#[test]
fn restrict_model_action_rejects_matching_request() {
let config = RoutingGroupConfig {
allowed_models: vec!["*".to_string()],
rules: vec![RoutingRule {
id: "restrict".to_string(),
priority: 1,
enabled: true,
phase: RoutingRulePhase::ClientRequest,
conditions: RoutingCondition::default(),
actions: vec![RoutingAction::RestrictModels {
models: vec!["gpt-5".to_string()],
}],
stop_processing: false,
}],
..RoutingGroupConfig::default()
};
let err = resolve_routing_policy(
&config,
RoutingPolicyInput {
group_id: None,
group_version: None,
selection_source: "test",
requested_model: "claude",
resolved_model: "claude",
api_format: "openai:chat",
user_id: None,
api_key_id: None,
headers: &json!({}),
body: &json!({}),
phase: RoutingRulePhase::ClientRequest,
},
)
.unwrap_err();
assert_eq!(
err,
RoutingPolicyError::ModelNotAllowed("claude".to_string())
);
}
}

View File

@@ -0,0 +1,181 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
pub const ROUTING_PRIORITY_UNSPECIFIED: i32 = i32::MAX;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CandidateKind {
Provider,
PoolGroup,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RankingOverlay {
#[serde(default)]
pub allowed_providers: Vec<String>,
#[serde(default)]
pub allowed_keys: Vec<String>,
#[serde(default)]
pub provider_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub key_priority_overrides: BTreeMap<String, i32>,
#[serde(default)]
pub pool_priority_overrides: BTreeMap<String, i32>,
}
impl RankingOverlay {
pub fn provider_priority(&self, provider_id: &str, fallback: i32) -> i32 {
self.provider_priority_overrides
.get(provider_id)
.copied()
.unwrap_or(fallback)
}
pub fn key_priority(&self, key_id: &str, fallback: i32) -> i32 {
self.key_priority_overrides
.get(key_id)
.copied()
.unwrap_or(fallback)
}
pub fn provider_priority_or_unspecified(&self, provider_id: &str) -> i32 {
self.provider_priority_overrides
.get(provider_id)
.copied()
.unwrap_or(ROUTING_PRIORITY_UNSPECIFIED)
}
pub fn key_priority_or_unspecified(&self, key_id: &str) -> i32 {
self.key_priority_overrides
.get(key_id)
.copied()
.unwrap_or(ROUTING_PRIORITY_UNSPECIFIED)
}
pub fn pool_priority_or_unspecified(&self, provider_id: &str) -> i32 {
self.pool_priority_overrides
.get(provider_id)
.copied()
.unwrap_or(ROUTING_PRIORITY_UNSPECIFIED)
}
pub fn provider_allowed(&self, provider_id: &str) -> bool {
self.allowed_providers.is_empty()
|| self
.allowed_providers
.iter()
.any(|item| item == provider_id)
}
pub fn key_allowed(&self, key_id: &str) -> bool {
self.allowed_keys.is_empty() || self.allowed_keys.iter().any(|item| item == key_id)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingCandidateFacts {
pub candidate_kind: CandidateKind,
pub provider_id: String,
pub endpoint_id: String,
pub model_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub key_id: Option<String>,
pub provider_priority: i32,
pub key_priority: i32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingCandidateRankVector {
pub provider_priority_before: i32,
pub provider_priority_after: i32,
pub key_priority_before: i32,
pub key_priority_after: i32,
}
pub fn rank_vector_for_candidate(
overlay: &RankingOverlay,
facts: &RoutingCandidateFacts,
) -> RoutingCandidateRankVector {
RoutingCandidateRankVector {
provider_priority_before: facts.provider_priority,
provider_priority_after: overlay.provider_priority_or_unspecified(&facts.provider_id),
key_priority_before: facts.key_priority,
key_priority_after: match facts.candidate_kind {
CandidateKind::Provider => facts
.key_id
.as_deref()
.map(|key_id| overlay.key_priority_or_unspecified(key_id))
.unwrap_or(ROUTING_PRIORITY_UNSPECIFIED),
CandidateKind::PoolGroup => overlay.pool_priority_or_unspecified(&facts.provider_id),
},
}
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::*;
#[test]
fn overlay_applies_provider_and_key_priority() {
let overlay = RankingOverlay {
provider_priority_overrides: BTreeMap::from([("provider-a".to_string(), 2)]),
key_priority_overrides: BTreeMap::from([("key-a".to_string(), 5)]),
..RankingOverlay::default()
};
let facts = RoutingCandidateFacts {
candidate_kind: CandidateKind::Provider,
provider_id: "provider-a".to_string(),
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: Some("key-a".to_string()),
provider_priority: 10,
key_priority: 20,
};
let vector = rank_vector_for_candidate(&overlay, &facts);
assert_eq!(vector.provider_priority_after, 2);
assert_eq!(vector.key_priority_after, 5);
}
#[test]
fn rank_vector_marks_missing_routing_priorities_unspecified() {
let facts = RoutingCandidateFacts {
candidate_kind: CandidateKind::Provider,
provider_id: "provider-a".to_string(),
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: Some("key-a".to_string()),
provider_priority: 10,
key_priority: 20,
};
let vector = rank_vector_for_candidate(&RankingOverlay::default(), &facts);
assert_eq!(vector.provider_priority_after, ROUTING_PRIORITY_UNSPECIFIED);
assert_eq!(vector.key_priority_after, ROUTING_PRIORITY_UNSPECIFIED);
}
#[test]
fn rank_vector_uses_pool_priority_for_pool_groups() {
let overlay = RankingOverlay {
pool_priority_overrides: BTreeMap::from([("provider-a".to_string(), 4)]),
..RankingOverlay::default()
};
let facts = RoutingCandidateFacts {
candidate_kind: CandidateKind::PoolGroup,
provider_id: "provider-a".to_string(),
endpoint_id: "endpoint-a".to_string(),
model_id: "model-a".to_string(),
key_id: None,
provider_priority: 10,
key_priority: 20,
};
let vector = rank_vector_for_candidate(&overlay, &facts);
assert_eq!(vector.key_priority_after, 4);
}
}

View File

@@ -0,0 +1,86 @@
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::actions::{RoutingSchedulingMode, RoutingSetPriorityMode};
use crate::ranking::{CandidateKind, RoutingCandidateRankVector};
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingPatchSummary {
#[serde(default)]
pub body_paths: Vec<String>,
#[serde(default)]
pub header_names: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub failed_action: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingCandidateTrace {
pub candidate_kind: CandidateKind,
pub provider_id: String,
pub endpoint_id: String,
pub model_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub key_id: Option<String>,
pub ranking_vector: RoutingCandidateRankVector,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub skip_reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub selected_order: Option<u32>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingPoolExpansionTrace {
pub pool_group_id: String,
pub key_id: String,
#[serde(default)]
pub pool_ranking_vector: Vec<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub pool_skip_reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub selected_order: Option<u32>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingRuntimeFacts {
#[serde(default)]
pub cache_affinity_hit: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub sticky_key_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub load_balance_seed: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub scheduler_mode: Option<RoutingSchedulingMode>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub priority_mode: Option<RoutingSetPriorityMode>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct RoutingDecisionTrace {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub group_version: Option<i64>,
pub selection_source: String,
#[serde(default)]
pub selected_rules: Vec<String>,
pub original_model: String,
pub resolved_model: String,
pub client_api_format: String,
#[serde(default)]
pub client_request_patch_summary: RoutingPatchSummary,
#[serde(default)]
pub provider_request_patch_summary: RoutingPatchSummary,
#[serde(default)]
pub global_candidates: Vec<RoutingCandidateTrace>,
#[serde(default)]
pub pool_expansion: Vec<RoutingPoolExpansionTrace>,
#[serde(default)]
pub runtime_facts: RoutingRuntimeFacts,
}
impl RoutingDecisionTrace {
pub fn to_extra_data_value(&self) -> Value {
serde_json::json!({ "routing_trace": self })
}
}

View File

@@ -0,0 +1,54 @@
use std::collections::BTreeSet;
use thiserror::Error;
use crate::model::RoutingGroupConfig;
use crate::mutations::{validate_header_patch, validate_json_patch_operations};
use crate::RoutingAction;
#[derive(Debug, Error, Clone, PartialEq, Eq)]
pub enum RoutingValidationError {
#[error("routing rule id is empty")]
EmptyRuleId,
#[error("duplicate routing rule id: {0}")]
DuplicateRuleId(String),
#[error("routing model policy selector is empty")]
EmptyModelSelector,
#[error("invalid mutation action: {0}")]
InvalidMutation(String),
}
pub fn validate_routing_group_config(
config: &RoutingGroupConfig,
) -> Result<(), RoutingValidationError> {
let mut rule_ids = BTreeSet::new();
for model_policy in &config.model_policies {
if model_policy.model.trim().is_empty() {
return Err(RoutingValidationError::EmptyModelSelector);
}
}
for rule in &config.rules {
if rule.id.trim().is_empty() {
return Err(RoutingValidationError::EmptyRuleId);
}
if !rule_ids.insert(rule.id.clone()) {
return Err(RoutingValidationError::DuplicateRuleId(rule.id.clone()));
}
for action in &rule.actions {
match action {
RoutingAction::JsonPatchBody { patch } => {
validate_json_patch_operations(patch).map_err(|error| {
RoutingValidationError::InvalidMutation(error.to_string())
})?;
}
RoutingAction::PatchHeaders { patch } => {
validate_header_patch(patch).map_err(|error| {
RoutingValidationError::InvalidMutation(error.to_string())
})?;
}
_ => {}
}
}
}
Ok(())
}