Aether/apps/aether-gateway/src/routing/resolver.rs

use aether_routing_core::{
    resolve_routing_policy, MutationPlan, RankingOverlay, ResolvedRoutingPolicy,
    RoutingDefaultPolicy, RoutingGroupConfig, RoutingPolicyInput, RoutingRulePhase,
    RoutingSchedulingMode, RoutingSetPriorityMode, DEFAULT_STICKY_KEY_ATTEMPTS,
};
use http::StatusCode;
use serde_json::Value;
use std::collections::BTreeMap;

use crate::GatewayError;

#[derive(Debug, Clone)]
pub(crate) struct GatewayRoutingPolicyInput<'a> {
    pub group_id: Option<&'a str>,
    pub group_version: Option<i64>,
    pub group_config_json: &'a Value,
    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)]
pub(crate) struct GatewayStaticRoutingPolicyInput<'a> {
    pub group_id: Option<&'a str>,
    pub group_version: Option<i64>,
    pub group_config_json: &'a Value,
    pub selection_source: &'a str,
    pub requested_model: &'a str,
    pub resolved_model: &'a str,
}

pub(crate) fn resolve_gateway_routing_policy(
    input: GatewayRoutingPolicyInput<'_>,
) -> Result<ResolvedRoutingPolicy, GatewayError> {
    if let Some(policy) =
        resolve_gateway_static_default_routing_policy(GatewayStaticRoutingPolicyInput {
            group_id: input.group_id,
            group_version: input.group_version,
            group_config_json: input.group_config_json,
            selection_source: input.selection_source,
            requested_model: input.requested_model,
            resolved_model: input.resolved_model,
        })?
    {
        return Ok(policy);
    }

    let config = serde_json::from_value::<RoutingGroupConfig>(input.group_config_json.clone())
        .map_err(|err| GatewayError::Client {
            status: StatusCode::BAD_REQUEST,
            message: format!("invalid routing group config: {err}"),
        })?;
    resolve_routing_policy(
        &config,
        RoutingPolicyInput {
            group_id: input.group_id,
            group_version: input.group_version,
            selection_source: input.selection_source,
            requested_model: input.requested_model,
            resolved_model: input.resolved_model,
            api_format: input.api_format,
            user_id: input.user_id,
            api_key_id: input.api_key_id,
            headers: input.headers,
            body: input.body,
            phase: input.phase,
        },
    )
    .map_err(|err| GatewayError::Client {
        status: StatusCode::BAD_REQUEST,
        message: err.to_string(),
    })
}

pub(crate) fn resolve_gateway_static_default_routing_policy(
    input: GatewayStaticRoutingPolicyInput<'_>,
) -> Result<Option<ResolvedRoutingPolicy>, GatewayError> {
    let Some(default_policy) = static_default_policy_fields(input.group_config_json)? else {
        return Ok(None);
    };

    Ok(Some(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: default_policy.priority_mode,
        scheduling_mode: default_policy.scheduling_mode,
        keep_priority_on_conversion: default_policy.keep_priority_on_conversion,
        sticky_key_attempts: default_policy.sticky_key_attempts,
        execution_policy: default_policy.execution_policy,
        ranking_overlay: RankingOverlay::default(),
        mutation_plan: MutationPlan::default(),
        pool_policy_overrides: BTreeMap::new(),
        matched_rules: Vec::new(),
    }))
}

fn static_default_policy_fields(
    config_json: &Value,
) -> Result<Option<RoutingDefaultPolicy>, GatewayError> {
    let Some(object) = config_json.as_object() else {
        return Ok(None);
    };
    // A strategy's default policy applies to every model. Only model policies
    // and rules require the request-context-aware resolver; unknown legacy
    // fields (including the removed group allowlist) are intentionally ignored.
    if !routing_array_field_is_missing_or_empty(object, "model_policies")
        || !routing_array_field_is_missing_or_empty(object, "rules")
    {
        return Ok(None);
    }

    let Some(default_policy) = object.get("default_policy") else {
        return Ok(Some(RoutingDefaultPolicy::default()));
    };
    let Some(default_policy) = default_policy.as_object() else {
        return Ok(None);
    };

    let priority_mode = routing_enum_field(
        default_policy.get("priority_mode"),
        RoutingSetPriorityMode::default,
    )?;
    let scheduling_mode = routing_enum_field(
        default_policy.get("scheduling_mode"),
        RoutingSchedulingMode::default,
    )?;
    let keep_priority_on_conversion = match default_policy.get("keep_priority_on_conversion") {
        Some(value) => value.as_bool().ok_or_else(|| {
            invalid_routing_group_config("keep_priority_on_conversion must be a boolean")
        })?,
        None => false,
    };
    let sticky_key_attempts = match default_policy.get("sticky_key_attempts") {
        Some(value) => value
            .as_u64()
            .and_then(|value| u32::try_from(value).ok())
            .ok_or_else(|| {
                invalid_routing_group_config("sticky_key_attempts must be a non-negative integer")
            })?,
        None => DEFAULT_STICKY_KEY_ATTEMPTS,
    };
    let enable_cf_heartbeat = routing_bool_field(
        default_policy.get("enable_cf_heartbeat"),
        "enable_cf_heartbeat",
    )?;
    // Older strategies stored separate image/text heartbeat flags. Treat
    // either legacy flag as enabling the unified CF heartbeat setting while
    // allowing newly saved strategies to use only the canonical key.
    let legacy_image_heartbeat = routing_bool_field(
        default_policy.get("enable_openai_image_sync_heartbeat"),
        "enable_openai_image_sync_heartbeat",
    )?;
    let legacy_text_heartbeat = routing_bool_field(
        default_policy.get("enable_standard_text_sync_heartbeat"),
        "enable_standard_text_sync_heartbeat",
    )?;
    let execution_policy = aether_routing_core::RoutingExecutionPolicy {
        enable_cf_heartbeat: enable_cf_heartbeat || legacy_image_heartbeat || legacy_text_heartbeat,
        cyber_continue_failover: routing_bool_field(
            default_policy.get("cyber_continue_failover"),
            "cyber_continue_failover",
        )?,
    };

    Ok(Some(RoutingDefaultPolicy {
        priority_mode,
        scheduling_mode,
        keep_priority_on_conversion,
        sticky_key_attempts,
        execution_policy,
    }))
}

