mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
feat: add routing profile scheduling policies
This commit is contained in:
fawney19
12
apps/aether-gateway/src/routing/mod.rs
Normal file
12
apps/aether-gateway/src/routing/mod.rs
Normal file
@@ -0,0 +1,12 @@
|
||||
pub(crate) mod mutations;
|
||||
pub(crate) mod resolver;
|
||||
pub(crate) mod selection;
|
||||
pub(crate) mod trace;
|
||||
|
||||
pub(crate) use mutations::apply_routing_mutation_plan;
|
||||
pub(crate) use resolver::{resolve_gateway_routing_policy, GatewayRoutingPolicyInput};
|
||||
pub(crate) use selection::{
|
||||
select_gateway_routing_group, GatewayRoutingGroupSelection, GatewayRoutingSelectionError,
|
||||
GatewayRoutingSelectionInput, ROUTING_GROUP_HEADER,
|
||||
};
|
||||
pub(crate) use trace::build_routing_trace_seed;
|
||||
49
apps/aether-gateway/src/routing/mutations.rs
Normal file
49
apps/aether-gateway/src/routing/mutations.rs
Normal file
@@ -0,0 +1,49 @@
|
||||
use aether_routing_core::{
|
||||
apply_json_patch_operations, validate_header_patch, MutationError, MutationPlan,
|
||||
RoutingHeaderPatch,
|
||||
};
|
||||
use http::StatusCode;
|
||||
use http::{HeaderMap, HeaderName, HeaderValue};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::GatewayError;
|
||||
|
||||
pub(crate) fn apply_routing_mutation_plan(
|
||||
body: &mut Value,
|
||||
headers: &mut HeaderMap,
|
||||
plan: &MutationPlan,
|
||||
) -> Result<(), GatewayError> {
|
||||
apply_json_patch_operations(body, &plan.body_patch).map_err(|err| GatewayError::Client {
|
||||
status: StatusCode::BAD_REQUEST,
|
||||
message: err.to_string(),
|
||||
})?;
|
||||
apply_header_patch(headers, &plan.header_patch).map_err(|err| GatewayError::Client {
|
||||
status: StatusCode::BAD_REQUEST,
|
||||
message: err.to_string(),
|
||||
})?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn apply_header_patch(
|
||||
headers: &mut HeaderMap,
|
||||
patch: &[RoutingHeaderPatch],
|
||||
) -> Result<(), MutationError> {
|
||||
validate_header_patch(patch)?;
|
||||
for item in patch {
|
||||
match item {
|
||||
RoutingHeaderPatch::Set { name, value } => {
|
||||
let name = HeaderName::from_bytes(name.as_bytes())
|
||||
.map_err(|_| MutationError::InvalidHeaderName(name.clone()))?;
|
||||
let value = HeaderValue::from_str(value)
|
||||
.map_err(|_| MutationError::InvalidHeaderName(name.to_string()))?;
|
||||
headers.insert(name, value);
|
||||
}
|
||||
RoutingHeaderPatch::Remove { name } => {
|
||||
let name = HeaderName::from_bytes(name.as_bytes())
|
||||
.map_err(|_| MutationError::InvalidHeaderName(name.clone()))?;
|
||||
headers.remove(name);
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
54
apps/aether-gateway/src/routing/resolver.rs
Normal file
54
apps/aether-gateway/src/routing/resolver.rs
Normal file
@@ -0,0 +1,54 @@
|
||||
use aether_routing_core::{
|
||||
resolve_routing_policy, ResolvedRoutingPolicy, RoutingGroupConfig, RoutingPolicyInput,
|
||||
RoutingRulePhase,
|
||||
};
|
||||
use http::StatusCode;
|
||||
use serde_json::Value;
|
||||
|
||||
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,
|
||||
}
|
||||
|
||||
pub(crate) fn resolve_gateway_routing_policy(
|
||||
input: GatewayRoutingPolicyInput<'_>,
|
||||
) -> Result<ResolvedRoutingPolicy, GatewayError> {
|
||||
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(),
|
||||
})
|
||||
}
|
||||
324
apps/aether-gateway/src/routing/selection.rs
Normal file
324
apps/aether-gateway/src/routing/selection.rs
Normal file
@@ -0,0 +1,324 @@
|
||||
use aether_data_contracts::repository::routing_profiles::{
|
||||
RoutingGroupBindingQuery, RoutingGroupBindingSubject, RoutingGroupLookupKey,
|
||||
RoutingGroupReadRepository, StoredRoutingGroup,
|
||||
};
|
||||
use thiserror::Error;
|
||||
|
||||
pub(crate) const ROUTING_GROUP_HEADER: &str = "x-aether-scheduler-group";
|
||||
|
||||
#[derive(Debug, Error, Clone, PartialEq, Eq)]
|
||||
pub(crate) enum GatewayRoutingSelectionError {
|
||||
#[error("routing group was explicitly requested but was not found: {0}")]
|
||||
NotFound(String),
|
||||
#[error("routing group was explicitly requested but is not enabled: {0}")]
|
||||
Disabled(String),
|
||||
#[error("routing group was explicitly requested but is not allowed for this principal: {0}")]
|
||||
Forbidden(String),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub(crate) struct GatewayRoutingSelectionInput<'a> {
|
||||
pub explicit_group: Option<&'a str>,
|
||||
pub user_id: Option<&'a str>,
|
||||
pub api_key_id: Option<&'a str>,
|
||||
pub user_group_ids: &'a [String],
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub(crate) struct GatewayRoutingGroupSelection {
|
||||
pub group: Option<StoredRoutingGroup>,
|
||||
pub source: String,
|
||||
}
|
||||
|
||||
pub(crate) async fn select_gateway_routing_group(
|
||||
repository: &(impl RoutingGroupReadRepository + ?Sized),
|
||||
input: GatewayRoutingSelectionInput<'_>,
|
||||
) -> Result<GatewayRoutingGroupSelection, GatewayRoutingSelectionError> {
|
||||
if let Some(explicit) = input
|
||||
.explicit_group
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
let group = repository
|
||||
.find_routing_group(RoutingGroupLookupKey::Id(explicit))
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.or({
|
||||
let group: Option<StoredRoutingGroup> = repository
|
||||
.find_routing_group(RoutingGroupLookupKey::Name(explicit))
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
group
|
||||
});
|
||||
let Some(group) = group else {
|
||||
return Err(GatewayRoutingSelectionError::NotFound(explicit.to_string()));
|
||||
};
|
||||
if !group.enabled {
|
||||
return Err(GatewayRoutingSelectionError::Disabled(group.id));
|
||||
}
|
||||
if !explicit_group_allowed(repository, &group.id, &input).await {
|
||||
return Err(GatewayRoutingSelectionError::Forbidden(group.id));
|
||||
}
|
||||
return Ok(GatewayRoutingGroupSelection {
|
||||
group: Some(group),
|
||||
source: "explicit_header".to_string(),
|
||||
});
|
||||
}
|
||||
|
||||
for (subject_type, subject_id, source) in default_binding_candidates(&input) {
|
||||
let bindings = repository
|
||||
.list_routing_group_bindings(&RoutingGroupBindingQuery {
|
||||
group_id: None,
|
||||
subject_type: Some(subject_type),
|
||||
subject_id: Some(subject_id.to_string()),
|
||||
})
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
for binding in bindings.into_iter().filter(|binding| binding.is_default) {
|
||||
let group = repository
|
||||
.find_routing_group(RoutingGroupLookupKey::Id(&binding.group_id))
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
if let Some(group) = group.filter(|group| group.enabled) {
|
||||
return Ok(GatewayRoutingGroupSelection {
|
||||
group: Some(group),
|
||||
source: source.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let system_default = repository
|
||||
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
|
||||
.await
|
||||
.ok()
|
||||
.flatten()
|
||||
.filter(|group| group.enabled);
|
||||
Ok(GatewayRoutingGroupSelection {
|
||||
group: system_default,
|
||||
source: "system_default".to_string(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn explicit_group_allowed(
|
||||
repository: &(impl RoutingGroupReadRepository + ?Sized),
|
||||
group_id: &str,
|
||||
input: &GatewayRoutingSelectionInput<'_>,
|
||||
) -> bool {
|
||||
if let Ok(Some(group)) = repository
|
||||
.find_routing_group(RoutingGroupLookupKey::Id(group_id))
|
||||
.await
|
||||
{
|
||||
if group.is_system_default {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
for (subject_type, subject_id, _) in default_binding_candidates(input) {
|
||||
let bindings = repository
|
||||
.list_routing_group_bindings(&RoutingGroupBindingQuery {
|
||||
group_id: Some(group_id.to_string()),
|
||||
subject_type: Some(subject_type),
|
||||
subject_id: Some(subject_id.to_string()),
|
||||
})
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
if bindings.iter().any(|binding| binding.allow_explicit_select) {
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
fn default_binding_candidates<'a>(
|
||||
input: &'a GatewayRoutingSelectionInput<'a>,
|
||||
) -> Vec<(RoutingGroupBindingSubject, &'a str, &'static str)> {
|
||||
let mut candidates = Vec::new();
|
||||
if let Some(api_key_id) = input.api_key_id {
|
||||
candidates.push((
|
||||
RoutingGroupBindingSubject::ApiKey,
|
||||
api_key_id,
|
||||
"api_key_default",
|
||||
));
|
||||
}
|
||||
if let Some(user_id) = input.user_id {
|
||||
candidates.push((RoutingGroupBindingSubject::User, user_id, "user_default"));
|
||||
}
|
||||
for group_id in input.user_group_ids {
|
||||
candidates.push((
|
||||
RoutingGroupBindingSubject::UserGroup,
|
||||
group_id.as_str(),
|
||||
"user_group_default",
|
||||
));
|
||||
}
|
||||
candidates
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use aether_data::repository::routing_profiles::InMemoryRoutingGroupRepository;
|
||||
use aether_data_contracts::repository::routing_profiles::{
|
||||
CreateRoutingGroupBindingRecord, CreateRoutingGroupRecord, RoutingGroupWriteRepository,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn selects_api_key_default_binding() {
|
||||
let repository = InMemoryRoutingGroupRepository::default();
|
||||
repository
|
||||
.create_routing_group(CreateRoutingGroupRecord {
|
||||
id: "group-1".to_string(),
|
||||
name: "default".to_string(),
|
||||
description: None,
|
||||
enabled: true,
|
||||
is_system_default: false,
|
||||
config_json: json!({}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
published_at: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
repository
|
||||
.create_routing_group_binding(CreateRoutingGroupBindingRecord {
|
||||
id: "binding-1".to_string(),
|
||||
group_id: "group-1".to_string(),
|
||||
subject_type: RoutingGroupBindingSubject::ApiKey,
|
||||
subject_id: "api-key-1".to_string(),
|
||||
is_default: true,
|
||||
allow_explicit_select: true,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let selection = select_gateway_routing_group(
|
||||
&repository,
|
||||
GatewayRoutingSelectionInput {
|
||||
explicit_group: None,
|
||||
user_id: None,
|
||||
api_key_id: Some("api-key-1"),
|
||||
user_group_ids: &[],
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
assert_eq!(selection.source, "api_key_default");
|
||||
assert_eq!(selection.group.unwrap().id, "group-1");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_explicit_group_that_does_not_exist() {
|
||||
let repository = InMemoryRoutingGroupRepository::default();
|
||||
|
||||
let error = select_gateway_routing_group(
|
||||
&repository,
|
||||
GatewayRoutingSelectionInput {
|
||||
explicit_group: Some("missing"),
|
||||
user_id: Some("user-1"),
|
||||
api_key_id: Some("api-key-1"),
|
||||
user_group_ids: &[],
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
GatewayRoutingSelectionError::NotFound("missing".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_explicit_disabled_group() {
|
||||
let repository = InMemoryRoutingGroupRepository::default();
|
||||
repository
|
||||
.create_routing_group(CreateRoutingGroupRecord {
|
||||
id: "disabled-group".to_string(),
|
||||
name: "disabled".to_string(),
|
||||
description: None,
|
||||
enabled: false,
|
||||
is_system_default: false,
|
||||
config_json: json!({}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
published_at: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let error = select_gateway_routing_group(
|
||||
&repository,
|
||||
GatewayRoutingSelectionInput {
|
||||
explicit_group: Some("disabled-group"),
|
||||
user_id: Some("user-1"),
|
||||
api_key_id: Some("api-key-1"),
|
||||
user_group_ids: &[],
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
GatewayRoutingSelectionError::Disabled("disabled-group".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_explicit_group_without_binding_permission() {
|
||||
let repository = InMemoryRoutingGroupRepository::default();
|
||||
repository
|
||||
.create_routing_group(CreateRoutingGroupRecord {
|
||||
id: "private-group".to_string(),
|
||||
name: "private".to_string(),
|
||||
description: None,
|
||||
enabled: true,
|
||||
is_system_default: false,
|
||||
config_json: json!({}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
published_at: None,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
repository
|
||||
.create_routing_group_binding(CreateRoutingGroupBindingRecord {
|
||||
id: "binding-1".to_string(),
|
||||
group_id: "private-group".to_string(),
|
||||
subject_type: RoutingGroupBindingSubject::ApiKey,
|
||||
subject_id: "api-key-1".to_string(),
|
||||
is_default: true,
|
||||
allow_explicit_select: false,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
})
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let error = select_gateway_routing_group(
|
||||
&repository,
|
||||
GatewayRoutingSelectionInput {
|
||||
explicit_group: Some("private-group"),
|
||||
user_id: Some("user-1"),
|
||||
api_key_id: Some("api-key-1"),
|
||||
user_group_ids: &[],
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap_err();
|
||||
|
||||
assert_eq!(
|
||||
error,
|
||||
GatewayRoutingSelectionError::Forbidden("private-group".to_string())
|
||||
);
|
||||
}
|
||||
}
|
||||
45
apps/aether-gateway/src/routing/trace.rs
Normal file
45
apps/aether-gateway/src/routing/trace.rs
Normal file
@@ -0,0 +1,45 @@
|
||||
use aether_routing_core::{ResolvedRoutingPolicy, RoutingDecisionTrace};
|
||||
|
||||
pub(crate) fn build_routing_trace_seed(
|
||||
policy: &ResolvedRoutingPolicy,
|
||||
client_api_format: &str,
|
||||
) -> RoutingDecisionTrace {
|
||||
RoutingDecisionTrace {
|
||||
group_id: policy.group_id.clone(),
|
||||
group_version: policy.group_version,
|
||||
selection_source: policy.selection_source.clone(),
|
||||
selected_rules: policy
|
||||
.matched_rules
|
||||
.iter()
|
||||
.map(|rule| rule.id.clone())
|
||||
.collect(),
|
||||
original_model: policy.requested_model.clone(),
|
||||
resolved_model: policy.resolved_model.clone(),
|
||||
client_api_format: client_api_format.to_string(),
|
||||
client_request_patch_summary: routing_patch_summary(&policy.mutation_plan),
|
||||
runtime_facts: aether_routing_core::RoutingRuntimeFacts {
|
||||
scheduler_mode: Some(policy.scheduling_mode),
|
||||
priority_mode: Some(policy.priority_mode),
|
||||
..Default::default()
|
||||
},
|
||||
..RoutingDecisionTrace::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn routing_patch_summary(
|
||||
plan: &aether_routing_core::MutationPlan,
|
||||
) -> aether_routing_core::RoutingPatchSummary {
|
||||
aether_routing_core::RoutingPatchSummary {
|
||||
body_paths: plan
|
||||
.body_patch
|
||||
.iter()
|
||||
.map(|operation| operation.path().to_string())
|
||||
.collect(),
|
||||
header_names: plan
|
||||
.header_patch
|
||||
.iter()
|
||||
.map(|operation| operation.name().to_string())
|
||||
.collect(),
|
||||
failed_action: None,
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user