feat: add routing profile scheduling policies

crates/aether-routing-core/src/policy.rs

@@ -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())
+        );
+    }
+}