fn routing_bool_field(value: Option<&Value>, field: &str) -> Result<bool, GatewayError> {
    match value {
        Some(value) => value
            .as_bool()
            .ok_or_else(|| invalid_routing_group_config(format!("{field} must be a boolean"))),
        None => Ok(false),
    }
}

fn routing_array_field_is_missing_or_empty(
    object: &serde_json::Map<String, Value>,
    key: &str,
) -> bool {
    match object.get(key) {
        None => true,
        Some(Value::Array(values)) => values.is_empty(),
        Some(_) => false,
    }
}

fn routing_enum_field<T>(
    value: Option<&Value>,
    default: impl FnOnce() -> T,
) -> Result<T, GatewayError>
where
    T: serde::de::DeserializeOwned,
{
    match value {
        Some(value) => serde_json::from_value(value.clone()).map_err(|err| {
            invalid_routing_group_config(format!("invalid default routing policy: {err}"))
        }),
        None => Ok(default()),
    }
}

fn invalid_routing_group_config(message: impl Into<String>) -> GatewayError {
    GatewayError::Client {
        status: StatusCode::BAD_REQUEST,
        message: format!("invalid routing group config: {}", message.into()),
    }
}

#[cfg(test)]
mod tests {
    use serde_json::json;

    use super::*;

    #[test]
    fn static_default_policy_matches_full_resolver_without_body_context() {
        let config = json!({
            "default_policy": {
                "priority_mode": "global_key",
                "scheduling_mode": "load_balance",
                "keep_priority_on_conversion": true
            },
            "allowed_models": ["legacy-model"],
            "model_policies": [],
            "rules": []
        });

        let static_policy =
            resolve_gateway_static_default_routing_policy(GatewayStaticRoutingPolicyInput {
                group_id: Some("group-1"),
                group_version: Some(7),
                group_config_json: &config,
                selection_source: "system_default",
                requested_model: "mock-model",
                resolved_model: "mock-model",
            })
            .expect("static default policy should resolve")
            .expect("static default policy should be detected");

        let full_policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput {
            group_id: Some("group-1"),
            group_version: Some(7),
            group_config_json: &config,
            selection_source: "system_default",
            requested_model: "mock-model",
            resolved_model: "mock-model",
            api_format: "openai:chat",
            user_id: Some("user-1"),
            api_key_id: Some("key-1"),
            headers: &json!({"x-test": "value"}),
            body: &json!({"model": "mock-model"}),
            phase: RoutingRulePhase::ClientRequest,
        })
        .expect("full policy should resolve");

        assert_eq!(static_policy, full_policy);
        assert_eq!(
            static_policy.priority_mode,
            RoutingSetPriorityMode::GlobalKey
        );
        assert_eq!(
            static_policy.scheduling_mode,
            RoutingSchedulingMode::LoadBalance
        );
        assert!(static_policy.keep_priority_on_conversion);
        assert!(static_policy.mutation_plan.is_empty());
        assert!(static_policy.matched_rules.is_empty());
    }

    #[test]
    fn dynamic_routing_config_is_not_static_default() {
        let config = json!({
            "rules": [{
                "id": "rule-1",
                "conditions": {},
                "actions": [{
                    "type": "restrict_providers",
                    "provider_ids": ["provider-1"]
                }]
            }]
        });

        let policy =
            resolve_gateway_static_default_routing_policy(GatewayStaticRoutingPolicyInput {
                group_id: Some("group-1"),
                group_version: Some(1),
                group_config_json: &config,
                selection_source: "system_default",
                requested_model: "mock-model",
                resolved_model: "mock-model",
            })
            .expect("dynamic config should not fail static detection");

        assert!(policy.is_none());
    }
}