mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-01 17:00:21 +08:00
feat: add routing profile scheduling policies
This commit is contained in:
12
Cargo.lock
generated
12
Cargo.lock
generated
@@ -148,6 +148,7 @@ name = "aether-data-contracts"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"aether-ai-formats",
|
||||
"aether-routing-core",
|
||||
"async-trait",
|
||||
"chrono",
|
||||
"serde",
|
||||
@@ -196,6 +197,7 @@ dependencies = [
|
||||
"aether-pool-core",
|
||||
"aether-provider-pool",
|
||||
"aether-provider-transport",
|
||||
"aether-routing-core",
|
||||
"aether-runtime",
|
||||
"aether-runtime-state",
|
||||
"aether-scheduler-core",
|
||||
@@ -377,6 +379,16 @@ dependencies = [
|
||||
"webpki-roots 0.26.11",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aether-routing-core"
|
||||
version = "0.1.0"
|
||||
dependencies = [
|
||||
"regex",
|
||||
"serde",
|
||||
"serde_json",
|
||||
"thiserror 2.0.18",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "aether-runtime"
|
||||
version = "0.1.0"
|
||||
|
||||
@@ -6,6 +6,7 @@ members = [
|
||||
"crates/aether-ai-serving",
|
||||
"crates/aether-pool-core",
|
||||
"crates/aether-provider-pool",
|
||||
"crates/aether-routing-core",
|
||||
"crates/aether-data-contracts",
|
||||
"crates/aether-data-schema",
|
||||
"crates/aether-dispatch-core",
|
||||
@@ -41,6 +42,7 @@ aether-ai-formats = { path = "crates/aether-ai-formats" }
|
||||
aether-ai-serving = { path = "crates/aether-ai-serving" }
|
||||
aether-pool-core = { path = "crates/aether-pool-core" }
|
||||
aether-provider-pool = { path = "crates/aether-provider-pool" }
|
||||
aether-routing-core = { path = "crates/aether-routing-core" }
|
||||
aether-data-contracts = { path = "crates/aether-data-contracts" }
|
||||
aether-data-schema = { path = "crates/aether-data-schema" }
|
||||
aether-dispatch-core = { path = "crates/aether-dispatch-core" }
|
||||
|
||||
@@ -23,6 +23,7 @@ aether-oauth.workspace = true
|
||||
aether-pool-core.workspace = true
|
||||
aether-provider-pool.workspace = true
|
||||
aether-provider-transport.workspace = true
|
||||
aether-routing-core.workspace = true
|
||||
aether-scheduler-core.workspace = true
|
||||
aether-runtime.workspace = true
|
||||
aether-runtime-state.workspace = true
|
||||
|
||||
@@ -7,7 +7,13 @@ use aether_ai_serving::{
|
||||
AiCandidatePreselectionOutcome, AiSkippedCandidatePersistencePort,
|
||||
};
|
||||
use aether_dispatch_core::{DispatchSequence, DispatchSequenceItem};
|
||||
use aether_scheduler_core::{ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate};
|
||||
use aether_routing_core::{
|
||||
rank_vector_for_candidate, CandidateKind, ResolvedRoutingPolicy, RoutingCandidateFacts,
|
||||
RoutingCandidateTrace, RoutingDecisionTrace,
|
||||
};
|
||||
use aether_scheduler_core::{
|
||||
ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome,
|
||||
};
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
use std::collections::VecDeque;
|
||||
@@ -19,6 +25,7 @@ use tracing::warn;
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::ai_serving::planner::candidate_affinity_cache::remember_scheduler_affinity_for_candidate_at_epoch;
|
||||
use crate::ai_serving::planner::candidate_ranking::scheduler_ordering_config_for_routing_policy;
|
||||
use crate::ai_serving::planner::candidate_resolution::{
|
||||
resolve_and_rank_logical_local_execution_candidates, EligibleLocalExecutionCandidate,
|
||||
LocalExecutionCandidateKind, SkippedLocalExecutionCandidate,
|
||||
@@ -36,7 +43,7 @@ use crate::dispatch::refs::dispatch_ref_for_local_candidate;
|
||||
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
|
||||
use crate::orchestration::{local_attempt_slot_count, ExecutionAttemptIdentity};
|
||||
use crate::scheduler::candidate::API_KEY_CONCURRENCY_LIMIT_SKIP_REASON;
|
||||
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
|
||||
use crate::scheduler::config::SchedulerSchedulingMode;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
const POOL_KEY_RETRY_INDEX_STRIDE: u32 = 100;
|
||||
@@ -189,6 +196,7 @@ struct GatewayLocalCandidateMaterializationPort<'a, F, G> {
|
||||
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&'a ClientSessionAffinity>,
|
||||
required_capabilities: Option<&'a Value>,
|
||||
routing_policy: Option<&'a ResolvedRoutingPolicy>,
|
||||
sticky_session_token: Option<&'a str>,
|
||||
request_auth_channel: Option<&'a str>,
|
||||
persistence_policy: LocalCandidatePersistencePolicy<'a>,
|
||||
@@ -243,6 +251,7 @@ where
|
||||
self.auth_snapshot,
|
||||
self.client_session_affinity,
|
||||
self.required_capabilities,
|
||||
self.routing_policy,
|
||||
self.sticky_session_token,
|
||||
self.request_auth_channel,
|
||||
self.resolution_mode,
|
||||
@@ -281,6 +290,8 @@ where
|
||||
.skipped
|
||||
.record_runtime_miss_diagnostic,
|
||||
candidates,
|
||||
self.routing_policy,
|
||||
self.client_api_format,
|
||||
self.sticky_session_token,
|
||||
self.requested_model,
|
||||
self.request_auth_channel,
|
||||
@@ -294,6 +305,12 @@ where
|
||||
starting_candidate_index: u32,
|
||||
skipped_candidates: Vec<Self::Skipped>,
|
||||
) -> Result<(), Self::Error> {
|
||||
let skipped_candidates = attach_routing_trace_to_skipped_candidates(
|
||||
self.routing_policy,
|
||||
self.client_api_format,
|
||||
starting_candidate_index,
|
||||
skipped_candidates,
|
||||
);
|
||||
persist_skipped_local_execution_candidates_with_context(
|
||||
self.state.app(),
|
||||
self.trace_id,
|
||||
@@ -432,6 +449,7 @@ pub(crate) async fn materialize_local_execution_candidates_with_serving<F, G>(
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
required_capabilities: Option<&Value>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
sticky_session_token: Option<&str>,
|
||||
request_auth_channel: Option<&str>,
|
||||
persistence_policy: LocalCandidatePersistencePolicy<'_>,
|
||||
@@ -445,7 +463,8 @@ where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
|
||||
{
|
||||
let scheduler_cache_affinity_enabled = scheduler_cache_affinity_enabled(state).await;
|
||||
let scheduler_cache_affinity_enabled =
|
||||
scheduler_cache_affinity_enabled(state, routing_policy).await;
|
||||
let port = GatewayLocalCandidateMaterializationPort {
|
||||
state,
|
||||
trace_id,
|
||||
@@ -454,6 +473,7 @@ where
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
required_capabilities,
|
||||
routing_policy,
|
||||
sticky_session_token,
|
||||
request_auth_channel,
|
||||
persistence_policy,
|
||||
@@ -478,6 +498,7 @@ pub(crate) async fn build_local_execution_candidate_attempt_source_with_serving<
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
required_capabilities: Option<&Value>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
sticky_session_token: Option<&str>,
|
||||
request_auth_channel: Option<&str>,
|
||||
persistence_policy: LocalCandidatePersistencePolicy<'_>,
|
||||
@@ -491,7 +512,8 @@ where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
|
||||
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync,
|
||||
{
|
||||
let scheduler_cache_affinity_enabled = scheduler_cache_affinity_enabled(state).await;
|
||||
let scheduler_cache_affinity_enabled =
|
||||
scheduler_cache_affinity_enabled(state, routing_policy).await;
|
||||
let _ = build_available_extra_data;
|
||||
let (candidates, resolved_skipped) = resolve_and_rank_logical_local_execution_candidates(
|
||||
state,
|
||||
@@ -501,6 +523,7 @@ where
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
required_capabilities,
|
||||
routing_policy,
|
||||
sticky_session_token,
|
||||
request_auth_channel,
|
||||
resolution_mode,
|
||||
@@ -529,7 +552,12 @@ where
|
||||
trace_id,
|
||||
persistence_policy.skipped,
|
||||
u32::try_from(candidates.len()).unwrap_or(u32::MAX),
|
||||
skipped_candidates,
|
||||
attach_routing_trace_to_skipped_candidates(
|
||||
routing_policy,
|
||||
client_api_format,
|
||||
u32::try_from(candidates.len()).unwrap_or(u32::MAX),
|
||||
skipped_candidates,
|
||||
),
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -542,6 +570,7 @@ where
|
||||
sticky_session_token,
|
||||
requested_model,
|
||||
request_auth_channel,
|
||||
routing_policy,
|
||||
);
|
||||
|
||||
(
|
||||
@@ -559,6 +588,7 @@ fn build_logical_candidate_items<'a>(
|
||||
sticky_session_token: Option<&str>,
|
||||
requested_model: Option<&str>,
|
||||
request_auth_channel: Option<&str>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
) -> (VecDeque<LocalExecutionCandidateAttemptSourceItem<'a>>, u32) {
|
||||
let mut items = VecDeque::new();
|
||||
let mut next_candidate_index = starting_candidate_index;
|
||||
@@ -578,12 +608,13 @@ fn build_logical_candidate_items<'a>(
|
||||
}
|
||||
}
|
||||
LocalExecutionCandidateKind::PoolGroup => {
|
||||
let cursor = PoolKeyCursor::new(
|
||||
let cursor = PoolKeyCursor::new_with_routing_policy(
|
||||
state,
|
||||
candidate,
|
||||
sticky_session_token,
|
||||
requested_model,
|
||||
request_auth_channel,
|
||||
routing_policy,
|
||||
);
|
||||
let cursor = if let Some(trace_id) = trace_id {
|
||||
cursor.with_runtime_miss_diagnostic(trace_id, record_runtime_miss_diagnostic)
|
||||
@@ -615,6 +646,7 @@ pub(crate) async fn build_lazy_requested_model_execution_candidate_attempt_sourc
|
||||
auth_snapshot: &GatewayAuthApiKeySnapshot,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
required_capabilities: Option<&Value>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
sticky_session_token: Option<&str>,
|
||||
request_auth_channel: Option<&str>,
|
||||
persistence_policy: LocalCandidatePersistencePolicy<'_>,
|
||||
@@ -628,7 +660,8 @@ where
|
||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync + 'a,
|
||||
G: Fn(SkippedLocalExecutionCandidate) -> SkippedLocalExecutionCandidate + Send + Sync + 'a,
|
||||
{
|
||||
let scheduler_cache_affinity_enabled = scheduler_cache_affinity_enabled(state).await;
|
||||
let scheduler_cache_affinity_enabled =
|
||||
scheduler_cache_affinity_enabled(state, routing_policy).await;
|
||||
let _ = build_available_extra_data;
|
||||
let decorate_skipped_candidate = Arc::new(decorate_skipped_candidate);
|
||||
let record_runtime_miss_diagnostic = persistence_policy.skipped.record_runtime_miss_diagnostic;
|
||||
@@ -639,6 +672,7 @@ where
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
routing_policy,
|
||||
client_session_affinity,
|
||||
use_api_format_alias_match,
|
||||
key_mode,
|
||||
@@ -652,6 +686,7 @@ where
|
||||
auth_snapshot: auth_snapshot.clone(),
|
||||
client_session_affinity: client_session_affinity.cloned(),
|
||||
required_capabilities: required_capabilities.cloned(),
|
||||
routing_policy: routing_policy.cloned(),
|
||||
sticky_session_token: sticky_session_token.map(str::to_string),
|
||||
request_auth_channel: request_auth_channel.map(str::to_string),
|
||||
skipped_user_id: persistence_policy.skipped.user_id.to_string(),
|
||||
@@ -693,6 +728,7 @@ struct RequestedModelAttemptPageCursor<'a> {
|
||||
auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
client_session_affinity: Option<ClientSessionAffinity>,
|
||||
required_capabilities: Option<Value>,
|
||||
routing_policy: Option<ResolvedRoutingPolicy>,
|
||||
sticky_session_token: Option<String>,
|
||||
request_auth_channel: Option<String>,
|
||||
skipped_user_id: String,
|
||||
@@ -756,6 +792,7 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
|
||||
Some(&self.auth_snapshot),
|
||||
self.client_session_affinity.as_ref(),
|
||||
self.required_capabilities.as_ref(),
|
||||
self.routing_policy.as_ref(),
|
||||
self.sticky_session_token.as_deref(),
|
||||
self.request_auth_channel.as_deref(),
|
||||
self.resolution_mode,
|
||||
@@ -794,6 +831,7 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
|
||||
self.sticky_session_token.as_deref(),
|
||||
Some(&self.requested_model),
|
||||
self.request_auth_channel.as_deref(),
|
||||
self.routing_policy.as_ref(),
|
||||
);
|
||||
self.next_candidate_index = next_candidate_index
|
||||
.saturating_add(u32::try_from(skipped_candidate_count).unwrap_or(u32::MAX));
|
||||
@@ -814,7 +852,12 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
|
||||
&self.trace_id,
|
||||
skipped_persistence,
|
||||
skipped_starting_candidate_index,
|
||||
skipped_candidates,
|
||||
attach_routing_trace_to_skipped_candidates(
|
||||
self.routing_policy.as_ref(),
|
||||
&self.client_api_format,
|
||||
skipped_starting_candidate_index,
|
||||
skipped_candidates,
|
||||
),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
@@ -858,7 +901,12 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
|
||||
&self.trace_id,
|
||||
skipped_persistence,
|
||||
self.next_candidate_index,
|
||||
skipped_candidates,
|
||||
attach_routing_trace_to_skipped_candidates(
|
||||
self.routing_policy.as_ref(),
|
||||
&self.client_api_format,
|
||||
self.next_candidate_index,
|
||||
skipped_candidates,
|
||||
),
|
||||
)
|
||||
.await;
|
||||
self.next_candidate_index = self
|
||||
@@ -925,19 +973,14 @@ async fn pop_attempt_from_items(
|
||||
}
|
||||
}
|
||||
|
||||
async fn scheduler_cache_affinity_enabled(state: PlannerAppState<'_>) -> bool {
|
||||
match read_scheduler_ordering_config(state.app()).await {
|
||||
Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
event_name = "planner_scheduler_affinity_config_load_failed",
|
||||
log_type = "event",
|
||||
error = ?error,
|
||||
"failed to load scheduler config while checking cache affinity mode"
|
||||
);
|
||||
SchedulerSchedulingMode::default() == SchedulerSchedulingMode::CacheAffinity
|
||||
}
|
||||
}
|
||||
async fn scheduler_cache_affinity_enabled(
|
||||
state: PlannerAppState<'_>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
) -> bool {
|
||||
scheduler_ordering_config_for_routing_policy(state, routing_policy)
|
||||
.await
|
||||
.scheduling_mode
|
||||
== SchedulerSchedulingMode::CacheAffinity
|
||||
}
|
||||
|
||||
pub(crate) fn remember_first_local_candidate_affinity(
|
||||
@@ -1038,6 +1081,8 @@ async fn materialize_logical_local_execution_candidate_attempts<F>(
|
||||
context: LocalAvailableCandidatePersistenceContext<'_>,
|
||||
record_runtime_miss_diagnostic: bool,
|
||||
candidates: Vec<EligibleLocalExecutionCandidate>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
client_api_format: &str,
|
||||
sticky_session_token: Option<&str>,
|
||||
requested_model: Option<&str>,
|
||||
request_auth_channel: Option<&str>,
|
||||
@@ -1059,18 +1104,21 @@ where
|
||||
context,
|
||||
candidate,
|
||||
candidate_index,
|
||||
routing_policy,
|
||||
client_api_format,
|
||||
build_extra_data,
|
||||
)
|
||||
.await,
|
||||
);
|
||||
}
|
||||
LocalExecutionCandidateKind::PoolGroup => {
|
||||
let mut cursor = PoolKeyCursor::new(
|
||||
let mut cursor = PoolKeyCursor::new_with_routing_policy(
|
||||
state,
|
||||
candidate,
|
||||
sticky_session_token,
|
||||
requested_model,
|
||||
request_auth_channel,
|
||||
routing_policy,
|
||||
)
|
||||
.with_runtime_miss_diagnostic(trace_id, record_runtime_miss_diagnostic);
|
||||
let attempt_count_before_pool = attempts.len();
|
||||
@@ -1097,6 +1145,8 @@ async fn persist_available_local_execution_candidate_at_index<F>(
|
||||
context: LocalAvailableCandidatePersistenceContext<'_>,
|
||||
candidate: EligibleLocalExecutionCandidate,
|
||||
candidate_index: u32,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
client_api_format: &str,
|
||||
build_extra_data: &F,
|
||||
) -> Vec<LocalExecutionCandidateAttempt>
|
||||
where
|
||||
@@ -1107,6 +1157,16 @@ where
|
||||
available_candidate_base_extra_data_with_dispatch_ref(&candidate, build_extra_data),
|
||||
candidate.ranking.as_ref(),
|
||||
);
|
||||
let extra_data = attach_routing_trace_to_extra_data(
|
||||
routing_policy,
|
||||
client_api_format,
|
||||
&candidate.candidate,
|
||||
candidate.kind,
|
||||
candidate.ranking.as_ref(),
|
||||
None,
|
||||
Some(candidate_index),
|
||||
extra_data,
|
||||
);
|
||||
let should_persist = should_persist_available_local_candidate(&candidate);
|
||||
let mut attempts = Vec::with_capacity(attempt_slots as usize);
|
||||
let mut owned_candidate = Some(candidate);
|
||||
@@ -1190,6 +1250,160 @@ where
|
||||
Some(Value::Object(object))
|
||||
}
|
||||
|
||||
fn attach_routing_trace_to_skipped_candidates(
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
client_api_format: &str,
|
||||
starting_candidate_index: u32,
|
||||
skipped_candidates: Vec<SkippedLocalExecutionCandidate>,
|
||||
) -> Vec<SkippedLocalExecutionCandidate> {
|
||||
skipped_candidates
|
||||
.into_iter()
|
||||
.enumerate()
|
||||
.map(|(offset, skipped)| {
|
||||
let selected_order =
|
||||
starting_candidate_index.saturating_add(u32::try_from(offset).unwrap_or(u32::MAX));
|
||||
attach_routing_trace_to_skipped_candidate(
|
||||
routing_policy,
|
||||
client_api_format,
|
||||
selected_order,
|
||||
skipped,
|
||||
)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn attach_routing_trace_to_skipped_candidate(
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
client_api_format: &str,
|
||||
selected_order: u32,
|
||||
mut skipped_candidate: SkippedLocalExecutionCandidate,
|
||||
) -> SkippedLocalExecutionCandidate {
|
||||
let kind = if skipped_candidate
|
||||
.transport
|
||||
.as_ref()
|
||||
.is_some_and(|transport| {
|
||||
admin_provider_pool_config_from_config_value(transport.provider.config.as_ref())
|
||||
.is_some()
|
||||
}) {
|
||||
LocalExecutionCandidateKind::PoolGroup
|
||||
} else {
|
||||
LocalExecutionCandidateKind::SingleKey
|
||||
};
|
||||
skipped_candidate.extra_data = attach_routing_trace_to_extra_data(
|
||||
routing_policy,
|
||||
client_api_format,
|
||||
&skipped_candidate.candidate,
|
||||
kind,
|
||||
skipped_candidate.ranking.as_ref(),
|
||||
Some(skipped_candidate.skip_reason),
|
||||
Some(selected_order),
|
||||
skipped_candidate.extra_data,
|
||||
);
|
||||
skipped_candidate
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn attach_routing_trace_to_extra_data(
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
client_api_format: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
kind: LocalExecutionCandidateKind,
|
||||
ranking: Option<&SchedulerRankingOutcome>,
|
||||
skip_reason: Option<&'static str>,
|
||||
selected_order: Option<u32>,
|
||||
extra_data: Option<Value>,
|
||||
) -> Option<Value> {
|
||||
let Some(policy) = routing_policy else {
|
||||
return extra_data;
|
||||
};
|
||||
let routing_trace = routing_trace_for_candidate(
|
||||
policy,
|
||||
client_api_format,
|
||||
candidate,
|
||||
kind,
|
||||
ranking,
|
||||
skip_reason,
|
||||
selected_order,
|
||||
);
|
||||
Some(merge_routing_trace_into_extra_data(
|
||||
extra_data,
|
||||
routing_trace,
|
||||
))
|
||||
}
|
||||
|
||||
fn merge_routing_trace_into_extra_data(
|
||||
extra_data: Option<Value>,
|
||||
routing_trace: RoutingDecisionTrace,
|
||||
) -> Value {
|
||||
let mut object = match extra_data {
|
||||
Some(Value::Object(object)) => object,
|
||||
Some(value) => {
|
||||
let mut object = serde_json::Map::new();
|
||||
object.insert("extra".to_string(), value);
|
||||
object
|
||||
}
|
||||
None => serde_json::Map::new(),
|
||||
};
|
||||
object.insert(
|
||||
"routing_trace".to_string(),
|
||||
serde_json::json!(routing_trace),
|
||||
);
|
||||
Value::Object(object)
|
||||
}
|
||||
|
||||
fn routing_trace_for_candidate(
|
||||
policy: &ResolvedRoutingPolicy,
|
||||
client_api_format: &str,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
kind: LocalExecutionCandidateKind,
|
||||
ranking: Option<&SchedulerRankingOutcome>,
|
||||
skip_reason: Option<&'static str>,
|
||||
selected_order: Option<u32>,
|
||||
) -> RoutingDecisionTrace {
|
||||
let candidate_kind = routing_candidate_kind(kind);
|
||||
let mut trace = crate::routing::build_routing_trace_seed(policy, client_api_format);
|
||||
trace.global_candidates.push(RoutingCandidateTrace {
|
||||
candidate_kind,
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
model_id: candidate.model_id.clone(),
|
||||
key_id: match candidate_kind {
|
||||
CandidateKind::Provider => Some(candidate.key_id.clone()),
|
||||
CandidateKind::PoolGroup => None,
|
||||
},
|
||||
ranking_vector: rank_vector_for_candidate(
|
||||
&policy.ranking_overlay,
|
||||
&RoutingCandidateFacts {
|
||||
candidate_kind,
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
model_id: candidate.model_id.clone(),
|
||||
key_id: match candidate_kind {
|
||||
CandidateKind::Provider => Some(candidate.key_id.clone()),
|
||||
CandidateKind::PoolGroup => None,
|
||||
},
|
||||
provider_priority: candidate.provider_priority,
|
||||
key_priority: candidate
|
||||
.key_global_priority_for_format
|
||||
.unwrap_or(candidate.key_internal_priority),
|
||||
},
|
||||
),
|
||||
skip_reason: skip_reason.map(str::to_string),
|
||||
selected_order,
|
||||
});
|
||||
if let Some(ranking) = ranking {
|
||||
trace.runtime_facts.cache_affinity_hit = ranking.promoted_by == Some("cached_affinity");
|
||||
}
|
||||
trace
|
||||
}
|
||||
|
||||
fn routing_candidate_kind(kind: LocalExecutionCandidateKind) -> CandidateKind {
|
||||
match kind {
|
||||
LocalExecutionCandidateKind::SingleKey => CandidateKind::Provider,
|
||||
LocalExecutionCandidateKind::PoolGroup => CandidateKind::PoolGroup,
|
||||
}
|
||||
}
|
||||
|
||||
fn dispatch_sequence_from_attempts(
|
||||
attempts: Vec<LocalExecutionCandidateAttempt>,
|
||||
) -> DispatchSequence<LocalExecutionCandidateAttempt> {
|
||||
@@ -1644,6 +1858,7 @@ mod tests {
|
||||
auth_snapshot: Some(&auth_snapshot),
|
||||
client_session_affinity: None,
|
||||
required_capabilities: None,
|
||||
routing_policy: None,
|
||||
sticky_session_token: None,
|
||||
request_auth_channel: None,
|
||||
persistence_policy: LocalCandidatePersistencePolicy {
|
||||
@@ -1718,6 +1933,8 @@ mod tests {
|
||||
false,
|
||||
vec![pool_group, sample_eligible("normal-key", None)],
|
||||
None,
|
||||
"openai:chat",
|
||||
None,
|
||||
Some("gpt-5"),
|
||||
None,
|
||||
&|_| None,
|
||||
|
||||
@@ -5,6 +5,7 @@ use aether_ai_serving::{
|
||||
AiCandidateRankingPort, AiRankableCandidateParts, AiRankingContextConfig,
|
||||
AiRankingSchedulingMode,
|
||||
};
|
||||
use aether_routing_core::{ResolvedRoutingPolicy, RoutingSchedulingMode, RoutingSetPriorityMode};
|
||||
use async_trait::async_trait;
|
||||
use tracing::warn;
|
||||
|
||||
@@ -16,12 +17,12 @@ use crate::scheduler::config::{
|
||||
};
|
||||
use aether_scheduler_core::{
|
||||
matches_affinity_target, ClientSessionAffinity, SchedulerAffinityTarget,
|
||||
SchedulerMinimalCandidateSelectionCandidate, SchedulerRankableCandidate,
|
||||
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, SchedulerRankableCandidate,
|
||||
SchedulerRankingContext, SchedulerRankingOutcome,
|
||||
};
|
||||
|
||||
use super::candidate_affinity_cache::read_cached_scheduler_affinity_target;
|
||||
use super::candidate_resolution::EligibleLocalExecutionCandidate;
|
||||
use super::candidate_resolution::{EligibleLocalExecutionCandidate, LocalExecutionCandidateKind};
|
||||
use super::candidate_transport_ranking_facts::{
|
||||
resolve_cached_transport_ranking_facts, CandidateTransportRankingFacts,
|
||||
};
|
||||
@@ -33,6 +34,7 @@ struct GatewayLocalCandidateRankingPort<'a> {
|
||||
client_session_affinity: Option<&'a ClientSessionAffinity>,
|
||||
required_capabilities: Option<&'a serde_json::Value>,
|
||||
ordering_config: SchedulerOrderingConfig,
|
||||
routing_policy: Option<&'a ResolvedRoutingPolicy>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
@@ -89,8 +91,10 @@ impl AiCandidateRankingPort for GatewayLocalCandidateRankingPort<'_> {
|
||||
self.ordering_config,
|
||||
)
|
||||
.await;
|
||||
let routing_overlaid_candidate =
|
||||
routing_overlaid_candidate(self.routing_policy, candidate.kind, &candidate.candidate);
|
||||
Ok(build_ai_rankable_candidate(AiRankableCandidateParts {
|
||||
candidate: &candidate.candidate,
|
||||
candidate: &routing_overlaid_candidate,
|
||||
original_index,
|
||||
normalized_client_api_format,
|
||||
provider_api_format: candidate.provider_api_format.as_str(),
|
||||
@@ -122,8 +126,9 @@ pub(crate) async fn rank_eligible_local_execution_candidates(
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
) -> Vec<EligibleLocalExecutionCandidate> {
|
||||
let ordering_config = read_scheduler_ordering_config_or_default(state).await;
|
||||
let ordering_config = scheduler_ordering_config_for_routing_policy(state, routing_policy).await;
|
||||
let port = GatewayLocalCandidateRankingPort {
|
||||
state,
|
||||
requested_model,
|
||||
@@ -131,6 +136,7 @@ pub(crate) async fn rank_eligible_local_execution_candidates(
|
||||
client_session_affinity,
|
||||
required_capabilities,
|
||||
ordering_config,
|
||||
routing_policy,
|
||||
};
|
||||
|
||||
match run_ai_candidate_ranking(&port, candidates, normalized_client_api_format).await {
|
||||
@@ -189,6 +195,58 @@ fn ai_ranking_scheduling_mode(mode: SchedulerSchedulingMode) -> AiRankingSchedul
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn scheduler_ordering_config_for_routing_policy(
|
||||
state: PlannerAppState<'_>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
) -> SchedulerOrderingConfig {
|
||||
match routing_policy {
|
||||
Some(policy) => scheduler_ordering_config_from_routing_policy(policy),
|
||||
None => read_scheduler_ordering_config_or_default(state).await,
|
||||
}
|
||||
}
|
||||
|
||||
fn scheduler_ordering_config_from_routing_policy(
|
||||
policy: &ResolvedRoutingPolicy,
|
||||
) -> SchedulerOrderingConfig {
|
||||
SchedulerOrderingConfig {
|
||||
priority_mode: match policy.priority_mode {
|
||||
RoutingSetPriorityMode::Provider => SchedulerPriorityMode::Provider,
|
||||
RoutingSetPriorityMode::GlobalKey => SchedulerPriorityMode::GlobalKey,
|
||||
},
|
||||
scheduling_mode: match policy.scheduling_mode {
|
||||
RoutingSchedulingMode::FixedOrder => SchedulerSchedulingMode::FixedOrder,
|
||||
RoutingSchedulingMode::CacheAffinity => SchedulerSchedulingMode::CacheAffinity,
|
||||
RoutingSchedulingMode::LoadBalance => SchedulerSchedulingMode::LoadBalance,
|
||||
},
|
||||
keep_priority_on_conversion: policy.keep_priority_on_conversion,
|
||||
}
|
||||
}
|
||||
|
||||
fn routing_overlaid_candidate(
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
kind: LocalExecutionCandidateKind,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> SchedulerMinimalCandidateSelectionCandidate {
|
||||
let Some(policy) = routing_policy else {
|
||||
return candidate.clone();
|
||||
};
|
||||
let mut overlaid = candidate.clone();
|
||||
overlaid.provider_priority = policy
|
||||
.ranking_overlay
|
||||
.provider_priority_or_unspecified(candidate.provider_id.as_str());
|
||||
let overlaid_key_priority = match kind {
|
||||
LocalExecutionCandidateKind::SingleKey => policy
|
||||
.ranking_overlay
|
||||
.key_priority_or_unspecified(candidate.key_id.as_str()),
|
||||
LocalExecutionCandidateKind::PoolGroup => policy
|
||||
.ranking_overlay
|
||||
.pool_priority_or_unspecified(candidate.provider_id.as_str()),
|
||||
};
|
||||
overlaid.key_internal_priority = overlaid_key_priority;
|
||||
overlaid.key_global_priority_for_format = Some(overlaid_key_priority);
|
||||
overlaid
|
||||
}
|
||||
|
||||
async fn read_scheduler_ordering_config_or_default(
|
||||
state: PlannerAppState<'_>,
|
||||
) -> SchedulerOrderingConfig {
|
||||
@@ -304,6 +362,82 @@ mod tests {
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn routing_policy_priorities_do_not_fall_back_to_candidate_priorities() {
|
||||
let mut candidate = sample_candidate("endpoint-1", "key-1");
|
||||
candidate.provider_priority = 7;
|
||||
candidate.key_internal_priority = 3;
|
||||
candidate.key_global_priority_for_format = Some(2);
|
||||
let policy = aether_routing_core::ResolvedRoutingPolicy {
|
||||
group_id: Some("group-1".to_string()),
|
||||
group_version: Some(1),
|
||||
selection_source: "system_default".to_string(),
|
||||
requested_model: "gpt-5".to_string(),
|
||||
resolved_model: "gpt-5".to_string(),
|
||||
priority_mode: aether_routing_core::RoutingSetPriorityMode::Provider,
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::CacheAffinity,
|
||||
keep_priority_on_conversion: false,
|
||||
ranking_overlay: aether_routing_core::RankingOverlay::default(),
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: BTreeMap::new(),
|
||||
matched_rules: Vec::new(),
|
||||
};
|
||||
|
||||
let overlaid = super::routing_overlaid_candidate(
|
||||
Some(&policy),
|
||||
LocalExecutionCandidateKind::SingleKey,
|
||||
&candidate,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
overlaid.provider_priority,
|
||||
aether_routing_core::ROUTING_PRIORITY_UNSPECIFIED
|
||||
);
|
||||
assert_eq!(
|
||||
overlaid.key_internal_priority,
|
||||
aether_routing_core::ROUTING_PRIORITY_UNSPECIFIED
|
||||
);
|
||||
assert_eq!(
|
||||
overlaid.key_global_priority_for_format,
|
||||
Some(aether_routing_core::ROUTING_PRIORITY_UNSPECIFIED)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn routing_policy_uses_pool_priority_for_pool_group_global_key_slot() {
|
||||
let mut candidate = sample_candidate("endpoint-1", "representative-key");
|
||||
candidate.provider_priority = 7;
|
||||
candidate.key_internal_priority = 3;
|
||||
candidate.key_global_priority_for_format = Some(2);
|
||||
let policy = aether_routing_core::ResolvedRoutingPolicy {
|
||||
group_id: Some("group-1".to_string()),
|
||||
group_version: Some(1),
|
||||
selection_source: "system_default".to_string(),
|
||||
requested_model: "gpt-5".to_string(),
|
||||
resolved_model: "gpt-5".to_string(),
|
||||
priority_mode: aether_routing_core::RoutingSetPriorityMode::GlobalKey,
|
||||
scheduling_mode: aether_routing_core::RoutingSchedulingMode::CacheAffinity,
|
||||
keep_priority_on_conversion: false,
|
||||
ranking_overlay: aether_routing_core::RankingOverlay {
|
||||
pool_priority_overrides: BTreeMap::from([("provider-1".to_string(), 4)]),
|
||||
key_priority_overrides: BTreeMap::from([("representative-key".to_string(), 1)]),
|
||||
..Default::default()
|
||||
},
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: BTreeMap::new(),
|
||||
matched_rules: Vec::new(),
|
||||
};
|
||||
|
||||
let overlaid = super::routing_overlaid_candidate(
|
||||
Some(&policy),
|
||||
LocalExecutionCandidateKind::PoolGroup,
|
||||
&candidate,
|
||||
);
|
||||
|
||||
assert_eq!(overlaid.key_internal_priority, 4);
|
||||
assert_eq!(overlaid.key_global_priority_for_format, Some(4));
|
||||
}
|
||||
|
||||
fn sample_provider() -> StoredProviderCatalogProvider {
|
||||
sample_provider_with_options("provider-1", false, 0)
|
||||
}
|
||||
@@ -1062,6 +1196,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1141,6 +1276,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1216,6 +1352,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1282,6 +1419,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1364,6 +1502,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1438,6 +1577,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1515,6 +1655,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1610,6 +1751,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1713,6 +1855,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -1809,6 +1952,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
aether_ai_serving::AiCandidateResolutionMode::Standard,
|
||||
)
|
||||
.await;
|
||||
@@ -1901,6 +2045,7 @@ mod tests {
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
aether_ai_serving::AiCandidateResolutionMode::Standard,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -4,6 +4,7 @@ use aether_ai_serving::{
|
||||
run_ai_candidate_resolution, AiCandidateResolutionMode, AiCandidateResolutionPort,
|
||||
AiCandidateResolutionRequest,
|
||||
};
|
||||
use aether_routing_core::ResolvedRoutingPolicy;
|
||||
use async_trait::async_trait;
|
||||
use std::convert::Infallible;
|
||||
use tracing::warn;
|
||||
@@ -60,6 +61,7 @@ struct GatewayLocalCandidateResolutionPort<'a> {
|
||||
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&'a ClientSessionAffinity>,
|
||||
required_capabilities: Option<&'a serde_json::Value>,
|
||||
routing_policy: Option<&'a ResolvedRoutingPolicy>,
|
||||
request_auth_channel: Option<&'a str>,
|
||||
}
|
||||
|
||||
@@ -97,6 +99,11 @@ impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> {
|
||||
transport: &Self::Transport,
|
||||
requested_model: Option<&str>,
|
||||
) -> Option<&'static str> {
|
||||
if let Some(skip_reason) =
|
||||
routing_policy_candidate_skip_reason(self.routing_policy, candidate, transport)
|
||||
{
|
||||
return Some(skip_reason);
|
||||
}
|
||||
if provider_transport_uses_pool(transport) {
|
||||
return pool_group_common_transport_skip_reason(candidate, transport);
|
||||
}
|
||||
@@ -172,6 +179,7 @@ impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> {
|
||||
self.auth_snapshot,
|
||||
self.client_session_affinity,
|
||||
self.required_capabilities,
|
||||
self.routing_policy,
|
||||
)
|
||||
.await)
|
||||
}
|
||||
@@ -192,6 +200,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates(
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
_sticky_session_token: Option<&str>,
|
||||
request_auth_channel: Option<&str>,
|
||||
) -> (
|
||||
@@ -207,6 +216,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates(
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
required_capabilities,
|
||||
routing_policy,
|
||||
None,
|
||||
request_auth_channel,
|
||||
AiCandidateResolutionMode::Standard,
|
||||
@@ -222,6 +232,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates_without_transpor
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
_sticky_session_token: Option<&str>,
|
||||
request_auth_channel: Option<&str>,
|
||||
) -> (
|
||||
@@ -237,6 +248,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates_without_transpor
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
required_capabilities,
|
||||
routing_policy,
|
||||
None,
|
||||
request_auth_channel,
|
||||
AiCandidateResolutionMode::WithoutTransportPairGate,
|
||||
@@ -252,6 +264,7 @@ pub(crate) async fn resolve_and_rank_logical_local_execution_candidates(
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
_sticky_session_token: Option<&str>,
|
||||
request_auth_channel: Option<&str>,
|
||||
mode: AiCandidateResolutionMode,
|
||||
@@ -267,6 +280,7 @@ pub(crate) async fn resolve_and_rank_logical_local_execution_candidates(
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
required_capabilities,
|
||||
routing_policy,
|
||||
None,
|
||||
request_auth_channel,
|
||||
mode,
|
||||
@@ -283,6 +297,7 @@ async fn resolve_and_rank_local_execution_candidates_with_mode(
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
_sticky_session_token: Option<&str>,
|
||||
request_auth_channel: Option<&str>,
|
||||
mode: AiCandidateResolutionMode,
|
||||
@@ -298,6 +313,7 @@ async fn resolve_and_rank_local_execution_candidates_with_mode(
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
required_capabilities,
|
||||
routing_policy,
|
||||
None,
|
||||
request_auth_channel,
|
||||
mode,
|
||||
@@ -315,6 +331,7 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
|
||||
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
_sticky_session_token: Option<&str>,
|
||||
request_auth_channel: Option<&str>,
|
||||
mode: AiCandidateResolutionMode,
|
||||
@@ -330,6 +347,7 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
|
||||
auth_snapshot,
|
||||
client_session_affinity,
|
||||
required_capabilities,
|
||||
routing_policy,
|
||||
request_auth_channel,
|
||||
};
|
||||
|
||||
@@ -369,6 +387,28 @@ fn provider_transport_uses_pool(transport: &GatewayProviderTransportSnapshot) ->
|
||||
.is_some()
|
||||
}
|
||||
|
||||
fn routing_policy_candidate_skip_reason(
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
) -> Option<&'static str> {
|
||||
let policy = routing_policy?;
|
||||
if !policy
|
||||
.ranking_overlay
|
||||
.provider_allowed(candidate.provider_id.as_str())
|
||||
{
|
||||
return Some("routing_profile_disallowed_provider");
|
||||
}
|
||||
if !provider_transport_uses_pool(transport)
|
||||
&& !policy
|
||||
.ranking_overlay
|
||||
.key_allowed(candidate.key_id.as_str())
|
||||
{
|
||||
return Some("routing_profile_disallowed_key");
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn pool_group_common_transport_skip_reason(
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
transport: &GatewayProviderTransportSnapshot,
|
||||
|
||||
@@ -2,6 +2,7 @@ use aether_ai_serving::{
|
||||
run_ai_candidate_preselection, AiCandidatePreselectionOutcome, AiCandidatePreselectionPort,
|
||||
};
|
||||
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
|
||||
use aether_routing_core::ResolvedRoutingPolicy;
|
||||
use aether_scheduler_core::{
|
||||
enumerate_minimal_candidate_selection_with_model_directives, normalize_api_format,
|
||||
resolve_requested_global_model_name_with_model_directives,
|
||||
@@ -35,6 +36,7 @@ struct GatewayLocalCandidatePreselectionPort<'a> {
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<&'a serde_json::Value>,
|
||||
auth_snapshot: &'a GatewayAuthApiKeySnapshot,
|
||||
routing_policy: Option<&'a ResolvedRoutingPolicy>,
|
||||
client_session_affinity: Option<&'a ClientSessionAffinity>,
|
||||
use_api_format_alias_match: bool,
|
||||
key_mode: LocalCandidatePreselectionKeyMode,
|
||||
@@ -100,13 +102,14 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
|
||||
let enable_model_directives = self.model_directive_enabled_api_formats.contains(
|
||||
&crate::ai_serving::normalize_api_format_alias(candidate_api_format),
|
||||
);
|
||||
matches_client_format
|
||||
|| auth_snapshot_allows_cross_format_candidate(
|
||||
self.auth_snapshot,
|
||||
self.requested_model,
|
||||
candidate,
|
||||
enable_model_directives,
|
||||
)
|
||||
routing_policy_allows_provider(self.routing_policy, candidate)
|
||||
&& (matches_client_format
|
||||
|| auth_snapshot_allows_cross_format_candidate(
|
||||
self.auth_snapshot,
|
||||
self.requested_model,
|
||||
candidate,
|
||||
enable_model_directives,
|
||||
))
|
||||
}
|
||||
|
||||
fn skipped_candidate_allowed(
|
||||
@@ -118,13 +121,14 @@ impl AiCandidatePreselectionPort for GatewayLocalCandidatePreselectionPort<'_> {
|
||||
let enable_model_directives = self.model_directive_enabled_api_formats.contains(
|
||||
&crate::ai_serving::normalize_api_format_alias(candidate_api_format),
|
||||
);
|
||||
matches_client_format
|
||||
|| auth_snapshot_allows_cross_format_candidate(
|
||||
self.auth_snapshot,
|
||||
self.requested_model,
|
||||
&skipped_candidate.candidate,
|
||||
enable_model_directives,
|
||||
)
|
||||
routing_policy_allows_provider(self.routing_policy, &skipped_candidate.candidate)
|
||||
&& (matches_client_format
|
||||
|| auth_snapshot_allows_cross_format_candidate(
|
||||
self.auth_snapshot,
|
||||
self.requested_model,
|
||||
&skipped_candidate.candidate,
|
||||
enable_model_directives,
|
||||
))
|
||||
}
|
||||
|
||||
fn candidate_key(&self, candidate: &Self::Candidate) -> String {
|
||||
@@ -144,6 +148,7 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
auth_snapshot: &GatewayAuthApiKeySnapshot,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
use_api_format_alias_match: bool,
|
||||
key_mode: LocalCandidatePreselectionKeyMode,
|
||||
@@ -166,6 +171,7 @@ pub(crate) async fn preselect_local_execution_candidates_with_serving(
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
routing_policy,
|
||||
client_session_affinity,
|
||||
use_api_format_alias_match,
|
||||
key_mode,
|
||||
@@ -182,6 +188,7 @@ pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_se
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
auth_snapshot: &GatewayAuthApiKeySnapshot,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
use_api_format_alias_match: bool,
|
||||
key_mode: LocalCandidatePreselectionKeyMode,
|
||||
@@ -213,6 +220,7 @@ pub(crate) async fn preselect_local_execution_candidates_for_api_formats_with_se
|
||||
require_streaming,
|
||||
required_capabilities,
|
||||
auth_snapshot,
|
||||
routing_policy,
|
||||
client_session_affinity,
|
||||
use_api_format_alias_match,
|
||||
key_mode,
|
||||
@@ -230,6 +238,7 @@ pub(crate) struct LocalCandidatePreselectionPageCursor<'a> {
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<serde_json::Value>,
|
||||
auth_snapshot: GatewayAuthApiKeySnapshot,
|
||||
routing_policy: Option<ResolvedRoutingPolicy>,
|
||||
client_session_affinity: Option<ClientSessionAffinity>,
|
||||
use_api_format_alias_match: bool,
|
||||
key_mode: LocalCandidatePreselectionKeyMode,
|
||||
@@ -253,6 +262,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
|
||||
require_streaming: bool,
|
||||
required_capabilities: Option<&serde_json::Value>,
|
||||
auth_snapshot: &GatewayAuthApiKeySnapshot,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
client_session_affinity: Option<&ClientSessionAffinity>,
|
||||
use_api_format_alias_match: bool,
|
||||
key_mode: LocalCandidatePreselectionKeyMode,
|
||||
@@ -283,6 +293,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
|
||||
require_streaming,
|
||||
required_capabilities: required_capabilities.cloned(),
|
||||
auth_snapshot: auth_snapshot.clone(),
|
||||
routing_policy: routing_policy.cloned(),
|
||||
client_session_affinity: client_session_affinity.cloned(),
|
||||
use_api_format_alias_match,
|
||||
key_mode,
|
||||
@@ -620,16 +631,17 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
|
||||
candidate_api_format: &str,
|
||||
enable_model_directives: bool,
|
||||
) -> bool {
|
||||
matches_client_api_format(
|
||||
self.use_api_format_alias_match,
|
||||
candidate_api_format,
|
||||
&self.client_api_format,
|
||||
) || auth_snapshot_allows_cross_format_candidate(
|
||||
&self.auth_snapshot,
|
||||
&self.requested_model,
|
||||
candidate,
|
||||
enable_model_directives,
|
||||
)
|
||||
routing_policy_allows_provider(self.routing_policy.as_ref(), candidate)
|
||||
&& (matches_client_api_format(
|
||||
self.use_api_format_alias_match,
|
||||
candidate_api_format,
|
||||
&self.client_api_format,
|
||||
) || auth_snapshot_allows_cross_format_candidate(
|
||||
&self.auth_snapshot,
|
||||
&self.requested_model,
|
||||
candidate,
|
||||
enable_model_directives,
|
||||
))
|
||||
}
|
||||
|
||||
fn skipped_candidate_allowed_for_page(
|
||||
@@ -638,16 +650,17 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
|
||||
candidate_api_format: &str,
|
||||
enable_model_directives: bool,
|
||||
) -> bool {
|
||||
matches_client_api_format(
|
||||
self.use_api_format_alias_match,
|
||||
candidate_api_format,
|
||||
&self.client_api_format,
|
||||
) || auth_snapshot_allows_cross_format_candidate(
|
||||
&self.auth_snapshot,
|
||||
&self.requested_model,
|
||||
&skipped_candidate.candidate,
|
||||
enable_model_directives,
|
||||
)
|
||||
routing_policy_allows_provider(self.routing_policy.as_ref(), &skipped_candidate.candidate)
|
||||
&& (matches_client_api_format(
|
||||
self.use_api_format_alias_match,
|
||||
candidate_api_format,
|
||||
&self.client_api_format,
|
||||
) || auth_snapshot_allows_cross_format_candidate(
|
||||
&self.auth_snapshot,
|
||||
&self.requested_model,
|
||||
&skipped_candidate.candidate,
|
||||
enable_model_directives,
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -739,6 +752,18 @@ pub(crate) fn auth_snapshot_allows_cross_format_candidate(
|
||||
true
|
||||
}
|
||||
|
||||
fn routing_policy_allows_provider(
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||
) -> bool {
|
||||
match routing_policy {
|
||||
Some(policy) => policy
|
||||
.ranking_overlay
|
||||
.provider_allowed(candidate.provider_id.as_str()),
|
||||
None => true,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
@@ -887,6 +912,7 @@ mod tests {
|
||||
None,
|
||||
&auth_snapshot,
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
|
||||
)
|
||||
@@ -943,6 +969,7 @@ mod tests {
|
||||
None,
|
||||
&auth_snapshot,
|
||||
None,
|
||||
None,
|
||||
true,
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
|
||||
)
|
||||
|
||||
@@ -1,10 +1,27 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_ai_serving::{run_ai_authenticated_decision_input, AiAuthenticatedDecisionInputPort};
|
||||
use aether_routing_core::{
|
||||
rank_vector_for_candidate, CandidateKind, ResolvedRoutingPolicy, RoutingCandidateFacts,
|
||||
RoutingCandidateTrace, RoutingDecisionTrace, RoutingPoolExpansionTrace, RoutingRulePhase,
|
||||
};
|
||||
use aether_scheduler_core::ClientSessionAffinity;
|
||||
use async_trait::async_trait;
|
||||
use http::StatusCode;
|
||||
use http::{HeaderMap, HeaderName, HeaderValue};
|
||||
use serde_json::{json, Value};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::planner::common::extract_standard_requested_model;
|
||||
use crate::ai_serving::{ExecutionRuntimeAuthContext, GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||
use crate::client_session_affinity::client_session_affinity_from_request;
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::{AppState, GatewayError};
|
||||
use crate::routing::{
|
||||
apply_routing_mutation_plan, build_routing_trace_seed, resolve_gateway_routing_policy,
|
||||
select_gateway_routing_group, GatewayRoutingPolicyInput, GatewayRoutingSelectionError,
|
||||
GatewayRoutingSelectionInput, ROUTING_GROUP_HEADER,
|
||||
};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct ResolvedLocalDecisionAuthInput {
|
||||
@@ -21,6 +38,9 @@ pub(crate) struct LocalRequestedModelDecisionInput {
|
||||
pub(crate) required_capabilities: Option<serde_json::Value>,
|
||||
pub(crate) request_auth_channel: Option<String>,
|
||||
pub(crate) client_session_affinity: Option<ClientSessionAffinity>,
|
||||
pub(crate) routing_policy: Option<ResolvedRoutingPolicy>,
|
||||
pub(crate) routing_trace_seed: Option<RoutingDecisionTrace>,
|
||||
pub(crate) routing_context: Option<LocalRoutingRequestContext>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -31,6 +51,92 @@ pub(crate) struct LocalAuthenticatedDecisionInput {
|
||||
pub(crate) client_session_affinity: Option<ClientSessionAffinity>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct LocalRoutingRequestContext {
|
||||
pub(crate) group_id: Option<String>,
|
||||
pub(crate) group_version: Option<i64>,
|
||||
pub(crate) group_config_json: Value,
|
||||
pub(crate) selection_source: String,
|
||||
pub(crate) client_api_format: String,
|
||||
pub(crate) effective_body_json: Value,
|
||||
pub(crate) effective_headers: HeaderMap,
|
||||
}
|
||||
|
||||
impl LocalRequestedModelDecisionInput {
|
||||
pub(crate) fn effective_body_json<'a>(&'a self, fallback: &'a Value) -> &'a Value {
|
||||
self.routing_context
|
||||
.as_ref()
|
||||
.map(|context| &context.effective_body_json)
|
||||
.unwrap_or(fallback)
|
||||
}
|
||||
|
||||
pub(crate) fn effective_headers<'a>(&'a self, fallback: &'a HeaderMap) -> &'a HeaderMap {
|
||||
self.routing_context
|
||||
.as_ref()
|
||||
.map(|context| &context.effective_headers)
|
||||
.unwrap_or(fallback)
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn apply_provider_request_routing_policy_to_decision(
|
||||
input: &LocalRequestedModelDecisionInput,
|
||||
decision: &mut AiExecutionDecision,
|
||||
) -> Result<(), GatewayError> {
|
||||
let Some(context) = input.routing_context.as_ref() else {
|
||||
return Ok(());
|
||||
};
|
||||
let provider_api_format = decision
|
||||
.provider_api_format
|
||||
.as_deref()
|
||||
.unwrap_or(context.client_api_format.as_str());
|
||||
let resolved_model = decision
|
||||
.mapped_model
|
||||
.as_deref()
|
||||
.or(decision.model_name.as_deref())
|
||||
.unwrap_or(input.requested_model.as_str());
|
||||
let original_provider_request_body = decision.provider_request_body.clone();
|
||||
let mut provider_request_body = original_provider_request_body
|
||||
.clone()
|
||||
.unwrap_or(serde_json::Value::Null);
|
||||
let mut provider_headers = btree_headers_to_header_map(&decision.provider_request_headers)?;
|
||||
let provider_headers_json = headers_to_routing_value(&provider_headers);
|
||||
let policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput {
|
||||
group_id: context.group_id.as_deref(),
|
||||
group_version: context.group_version,
|
||||
group_config_json: &context.group_config_json,
|
||||
selection_source: context.selection_source.as_str(),
|
||||
requested_model: input.requested_model.as_str(),
|
||||
resolved_model,
|
||||
api_format: provider_api_format,
|
||||
user_id: Some(input.auth_context.user_id.as_str()),
|
||||
api_key_id: Some(input.auth_context.api_key_id.as_str()),
|
||||
headers: &provider_headers_json,
|
||||
body: &provider_request_body,
|
||||
phase: RoutingRulePhase::ProviderRequest,
|
||||
})?;
|
||||
ensure_report_context_routing_trace(input, decision, &policy);
|
||||
if policy.mutation_plan.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
if original_provider_request_body.is_none() && !policy.mutation_plan.body_patch.is_empty() {
|
||||
return Err(GatewayError::Client {
|
||||
status: StatusCode::BAD_REQUEST,
|
||||
message: "routing provider_request body patch cannot be applied to a binary or empty upstream body".to_string(),
|
||||
});
|
||||
}
|
||||
apply_routing_mutation_plan(
|
||||
&mut provider_request_body,
|
||||
&mut provider_headers,
|
||||
&policy.mutation_plan,
|
||||
)?;
|
||||
decision.provider_request_headers = header_map_to_btree_headers(&provider_headers);
|
||||
if original_provider_request_body.is_some() {
|
||||
decision.provider_request_body = Some(provider_request_body);
|
||||
}
|
||||
update_report_context_provider_request_mutation(decision, &policy);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
struct GatewayAuthenticatedDecisionInputPort<'a> {
|
||||
state: PlannerAppState<'a>,
|
||||
now_unix_secs: u64,
|
||||
@@ -99,9 +205,154 @@ pub(crate) fn build_local_requested_model_decision_input(
|
||||
required_capabilities: resolved_input.required_capabilities,
|
||||
request_auth_channel: None,
|
||||
client_session_affinity: None,
|
||||
routing_policy: None,
|
||||
routing_trace_seed: None,
|
||||
routing_context: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
|
||||
state: &AppState,
|
||||
parts: &http::request::Parts,
|
||||
input: &mut LocalRequestedModelDecisionInput,
|
||||
body_json: &Value,
|
||||
client_api_format: &str,
|
||||
) -> Result<(), GatewayError> {
|
||||
let explicit_group = routing_header_value_str(&parts.headers, ROUTING_GROUP_HEADER);
|
||||
let selected_group = match state.routing_group_read_repository() {
|
||||
Some(repository) => {
|
||||
let user_group_ids = match state
|
||||
.list_user_groups_for_user(&input.auth_context.user_id)
|
||||
.await
|
||||
{
|
||||
Ok(groups) => groups.into_iter().map(|group| group.id).collect::<Vec<_>>(),
|
||||
Err(error) => {
|
||||
warn!(
|
||||
user_id = %input.auth_context.user_id,
|
||||
error = ?error,
|
||||
"gateway routing profile user group lookup failed"
|
||||
);
|
||||
Vec::new()
|
||||
}
|
||||
};
|
||||
let selection = select_gateway_routing_group(
|
||||
repository.as_ref(),
|
||||
GatewayRoutingSelectionInput {
|
||||
explicit_group: explicit_group.as_deref(),
|
||||
user_id: Some(input.auth_context.user_id.as_str()),
|
||||
api_key_id: Some(input.auth_context.api_key_id.as_str()),
|
||||
user_group_ids: &user_group_ids,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.map_err(routing_selection_error)?;
|
||||
selection.group.map(|group| {
|
||||
(
|
||||
Some(group.id),
|
||||
Some(group.version),
|
||||
group.config_json,
|
||||
selection.source,
|
||||
)
|
||||
})
|
||||
}
|
||||
None => {
|
||||
if explicit_group
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
{
|
||||
return Err(routing_selection_error(
|
||||
GatewayRoutingSelectionError::NotFound(explicit_group.unwrap_or_default()),
|
||||
));
|
||||
}
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
let Some((group_id, group_version, group_config_json, selection_source)) = selected_group
|
||||
else {
|
||||
input.client_session_affinity =
|
||||
client_session_affinity_from_request(&parts.headers, Some(body_json));
|
||||
input.routing_policy = None;
|
||||
input.routing_trace_seed = None;
|
||||
input.routing_context = None;
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
let headers_json = headers_to_routing_value(&parts.headers);
|
||||
let policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput {
|
||||
group_id: group_id.as_deref(),
|
||||
group_version,
|
||||
group_config_json: &group_config_json,
|
||||
selection_source: selection_source.as_str(),
|
||||
requested_model: input.requested_model.as_str(),
|
||||
resolved_model: input.requested_model.as_str(),
|
||||
api_format: client_api_format,
|
||||
user_id: Some(input.auth_context.user_id.as_str()),
|
||||
api_key_id: Some(input.auth_context.api_key_id.as_str()),
|
||||
headers: &headers_json,
|
||||
body: body_json,
|
||||
phase: RoutingRulePhase::ClientRequest,
|
||||
})?;
|
||||
let mut effective_body_json = body_json.clone();
|
||||
let mut effective_headers = parts.headers.clone();
|
||||
apply_routing_mutation_plan(
|
||||
&mut effective_body_json,
|
||||
&mut effective_headers,
|
||||
&policy.mutation_plan,
|
||||
)?;
|
||||
|
||||
let mut requested_model_changed = false;
|
||||
if let Some(mut mutated_model) = extract_standard_requested_model(&effective_body_json) {
|
||||
mutated_model = mutated_model.trim().to_string();
|
||||
if !mutated_model.is_empty() && mutated_model != input.requested_model {
|
||||
input.requested_model = mutated_model;
|
||||
requested_model_changed = true;
|
||||
}
|
||||
}
|
||||
if requested_model_changed {
|
||||
input.required_capabilities = PlannerAppState::new(state)
|
||||
.resolve_request_candidate_required_capabilities(
|
||||
&input.auth_context.user_id,
|
||||
&input.auth_context.api_key_id,
|
||||
Some(input.requested_model.as_str()),
|
||||
input.required_capabilities.as_ref(),
|
||||
)
|
||||
.await;
|
||||
}
|
||||
|
||||
let effective_headers_json = headers_to_routing_value(&effective_headers);
|
||||
input.client_session_affinity =
|
||||
client_session_affinity_from_request(&effective_headers, Some(&effective_body_json));
|
||||
let mut final_policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput {
|
||||
group_id: group_id.as_deref(),
|
||||
group_version,
|
||||
group_config_json: &group_config_json,
|
||||
selection_source: selection_source.as_str(),
|
||||
requested_model: input.requested_model.as_str(),
|
||||
resolved_model: input.requested_model.as_str(),
|
||||
api_format: client_api_format,
|
||||
user_id: Some(input.auth_context.user_id.as_str()),
|
||||
api_key_id: Some(input.auth_context.api_key_id.as_str()),
|
||||
headers: &effective_headers_json,
|
||||
body: &effective_body_json,
|
||||
phase: RoutingRulePhase::ClientRequest,
|
||||
})?;
|
||||
final_policy.mutation_plan = policy.mutation_plan.clone();
|
||||
input.routing_trace_seed = Some(build_routing_trace_seed(&final_policy, client_api_format));
|
||||
input.routing_policy = Some(final_policy);
|
||||
input.routing_context = Some(LocalRoutingRequestContext {
|
||||
group_id,
|
||||
group_version,
|
||||
group_config_json,
|
||||
selection_source,
|
||||
client_api_format: client_api_format.to_string(),
|
||||
effective_body_json,
|
||||
effective_headers,
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(crate) fn build_local_authenticated_decision_input(
|
||||
resolved_input: ResolvedLocalDecisionAuthInput,
|
||||
) -> LocalAuthenticatedDecisionInput {
|
||||
@@ -132,3 +383,506 @@ pub(crate) async fn resolve_local_authenticated_decision_input(
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
fn routing_selection_error(error: GatewayRoutingSelectionError) -> GatewayError {
|
||||
GatewayError::Client {
|
||||
status: StatusCode::FORBIDDEN,
|
||||
message: error.to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
fn headers_to_routing_value(headers: &http::HeaderMap) -> Value {
|
||||
let mut object = serde_json::Map::new();
|
||||
for (name, value) in headers {
|
||||
if let Ok(value) = value.to_str() {
|
||||
object.insert(name.as_str().to_ascii_lowercase(), json!(value));
|
||||
}
|
||||
}
|
||||
Value::Object(object)
|
||||
}
|
||||
|
||||
fn routing_header_value_str(headers: &http::HeaderMap, key: &str) -> Option<String> {
|
||||
headers
|
||||
.get(key)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn btree_headers_to_header_map(
|
||||
headers: &BTreeMap<String, String>,
|
||||
) -> Result<HeaderMap, GatewayError> {
|
||||
let mut output = HeaderMap::new();
|
||||
for (name, value) in headers {
|
||||
let name = HeaderName::from_bytes(name.as_bytes()).map_err(|err| GatewayError::Client {
|
||||
status: StatusCode::BAD_REQUEST,
|
||||
message: format!("invalid provider request header name in routing mutation: {err}"),
|
||||
})?;
|
||||
let value = HeaderValue::from_str(value).map_err(|err| GatewayError::Client {
|
||||
status: StatusCode::BAD_REQUEST,
|
||||
message: format!("invalid provider request header value in routing mutation: {err}"),
|
||||
})?;
|
||||
output.insert(name, value);
|
||||
}
|
||||
Ok(output)
|
||||
}
|
||||
|
||||
fn header_map_to_btree_headers(headers: &HeaderMap) -> BTreeMap<String, String> {
|
||||
headers
|
||||
.iter()
|
||||
.filter_map(|(name, value)| {
|
||||
value
|
||||
.to_str()
|
||||
.ok()
|
||||
.map(|value| (name.as_str().to_string(), value.to_string()))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn update_report_context_provider_request_mutation(
|
||||
decision: &mut AiExecutionDecision,
|
||||
policy: &ResolvedRoutingPolicy,
|
||||
) {
|
||||
let Some(serde_json::Value::Object(object)) = decision.report_context.as_mut() else {
|
||||
return;
|
||||
};
|
||||
let body_paths = policy
|
||||
.mutation_plan
|
||||
.body_patch
|
||||
.iter()
|
||||
.map(|operation| operation.path().to_string())
|
||||
.collect::<Vec<_>>();
|
||||
let header_names = policy
|
||||
.mutation_plan
|
||||
.header_patch
|
||||
.iter()
|
||||
.map(|operation| operation.name().to_string())
|
||||
.collect::<Vec<_>>();
|
||||
let trace_patch_summary = serde_json::json!({
|
||||
"body_paths": body_paths,
|
||||
"header_names": header_names,
|
||||
});
|
||||
if let Some(serde_json::Value::Object(routing_trace)) = object.get_mut("routing_trace") {
|
||||
routing_trace.insert(
|
||||
"provider_request_patch_summary".to_string(),
|
||||
trace_patch_summary.clone(),
|
||||
);
|
||||
}
|
||||
object.insert(
|
||||
"provider_request_headers".to_string(),
|
||||
serde_json::json!(decision.provider_request_headers),
|
||||
);
|
||||
object.insert(
|
||||
"routing_provider_request_patch_summary".to_string(),
|
||||
serde_json::json!({
|
||||
"body_paths": trace_patch_summary["body_paths"].clone(),
|
||||
"header_names": trace_patch_summary["header_names"].clone(),
|
||||
"matched_rules": policy
|
||||
.matched_rules
|
||||
.iter()
|
||||
.map(|rule| rule.id.clone())
|
||||
.collect::<Vec<_>>()
|
||||
}),
|
||||
);
|
||||
}
|
||||
|
||||
fn ensure_report_context_routing_trace(
|
||||
input: &LocalRequestedModelDecisionInput,
|
||||
decision: &mut AiExecutionDecision,
|
||||
policy: &ResolvedRoutingPolicy,
|
||||
) {
|
||||
let Some(serde_json::Value::Object(object)) = decision.report_context.as_mut() else {
|
||||
return;
|
||||
};
|
||||
if object.get("routing_trace").is_some() {
|
||||
return;
|
||||
}
|
||||
|
||||
let client_api_format = decision
|
||||
.client_api_format
|
||||
.as_deref()
|
||||
.or_else(|| {
|
||||
input
|
||||
.routing_context
|
||||
.as_ref()
|
||||
.map(|context| context.client_api_format.as_str())
|
||||
})
|
||||
.unwrap_or_default();
|
||||
let mut trace = input
|
||||
.routing_trace_seed
|
||||
.clone()
|
||||
.unwrap_or_else(|| build_routing_trace_seed(policy, client_api_format));
|
||||
|
||||
let candidate_group_id = object
|
||||
.get("candidate_group_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let pool_key_index = object
|
||||
.get("pool_key_index")
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|value| u32::try_from(value).ok());
|
||||
let is_pool_expansion = candidate_group_id.is_some() && pool_key_index.is_some();
|
||||
let candidate_kind = if is_pool_expansion {
|
||||
CandidateKind::PoolGroup
|
||||
} else {
|
||||
CandidateKind::Provider
|
||||
};
|
||||
let provider_id = candidate_group_id
|
||||
.clone()
|
||||
.or_else(|| decision.provider_id.clone())
|
||||
.unwrap_or_default();
|
||||
let endpoint_id = decision.endpoint_id.clone().unwrap_or_default();
|
||||
let model_id = object
|
||||
.get("model_id")
|
||||
.and_then(Value::as_str)
|
||||
.map(ToOwned::to_owned)
|
||||
.or_else(|| decision.mapped_model.clone())
|
||||
.or_else(|| decision.model_name.clone())
|
||||
.unwrap_or_else(|| input.requested_model.clone());
|
||||
let key_id = decision.key_id.clone().filter(|_| !is_pool_expansion);
|
||||
let provider_priority = object
|
||||
.get("provider_priority")
|
||||
.and_then(Value::as_i64)
|
||||
.and_then(|value| i32::try_from(value).ok())
|
||||
.unwrap_or_default();
|
||||
let key_priority = object
|
||||
.get("priority_slot")
|
||||
.and_then(Value::as_i64)
|
||||
.and_then(|value| i32::try_from(value).ok())
|
||||
.unwrap_or_default();
|
||||
trace.global_candidates.push(RoutingCandidateTrace {
|
||||
candidate_kind,
|
||||
provider_id: provider_id.clone(),
|
||||
endpoint_id,
|
||||
model_id: model_id.clone(),
|
||||
key_id: key_id.clone(),
|
||||
ranking_vector: rank_vector_for_candidate(
|
||||
&policy.ranking_overlay,
|
||||
&RoutingCandidateFacts {
|
||||
candidate_kind,
|
||||
provider_id: provider_id.clone(),
|
||||
endpoint_id: decision.endpoint_id.clone().unwrap_or_default(),
|
||||
model_id,
|
||||
key_id,
|
||||
provider_priority,
|
||||
key_priority,
|
||||
},
|
||||
),
|
||||
skip_reason: None,
|
||||
selected_order: object
|
||||
.get("candidate_index")
|
||||
.and_then(Value::as_u64)
|
||||
.and_then(|value| u32::try_from(value).ok()),
|
||||
});
|
||||
|
||||
if is_pool_expansion {
|
||||
if let (Some(pool_group_id), Some(key_id)) = (candidate_group_id, decision.key_id.clone()) {
|
||||
trace.pool_expansion.push(RoutingPoolExpansionTrace {
|
||||
pool_group_id,
|
||||
key_id,
|
||||
pool_ranking_vector: Vec::new(),
|
||||
pool_skip_reason: None,
|
||||
selected_order: pool_key_index,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
object.insert("routing_trace".to_string(), serde_json::json!(trace));
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn sample_auth_context() -> ExecutionRuntimeAuthContext {
|
||||
ExecutionRuntimeAuthContext {
|
||||
user_id: "user-1".to_string(),
|
||||
api_key_id: "api-key-1".to_string(),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
balance_remaining: None,
|
||||
access_allowed: true,
|
||||
api_key_is_standalone: false,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_auth_snapshot() -> GatewayAuthApiKeySnapshot {
|
||||
GatewayAuthApiKeySnapshot {
|
||||
user_id: "user-1".to_string(),
|
||||
username: "alice".to_string(),
|
||||
email: None,
|
||||
user_role: "user".to_string(),
|
||||
user_auth_source: "local".to_string(),
|
||||
user_is_active: true,
|
||||
user_is_deleted: false,
|
||||
user_rate_limit: None,
|
||||
user_allowed_providers: None,
|
||||
user_allowed_api_formats: None,
|
||||
user_allowed_models: None,
|
||||
api_key_id: "api-key-1".to_string(),
|
||||
api_key_name: Some("default".to_string()),
|
||||
api_key_is_active: true,
|
||||
api_key_is_locked: false,
|
||||
api_key_is_standalone: false,
|
||||
api_key_rate_limit: None,
|
||||
api_key_concurrent_limit: None,
|
||||
api_key_expires_at_unix_secs: None,
|
||||
api_key_allowed_providers: None,
|
||||
api_key_allowed_api_formats: None,
|
||||
api_key_allowed_models: None,
|
||||
currently_usable: true,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_decision_input() -> LocalRequestedModelDecisionInput {
|
||||
LocalRequestedModelDecisionInput {
|
||||
auth_context: sample_auth_context(),
|
||||
requested_model: "gpt-5".to_string(),
|
||||
auth_snapshot: sample_auth_snapshot(),
|
||||
required_capabilities: None,
|
||||
request_auth_channel: None,
|
||||
client_session_affinity: None,
|
||||
routing_policy: None,
|
||||
routing_trace_seed: None,
|
||||
routing_context: Some(LocalRoutingRequestContext {
|
||||
group_id: Some("group-1".to_string()),
|
||||
group_version: Some(3),
|
||||
selection_source: "explicit_header".to_string(),
|
||||
client_api_format: "openai:chat".to_string(),
|
||||
effective_body_json: json!({"model":"gpt-5"}),
|
||||
effective_headers: HeaderMap::new(),
|
||||
group_config_json: json!({
|
||||
"allowed_models": ["gpt-5"],
|
||||
"rules": [{
|
||||
"id": "provider-patch",
|
||||
"priority": 1,
|
||||
"enabled": true,
|
||||
"phase": "provider_request",
|
||||
"conditions": {},
|
||||
"actions": [
|
||||
{
|
||||
"type": "json_patch_body",
|
||||
"patch": [{
|
||||
"op": "add",
|
||||
"path": "/metadata/routing",
|
||||
"value": "provider"
|
||||
}]
|
||||
},
|
||||
{
|
||||
"type": "patch_headers",
|
||||
"patch": [{
|
||||
"op": "set",
|
||||
"name": "x-provider-route",
|
||||
"value": "provider"
|
||||
}]
|
||||
}
|
||||
]
|
||||
}]
|
||||
}),
|
||||
}),
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_decision() -> AiExecutionDecision {
|
||||
AiExecutionDecision {
|
||||
action: "execution_runtime_sync_decision".to_string(),
|
||||
decision_kind: Some("openai_chat_sync".to_string()),
|
||||
execution_strategy: None,
|
||||
conversion_mode: None,
|
||||
request_id: Some("trace-1".to_string()),
|
||||
candidate_id: Some("candidate-1".to_string()),
|
||||
provider_name: Some("provider".to_string()),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
endpoint_id: Some("endpoint-1".to_string()),
|
||||
key_id: Some("key-1".to_string()),
|
||||
upstream_base_url: None,
|
||||
upstream_url: None,
|
||||
provider_request_method: None,
|
||||
auth_header: None,
|
||||
auth_value: None,
|
||||
provider_api_format: Some("openai:chat".to_string()),
|
||||
client_api_format: Some("openai:chat".to_string()),
|
||||
provider_contract: None,
|
||||
client_contract: None,
|
||||
model_name: Some("gpt-5".to_string()),
|
||||
mapped_model: Some("gpt-5".to_string()),
|
||||
prompt_cache_key: None,
|
||||
extra_headers: BTreeMap::new(),
|
||||
provider_request_headers: BTreeMap::from([(
|
||||
"content-type".to_string(),
|
||||
"application/json".to_string(),
|
||||
)]),
|
||||
provider_request_body: Some(json!({"model":"gpt-5","metadata":{}})),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
upstream_is_stream: false,
|
||||
report_kind: Some("local_sync_success".to_string()),
|
||||
report_context: Some(json!({
|
||||
"candidate_index": 0,
|
||||
"retry_index": 0,
|
||||
"model_id": "model-1"
|
||||
})),
|
||||
auth_context: Some(sample_auth_context()),
|
||||
}
|
||||
}
|
||||
|
||||
fn set_provider_request_rules(input: &mut LocalRequestedModelDecisionInput, actions: Value) {
|
||||
let config = json!({
|
||||
"allowed_models": ["gpt-5"],
|
||||
"rules": [{
|
||||
"id": "provider-patch",
|
||||
"priority": 1,
|
||||
"enabled": true,
|
||||
"phase": "provider_request",
|
||||
"conditions": {},
|
||||
"actions": actions
|
||||
}]
|
||||
});
|
||||
input
|
||||
.routing_context
|
||||
.as_mut()
|
||||
.expect("sample input should include routing context")
|
||||
.group_config_json = config;
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_request_routing_policy_mutates_decision_body_headers_and_report_context() {
|
||||
let input = sample_decision_input();
|
||||
let mut decision = sample_decision();
|
||||
|
||||
apply_provider_request_routing_policy_to_decision(&input, &mut decision)
|
||||
.expect("provider routing mutation should apply");
|
||||
|
||||
assert_eq!(
|
||||
decision.provider_request_body.as_ref().unwrap()["metadata"]["routing"],
|
||||
json!("provider")
|
||||
);
|
||||
assert_eq!(
|
||||
decision
|
||||
.provider_request_headers
|
||||
.get("x-provider-route")
|
||||
.map(String::as_str),
|
||||
Some("provider")
|
||||
);
|
||||
let report_context = decision.report_context.as_ref().unwrap();
|
||||
assert_eq!(
|
||||
report_context["routing_provider_request_patch_summary"]["matched_rules"],
|
||||
json!(["provider-patch"])
|
||||
);
|
||||
assert_eq!(
|
||||
report_context["routing_trace"]["provider_request_patch_summary"]["body_paths"],
|
||||
json!(["/metadata/routing"])
|
||||
);
|
||||
assert_eq!(
|
||||
report_context["routing_trace"]["global_candidates"][0]["provider_id"],
|
||||
json!("provider-1")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_request_routing_policy_rejects_body_patch_without_json_body() {
|
||||
let input = sample_decision_input();
|
||||
let mut decision = sample_decision();
|
||||
decision.provider_request_body = None;
|
||||
decision.provider_request_body_base64 = Some("AA==".to_string());
|
||||
|
||||
let error = apply_provider_request_routing_policy_to_decision(&input, &mut decision)
|
||||
.expect_err("provider body patch should reject binary upstream bodies");
|
||||
|
||||
match error {
|
||||
GatewayError::Client { status, message } => {
|
||||
assert_eq!(status, StatusCode::BAD_REQUEST);
|
||||
assert!(message.contains("binary or empty upstream body"));
|
||||
}
|
||||
other => panic!("unexpected error: {other:?}"),
|
||||
}
|
||||
assert!(
|
||||
decision
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|context| context.get("routing_trace"))
|
||||
.is_some(),
|
||||
"failed provider_request mutation should still seed routing trace"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_request_routing_policy_allows_header_patch_without_json_body() {
|
||||
let mut input = sample_decision_input();
|
||||
set_provider_request_rules(
|
||||
&mut input,
|
||||
json!([{
|
||||
"type": "patch_headers",
|
||||
"patch": [{
|
||||
"op": "set",
|
||||
"name": "x-provider-route",
|
||||
"value": "header-only"
|
||||
}]
|
||||
}]),
|
||||
);
|
||||
let mut decision = sample_decision();
|
||||
decision.provider_request_body = None;
|
||||
decision.provider_request_body_base64 = Some("AA==".to_string());
|
||||
|
||||
apply_provider_request_routing_policy_to_decision(&input, &mut decision)
|
||||
.expect("header-only provider routing mutation should apply without JSON body");
|
||||
|
||||
assert_eq!(decision.provider_request_body, None);
|
||||
assert_eq!(
|
||||
decision
|
||||
.provider_request_headers
|
||||
.get("x-provider-route")
|
||||
.map(String::as_str),
|
||||
Some("header-only")
|
||||
);
|
||||
assert_eq!(
|
||||
decision.report_context.as_ref().unwrap()["routing_trace"]
|
||||
["provider_request_patch_summary"]["header_names"],
|
||||
json!(["x-provider-route"])
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_request_routing_trace_records_pool_expansion_candidate() {
|
||||
let input = sample_decision_input();
|
||||
let mut decision = sample_decision();
|
||||
decision.report_context = Some(json!({
|
||||
"candidate_index": 2,
|
||||
"retry_index": 2,
|
||||
"model_id": "model-1",
|
||||
"candidate_group_id": "pool-group-1",
|
||||
"pool_key_index": 1,
|
||||
"provider_priority": 7,
|
||||
"priority_slot": 3
|
||||
}));
|
||||
|
||||
apply_provider_request_routing_policy_to_decision(&input, &mut decision)
|
||||
.expect("provider routing mutation should seed pool trace");
|
||||
|
||||
let routing_trace = &decision.report_context.as_ref().unwrap()["routing_trace"];
|
||||
assert_eq!(
|
||||
routing_trace["global_candidates"][0]["candidate_kind"],
|
||||
json!("pool_group")
|
||||
);
|
||||
assert_eq!(
|
||||
routing_trace["global_candidates"][0]["provider_id"],
|
||||
json!("pool-group-1")
|
||||
);
|
||||
assert_eq!(routing_trace["global_candidates"][0]["key_id"], Value::Null);
|
||||
assert_eq!(
|
||||
routing_trace["pool_expansion"][0]["pool_group_id"],
|
||||
json!("pool-group-1")
|
||||
);
|
||||
assert_eq!(routing_trace["pool_expansion"][0]["key_id"], json!("key-1"));
|
||||
assert_eq!(
|
||||
routing_trace["pool_expansion"][0]["selected_order"],
|
||||
json!(1)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -33,7 +33,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
@@ -55,6 +55,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
@@ -70,7 +71,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
@@ -100,7 +101,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
@@ -122,6 +123,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
@@ -137,7 +139,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
|
||||
maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ use crate::ai_serving::planner::candidate_metadata::{
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
attach_routing_policy_to_local_requested_model_input,
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
@@ -39,19 +40,21 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<LocalSameFormatProviderDecisionInput> {
|
||||
) -> Result<Option<LocalSameFormatProviderDecisionInput>, GatewayError> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
return None;
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let requested_model = extract_requested_model_from_request(
|
||||
let Some(requested_model) = extract_requested_model_from_request(
|
||||
parts,
|
||||
body_json,
|
||||
spec_metadata
|
||||
.requested_model_family
|
||||
.expect("same-format provider specs should declare requested-model family"),
|
||||
)?;
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
@@ -62,7 +65,7 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Ok(None) => return Ok(None),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
@@ -70,14 +73,31 @@ pub(crate) async fn resolve_local_same_format_provider_decision_input(
|
||||
error = ?err,
|
||||
"gateway local same-format decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
|
||||
input.request_auth_channel = decision.request_auth_channel.clone();
|
||||
input.client_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
|
||||
Some(input)
|
||||
if let Err(err) = attach_routing_policy_to_local_requested_model_input(
|
||||
state,
|
||||
parts,
|
||||
&mut input,
|
||||
body_json,
|
||||
spec_metadata.api_format,
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local same-format decision routing profile resolution failed"
|
||||
);
|
||||
return Err(err);
|
||||
}
|
||||
Ok(Some(input))
|
||||
}
|
||||
|
||||
pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
@@ -114,6 +134,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
@@ -211,6 +232,7 @@ pub(crate) async fn build_local_same_format_provider_candidate_attempt_source<'a
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
|
||||
@@ -6,6 +6,7 @@ use crate::ai_serving::planner::candidate_materialization::{
|
||||
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic,
|
||||
};
|
||||
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
@@ -22,7 +23,7 @@ use crate::ai_serving::transport::{
|
||||
};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
|
||||
AiExecutionDecision, AppState,
|
||||
AiExecutionDecision, AppState, GatewayError,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
|
||||
@@ -40,7 +41,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
input: &LocalSameFormatProviderDecisionInput,
|
||||
attempt: LocalSameFormatProviderCandidateAttempt,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let spec_metadata = local_same_format_provider_spec_metadata(spec);
|
||||
let LocalSameFormatProviderCandidateAttempt {
|
||||
eligible,
|
||||
@@ -51,10 +52,13 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
let candidate = &eligible.candidate;
|
||||
let (execution_strategy, conversion_mode) =
|
||||
ai_local_execution_contract_for_formats(spec_metadata.api_format, spec_metadata.api_format);
|
||||
let resolved = resolve_local_same_format_provider_candidate_payload_parts(
|
||||
let Some(resolved) = resolve_local_same_format_provider_candidate_payload_parts(
|
||||
state, parts, trace_id, body_json, input, &attempt, spec,
|
||||
)
|
||||
.await?;
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let prompt_cache_key = resolved
|
||||
.provider_request_body
|
||||
@@ -85,6 +89,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
);
|
||||
}
|
||||
let provider_api_format = resolved.provider_api_format.clone();
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let report_context = append_local_failover_policy_to_value(
|
||||
append_execution_contract_fields_to_value(
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
@@ -112,7 +117,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
body_rules: resolved.transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
original_headers: effective_headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
@@ -151,41 +156,41 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
||||
provider_request_body,
|
||||
} = resolved;
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.to_string(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header,
|
||||
auth_value,
|
||||
provider_api_format,
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream,
|
||||
report_kind: Some(report_kind.to_string()),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.to_string(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header,
|
||||
auth_value,
|
||||
provider_api_format,
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream,
|
||||
report_kind: Some(report_kind.to_string()),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
});
|
||||
apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
|
||||
Ok(Some(decision))
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_same_format_provider_candidate(
|
||||
|
||||
@@ -125,6 +125,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
|
||||
Some(&input.requested_model),
|
||||
)
|
||||
.await;
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
|
||||
let Some(mut base_provider_request_body) =
|
||||
super::super::request::build_same_format_provider_request_body(
|
||||
@@ -133,7 +134,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
|
||||
&prepared.mapped_model,
|
||||
spec,
|
||||
prepared.transport.endpoint.body_rules.as_ref(),
|
||||
Some(&parts.headers),
|
||||
Some(effective_headers),
|
||||
prepared.upstream_is_stream,
|
||||
prepared.force_body_stream_field,
|
||||
prepared.kiro_auth.as_ref(),
|
||||
@@ -279,7 +280,7 @@ pub(crate) async fn resolve_local_same_format_provider_candidate_payload_parts(
|
||||
.unwrap_or_default();
|
||||
let Some(provider_request_headers) =
|
||||
build_same_format_provider_headers(SameFormatProviderHeadersInput {
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: body_json,
|
||||
header_rules: prepared.transport.endpoint.header_rules.as_ref(),
|
||||
|
||||
@@ -31,7 +31,7 @@ pub(crate) struct LocalSameFormatProviderSyncAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
trace_id: &'a str,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_json: serde_json::Value,
|
||||
input: LocalSameFormatProviderDecisionInput,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
requested_model_family: RequestedModelFamily,
|
||||
@@ -42,7 +42,7 @@ pub(crate) struct LocalSameFormatProviderStreamAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
trace_id: &'a str,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_json: serde_json::Value,
|
||||
input: LocalSameFormatProviderDecisionInput,
|
||||
spec: LocalSameFormatProviderSpec,
|
||||
requested_model_family: RequestedModelFamily,
|
||||
@@ -64,7 +64,7 @@ pub(crate) async fn build_local_sync_attempt_source<'a>(
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
@@ -85,8 +85,13 @@ pub(crate) async fn build_local_sync_attempt_source<'a>(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
let (candidates, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
&effective_body_json,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
@@ -103,7 +108,7 @@ pub(crate) async fn build_local_sync_attempt_source<'a>(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
body_json: effective_body_json,
|
||||
input,
|
||||
spec,
|
||||
requested_model_family,
|
||||
@@ -128,7 +133,7 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
@@ -149,8 +154,13 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
let (candidates, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
&effective_body_json,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress_preserving_candidate_signal(
|
||||
@@ -167,7 +177,7 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
body_json: effective_body_json,
|
||||
input,
|
||||
spec,
|
||||
requested_model_family,
|
||||
@@ -244,12 +254,12 @@ impl LocalSameFormatProviderSyncAttemptSource<'_> {
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
&self.input,
|
||||
attempt,
|
||||
self.spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -257,7 +267,7 @@ impl LocalSameFormatProviderSyncAttemptSource<'_> {
|
||||
match build_sync_plan_from_requested_model_family(
|
||||
self.requested_model_family,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
payload,
|
||||
) {
|
||||
Ok(value) => Ok(value),
|
||||
@@ -282,12 +292,12 @@ impl LocalSameFormatProviderStreamAttemptSource<'_> {
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
&self.input,
|
||||
attempt,
|
||||
self.spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -295,7 +305,7 @@ impl LocalSameFormatProviderStreamAttemptSource<'_> {
|
||||
match build_stream_plan_from_requested_model_family(
|
||||
self.requested_model_family,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
payload,
|
||||
) {
|
||||
Ok(value) => Ok(value),
|
||||
@@ -326,7 +336,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
@@ -347,6 +357,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
@@ -365,7 +376,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
@@ -411,7 +422,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
let Some(input) = resolve_local_same_format_provider_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
@@ -432,6 +443,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
)
|
||||
@@ -450,7 +462,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
@@ -29,7 +29,7 @@ use self::support::{
|
||||
pub(crate) struct LocalGeminiFilesSyncAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_json: serde_json::Value,
|
||||
body_base64: Option<&'a str>,
|
||||
body_is_empty: bool,
|
||||
trace_id: &'a str,
|
||||
@@ -110,10 +110,11 @@ pub(crate) async fn build_local_gemini_files_sync_attempt_source_for_kind<'a>(
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
let (candidates, candidate_count) =
|
||||
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
|
||||
if candidate_count == 0 {
|
||||
@@ -124,7 +125,7 @@ pub(crate) async fn build_local_gemini_files_sync_attempt_source_for_kind<'a>(
|
||||
LocalGeminiFilesSyncAttemptSource {
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_json: effective_body_json,
|
||||
body_base64,
|
||||
body_is_empty,
|
||||
trace_id,
|
||||
@@ -148,7 +149,7 @@ pub(crate) async fn build_local_gemini_files_stream_attempt_source_for_kind<'a>(
|
||||
};
|
||||
|
||||
let Some(input) =
|
||||
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await
|
||||
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -226,7 +227,7 @@ impl LocalGeminiFilesSyncAttemptSource<'_> {
|
||||
let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
self.body_base64,
|
||||
self.body_is_empty,
|
||||
self.trace_id,
|
||||
@@ -234,7 +235,7 @@ impl LocalGeminiFilesSyncAttemptSource<'_> {
|
||||
attempt,
|
||||
self.spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -272,7 +273,7 @@ impl LocalGeminiFilesStreamAttemptSource<'_> {
|
||||
attempt,
|
||||
self.spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -313,10 +314,11 @@ pub(crate) async fn maybe_build_sync_local_gemini_files_decision_payload(
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
|
||||
let (mut source, _) =
|
||||
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
|
||||
@@ -333,7 +335,7 @@ pub(crate) async fn maybe_build_sync_local_gemini_files_decision_payload(
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
@@ -354,7 +356,7 @@ pub(crate) async fn maybe_build_stream_local_gemini_files_decision_payload(
|
||||
};
|
||||
|
||||
let Some(input) =
|
||||
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await
|
||||
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -375,7 +377,7 @@ pub(crate) async fn maybe_build_stream_local_gemini_files_decision_payload(
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
@@ -402,10 +404,11 @@ async fn build_local_sync_plan_and_reports(
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
|
||||
let (mut source, _) =
|
||||
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
|
||||
@@ -423,7 +426,7 @@ async fn build_local_sync_plan_and_reports(
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
@@ -454,7 +457,7 @@ async fn build_local_stream_plan_and_reports(
|
||||
) -> Result<Vec<AiStreamAttempt>, GatewayError> {
|
||||
let spec_metadata = local_gemini_files_spec_metadata(spec);
|
||||
let Some(input) =
|
||||
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await
|
||||
resolve_local_gemini_files_decision_input(state, parts, None, trace_id, decision).await?
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
@@ -476,7 +479,7 @@ async fn build_local_stream_plan_and_reports(
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
use serde_json::json;
|
||||
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, LocalExecutionReportContextParts,
|
||||
};
|
||||
@@ -12,7 +13,7 @@ use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_profile,
|
||||
};
|
||||
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
|
||||
use crate::{AiExecutionDecision, AppState};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use super::request::resolve_local_gemini_files_candidate_payload_parts;
|
||||
use super::support::{
|
||||
@@ -31,7 +32,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
input: &LocalGeminiFilesDecisionInput,
|
||||
attempt: LocalGeminiFilesCandidateAttempt,
|
||||
spec: LocalGeminiFilesSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let spec_metadata = local_gemini_files_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let attempt_identity = attempt.attempt_identity();
|
||||
@@ -46,7 +47,10 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
&attempt,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
.await;
|
||||
let Some(resolved) = resolved else {
|
||||
return Ok(None);
|
||||
};
|
||||
let LocalGeminiFilesCandidateAttempt {
|
||||
eligible,
|
||||
candidate_id,
|
||||
@@ -69,6 +73,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
}
|
||||
extra_fields.insert("file_key_id".to_string(), json!(candidate.key_id));
|
||||
extra_fields.insert("file_name".to_string(), json!(resolved.file_name));
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let report_context = build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &input.auth_context,
|
||||
request_id: trace_id,
|
||||
@@ -94,7 +99,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
body_rules: transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: None,
|
||||
provider_request_headers: None,
|
||||
original_headers: &parts.headers,
|
||||
original_headers: effective_headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
@@ -119,45 +124,44 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
||||
file_name: _,
|
||||
} = resolved;
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT.to_string(),
|
||||
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT.to_string(),
|
||||
model_name: "gemini-files".to_string(),
|
||||
mapped_model: candidate.selected_provider_model_name.clone(),
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
provider_request_body_base64,
|
||||
content_type: parts
|
||||
.headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream: spec_metadata.require_streaming,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT.to_string(),
|
||||
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT.to_string(),
|
||||
model_name: "gemini-files".to_string(),
|
||||
mapped_model: candidate.selected_provider_model_name.clone(),
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
provider_request_body_base64,
|
||||
content_type: effective_headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream: spec_metadata.require_streaming,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
});
|
||||
apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
|
||||
Ok(Some(decision))
|
||||
}
|
||||
|
||||
@@ -45,6 +45,7 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
|
||||
let spec_metadata = local_gemini_files_spec_metadata(spec);
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
|
||||
if let Some(skip_reason) =
|
||||
gemini_files_transport_unsupported_reason(transport, GEMINI_FILES_CANDIDATE_API_FORMAT)
|
||||
@@ -103,7 +104,7 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
|
||||
body_is_empty,
|
||||
spec_metadata.decision_kind == GEMINI_FILES_UPLOAD_PLAN_KIND,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
Some(&parts.headers),
|
||||
Some(effective_headers),
|
||||
) {
|
||||
Ok(parts) => parts,
|
||||
Err(GeminiFilesRequestBodyError::BodyRulesUnsupportedForBinaryUpload) => {
|
||||
@@ -145,7 +146,7 @@ pub(super) async fn resolve_local_gemini_files_candidate_payload_parts(
|
||||
};
|
||||
|
||||
let Some(provider_request_headers) = build_gemini_files_headers(GeminiFilesHeadersInput {
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
auth_header: &auth_header,
|
||||
auth_value: &auth_value,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
|
||||
@@ -14,7 +14,8 @@ use crate::ai_serving::planner::candidate_metadata::{
|
||||
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
|
||||
};
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
build_local_authenticated_decision_input, resolve_local_authenticated_decision_input,
|
||||
attach_routing_policy_to_local_requested_model_input,
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
@@ -29,11 +30,12 @@ use crate::{AppState, GatewayError};
|
||||
|
||||
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttempt as LocalGeminiFilesCandidateAttempt;
|
||||
pub(super) use crate::ai_serving::planner::candidate_materialization::LocalExecutionCandidateAttemptSource as LocalGeminiFilesCandidateAttemptSource;
|
||||
pub(super) use crate::ai_serving::planner::decision_input::LocalAuthenticatedDecisionInput as LocalGeminiFilesDecisionInput;
|
||||
pub(super) use crate::ai_serving::planner::decision_input::LocalRequestedModelDecisionInput as LocalGeminiFilesDecisionInput;
|
||||
|
||||
pub(super) const GEMINI_FILES_CANDIDATE_API_FORMAT: &str = "gemini:files";
|
||||
pub(super) const GEMINI_FILES_CLIENT_API_FORMAT: &str = "gemini:files";
|
||||
pub(super) const GEMINI_FILES_REQUIRED_CAPABILITY: &str = "gemini_files";
|
||||
pub(super) const GEMINI_FILES_ROUTING_MODEL: &str = "gemini-files";
|
||||
|
||||
pub(super) async fn resolve_local_gemini_files_decision_input(
|
||||
state: &AppState,
|
||||
@@ -41,9 +43,9 @@ pub(super) async fn resolve_local_gemini_files_decision_input(
|
||||
body_json: Option<&serde_json::Value>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<LocalGeminiFilesDecisionInput> {
|
||||
) -> Result<Option<LocalGeminiFilesDecisionInput>, GatewayError> {
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
return None;
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let explicit_required_capabilities = json!({ "gemini_files": true });
|
||||
@@ -56,20 +58,33 @@ pub(super) async fn resolve_local_gemini_files_decision_input(
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Ok(None) => return Ok(None),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local gemini files decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
let mut input = build_local_authenticated_decision_input(resolved_input);
|
||||
let routing_body_json = body_json.cloned().unwrap_or(serde_json::Value::Null);
|
||||
let mut input = build_local_requested_model_decision_input(
|
||||
resolved_input,
|
||||
GEMINI_FILES_ROUTING_MODEL.to_string(),
|
||||
);
|
||||
input.request_auth_channel = decision.request_auth_channel.clone();
|
||||
input.client_session_affinity = client_session_affinity_from_parts(parts, body_json);
|
||||
Some(input)
|
||||
attach_routing_policy_to_local_requested_model_input(
|
||||
state,
|
||||
parts,
|
||||
&mut input,
|
||||
&routing_body_json,
|
||||
GEMINI_FILES_CLIENT_API_FORMAT,
|
||||
)
|
||||
.await?;
|
||||
Ok(Some(input))
|
||||
}
|
||||
|
||||
pub(super) async fn materialize_local_gemini_files_candidate_attempts(
|
||||
@@ -101,8 +116,9 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
None,
|
||||
None,
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
candidates,
|
||||
Vec::new(),
|
||||
@@ -173,8 +189,9 @@ pub(super) async fn build_local_gemini_files_candidate_attempt_source<'a>(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
None,
|
||||
None,
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
candidates,
|
||||
Vec::new(),
|
||||
|
||||
@@ -32,7 +32,7 @@ pub(super) use crate::ai_serving::LocalOpenAiImageSpec;
|
||||
pub(crate) struct LocalOpenAiImageSyncAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_json: serde_json::Value,
|
||||
body_base64: Option<&'a str>,
|
||||
trace_id: &'a str,
|
||||
input: LocalOpenAiImageDecisionInput,
|
||||
@@ -43,7 +43,7 @@ pub(crate) struct LocalOpenAiImageSyncAttemptSource<'a> {
|
||||
pub(crate) struct LocalOpenAiImageStreamAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_json: serde_json::Value,
|
||||
body_base64: Option<&'a str>,
|
||||
trace_id: &'a str,
|
||||
input: LocalOpenAiImageDecisionInput,
|
||||
@@ -152,16 +152,17 @@ pub(crate) async fn build_local_image_sync_attempt_source_for_kind<'a>(
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
let Some((candidates, candidate_count)) = build_local_openai_image_candidate_attempt_source(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
&effective_body_json,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
@@ -178,7 +179,7 @@ pub(crate) async fn build_local_image_sync_attempt_source_for_kind<'a>(
|
||||
LocalOpenAiImageSyncAttemptSource {
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_json: effective_body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
input,
|
||||
@@ -211,16 +212,17 @@ pub(crate) async fn build_local_image_stream_attempt_source_for_kind<'a>(
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
let Some((candidates, candidate_count)) = build_local_openai_image_candidate_attempt_source(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
&effective_body_json,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
@@ -237,7 +239,7 @@ pub(crate) async fn build_local_image_stream_attempt_source_for_kind<'a>(
|
||||
LocalOpenAiImageStreamAttemptSource {
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_json: effective_body_json,
|
||||
body_base64,
|
||||
trace_id,
|
||||
input,
|
||||
@@ -303,21 +305,21 @@ impl LocalOpenAiImageSyncAttemptSource<'_> {
|
||||
let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
self.body_base64,
|
||||
self.trace_id,
|
||||
&self.input,
|
||||
attempt,
|
||||
self.spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
|
||||
let built = if provider_api_format == "gemini:generate_content" {
|
||||
build_gemini_sync_plan_from_decision(self.parts, self.body_json, payload)
|
||||
build_gemini_sync_plan_from_decision(self.parts, &self.body_json, payload)
|
||||
} else {
|
||||
build_passthrough_sync_plan_from_decision(self.parts, payload)
|
||||
};
|
||||
@@ -345,23 +347,23 @@ impl LocalOpenAiImageStreamAttemptSource<'_> {
|
||||
let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
self.body_base64,
|
||||
self.trace_id,
|
||||
&self.input,
|
||||
attempt,
|
||||
self.spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let provider_api_format = payload.provider_api_format.as_deref().unwrap_or_default();
|
||||
let built = if provider_api_format == "gemini:generate_content" {
|
||||
build_gemini_stream_plan_from_decision(self.parts, self.body_json, payload)
|
||||
build_gemini_stream_plan_from_decision(self.parts, &self.body_json, payload)
|
||||
} else {
|
||||
build_standard_stream_plan_from_decision(self.parts, self.body_json, payload, false)
|
||||
build_standard_stream_plan_from_decision(self.parts, &self.body_json, payload, false)
|
||||
};
|
||||
match built {
|
||||
Ok(value) => Ok(value),
|
||||
@@ -400,10 +402,11 @@ pub(crate) async fn maybe_build_sync_local_image_decision_payload(
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
|
||||
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
|
||||
state,
|
||||
@@ -429,7 +432,7 @@ pub(crate) async fn maybe_build_sync_local_image_decision_payload(
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
@@ -460,10 +463,11 @@ pub(crate) async fn maybe_build_stream_local_image_decision_payload(
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
|
||||
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
|
||||
state,
|
||||
@@ -489,7 +493,7 @@ pub(crate) async fn maybe_build_stream_local_image_decision_payload(
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
@@ -516,10 +520,11 @@ async fn build_local_sync_plan_and_reports(
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
|
||||
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
|
||||
state,
|
||||
@@ -546,7 +551,7 @@ async fn build_local_sync_plan_and_reports(
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
@@ -592,10 +597,11 @@ async fn build_local_stream_plan_and_reports(
|
||||
trace_id,
|
||||
decision,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
|
||||
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
|
||||
state,
|
||||
@@ -622,7 +628,7 @@ async fn build_local_stream_plan_and_reports(
|
||||
attempt,
|
||||
spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, LocalExecutionReportContextParts,
|
||||
};
|
||||
@@ -10,7 +11,9 @@ use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_profile,
|
||||
};
|
||||
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
|
||||
use crate::{append_execution_contract_fields_to_value, AiExecutionDecision, AppState};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, AiExecutionDecision, AppState, GatewayError,
|
||||
};
|
||||
|
||||
use super::request::resolve_local_openai_image_candidate_payload_parts;
|
||||
use super::support::{LocalOpenAiImageCandidateAttempt, LocalOpenAiImageDecisionInput};
|
||||
@@ -25,11 +28,11 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
input: &LocalOpenAiImageDecisionInput,
|
||||
attempt: LocalOpenAiImageCandidateAttempt,
|
||||
spec: LocalOpenAiImageSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let spec_metadata = local_openai_image_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let attempt_identity = attempt.attempt_identity();
|
||||
let resolved = resolve_local_openai_image_candidate_payload_parts(
|
||||
let Some(resolved) = resolve_local_openai_image_candidate_payload_parts(
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
@@ -39,7 +42,10 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
&attempt,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let LocalOpenAiImageCandidateAttempt {
|
||||
eligible,
|
||||
candidate_id,
|
||||
@@ -88,6 +94,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
.get("stream")
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(spec_metadata.require_streaming);
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let report_context = append_execution_contract_fields_to_value(
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &input.auth_context,
|
||||
@@ -114,7 +121,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
body_rules: transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: Some(serde_json::Value::String(parts.method.to_string())),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
original_headers: effective_headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
@@ -134,39 +141,39 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
|
||||
provider_api_format.as_str(),
|
||||
);
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url: resolved.upstream_url,
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(resolved.auth_header),
|
||||
auth_value: Some(resolved.auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: resolved.requested_model,
|
||||
mapped_model: resolved.mapped_model,
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers: resolved.provider_request_headers,
|
||||
provider_request_body: Some(resolved.provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url: resolved.upstream_url,
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(resolved.auth_header),
|
||||
auth_value: Some(resolved.auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: resolved.requested_model,
|
||||
mapped_model: resolved.mapped_model,
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers: resolved.provider_request_headers,
|
||||
provider_request_body: Some(resolved.provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
});
|
||||
apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
|
||||
Ok(Some(decision))
|
||||
}
|
||||
|
||||
@@ -59,6 +59,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
let provider_api_format = attempt.eligible.provider_api_format.as_str();
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
|
||||
if provider_api_format == "gemini:generate_content" {
|
||||
return resolve_local_openai_image_to_gemini_candidate_payload_parts(
|
||||
@@ -170,7 +171,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
|
||||
let Some(mut provider_request_headers) =
|
||||
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
auth_header: &auth_header,
|
||||
auth_value: &auth_value,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
@@ -201,7 +202,7 @@ pub(super) async fn resolve_local_openai_image_candidate_payload_parts(
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
spec_metadata.api_format,
|
||||
Some(trace_id),
|
||||
@@ -256,6 +257,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
let provider_api_format = "gemini:generate_content";
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
|
||||
let prepared_candidate = match prepare_header_authenticated_candidate(
|
||||
PlannerAppState::new(state),
|
||||
@@ -332,7 +334,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
|
||||
converted.body_json,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
body_json,
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
) {
|
||||
Some(body) => body,
|
||||
None => {
|
||||
@@ -384,7 +386,7 @@ async fn resolve_local_openai_image_to_gemini_candidate_payload_parts(
|
||||
transport,
|
||||
provider_api_format,
|
||||
same_format: false,
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
auth_header: &prepared_candidate.auth_header,
|
||||
auth_value: &prepared_candidate.auth_value,
|
||||
extra_headers: &BTreeMap::new(),
|
||||
|
||||
@@ -13,6 +13,7 @@ use crate::ai_serving::planner::candidate_metadata::{
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
attach_routing_policy_to_local_requested_model_input,
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
@@ -42,12 +43,16 @@ pub(super) async fn resolve_local_openai_image_decision_input(
|
||||
body_base64: Option<&str>,
|
||||
trace_id: &str,
|
||||
decision: &GatewayControlDecision,
|
||||
) -> Option<LocalOpenAiImageDecisionInput> {
|
||||
) -> Result<Option<LocalOpenAiImageDecisionInput>, GatewayError> {
|
||||
let Some(auth_context) = resolve_local_openai_image_auth_context(decision) else {
|
||||
return None;
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let requested_model = resolve_requested_image_model_for_request(parts, body_json, body_base64)?;
|
||||
let Some(requested_model) =
|
||||
resolve_requested_image_model_for_request(parts, body_json, body_base64)
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
@@ -58,21 +63,37 @@ pub(super) async fn resolve_local_openai_image_decision_input(
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Ok(None) => return Ok(None),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai image decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
|
||||
input.request_auth_channel = decision.request_auth_channel.clone();
|
||||
input.client_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
|
||||
Some(input)
|
||||
if let Err(err) = attach_routing_policy_to_local_requested_model_input(
|
||||
state,
|
||||
parts,
|
||||
&mut input,
|
||||
body_json,
|
||||
"openai:image",
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai image decision routing profile resolution failed"
|
||||
);
|
||||
return Err(err);
|
||||
}
|
||||
Ok(Some(input))
|
||||
}
|
||||
|
||||
fn resolve_local_openai_image_auth_context(
|
||||
@@ -229,6 +250,7 @@ pub(super) async fn build_local_openai_image_candidate_attempt_source<'a>(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
@@ -305,6 +327,7 @@ async fn materialize_local_openai_image_candidate_attempts(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
|
||||
@@ -27,7 +27,7 @@ use self::support::{
|
||||
pub(crate) struct LocalVideoCreateSyncAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_json: serde_json::Value,
|
||||
trace_id: &'a str,
|
||||
input: LocalVideoCreateDecisionInput,
|
||||
spec: LocalVideoCreateSpec,
|
||||
@@ -65,16 +65,17 @@ pub(crate) async fn build_local_video_sync_attempt_source_for_kind<'a>(
|
||||
let Some(input) = resolve_local_video_create_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
let Some((candidates, candidate_count)) = build_local_video_create_candidate_attempt_source(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
body_json,
|
||||
&effective_body_json,
|
||||
spec_metadata.api_format,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
@@ -91,7 +92,7 @@ pub(crate) async fn build_local_video_sync_attempt_source_for_kind<'a>(
|
||||
LocalVideoCreateSyncAttemptSource {
|
||||
state,
|
||||
parts,
|
||||
body_json,
|
||||
body_json: effective_body_json,
|
||||
trace_id,
|
||||
input,
|
||||
spec,
|
||||
@@ -133,13 +134,13 @@ impl LocalVideoCreateSyncAttemptSource<'_> {
|
||||
let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
|
||||
self.state,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
self.trace_id,
|
||||
&self.input,
|
||||
attempt,
|
||||
self.spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -175,10 +176,11 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
|
||||
let Some(input) = resolve_local_video_create_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
|
||||
let Some((mut source, _)) = build_local_video_create_candidate_attempt_source(
|
||||
state,
|
||||
@@ -197,7 +199,7 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
|
||||
if let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
|
||||
state, parts, body_json, trace_id, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
@@ -218,10 +220,11 @@ async fn build_local_sync_plan_and_reports(
|
||||
let Some(input) = resolve_local_video_create_decision_input(
|
||||
state, parts, trace_id, decision, body_json, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
|
||||
let Some((mut source, _)) = build_local_video_create_candidate_attempt_source(
|
||||
state,
|
||||
@@ -241,7 +244,7 @@ async fn build_local_sync_plan_and_reports(
|
||||
let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
|
||||
state, parts, body_json, trace_id, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, LocalExecutionReportContextParts,
|
||||
};
|
||||
@@ -10,7 +11,7 @@ use crate::ai_serving::transport::{
|
||||
resolve_transport_execution_timeouts, resolve_transport_profile,
|
||||
};
|
||||
use crate::ai_serving::{ai_local_execution_contract_for_formats, PlannerAppState};
|
||||
use crate::{AiExecutionDecision, AppState};
|
||||
use crate::{AiExecutionDecision, AppState, GatewayError};
|
||||
|
||||
use super::request::resolve_local_video_create_candidate_payload_parts;
|
||||
use super::support::{LocalVideoCreateCandidateAttempt, LocalVideoCreateDecisionInput};
|
||||
@@ -24,14 +25,17 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
||||
input: &LocalVideoCreateDecisionInput,
|
||||
attempt: LocalVideoCreateCandidateAttempt,
|
||||
spec: LocalVideoCreateSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let spec_metadata = local_video_create_spec_metadata(spec);
|
||||
let planner_state = PlannerAppState::new(state);
|
||||
let attempt_identity = attempt.attempt_identity();
|
||||
let resolved = resolve_local_video_create_candidate_payload_parts(
|
||||
let Some(resolved) = resolve_local_video_create_candidate_payload_parts(
|
||||
state, parts, body_json, trace_id, input, &attempt, spec,
|
||||
)
|
||||
.await?;
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let LocalVideoCreateCandidateAttempt {
|
||||
eligible,
|
||||
candidate_id,
|
||||
@@ -50,6 +54,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
||||
if let Some(proxy_value) = build_request_trace_proxy_value(Some(&transport), proxy.as_ref()) {
|
||||
extra_fields.insert("proxy".to_string(), proxy_value);
|
||||
}
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let report_context = build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
auth_context: &input.auth_context,
|
||||
request_id: trace_id,
|
||||
@@ -75,7 +80,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
||||
body_rules: transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: None,
|
||||
provider_request_headers: None,
|
||||
original_headers: &parts.headers,
|
||||
original_headers: effective_headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
@@ -99,45 +104,45 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
||||
upstream_url,
|
||||
} = resolved;
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: false,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: spec_metadata.api_format.to_string(),
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: parts
|
||||
.headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream: false,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: false,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: Some(parts.method.to_string()),
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format: spec_metadata.api_format.to_string(),
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: parts
|
||||
.headers
|
||||
.get(http::header::CONTENT_TYPE)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts: resolve_transport_execution_timeouts(&transport),
|
||||
upstream_is_stream: false,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
});
|
||||
apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
|
||||
Ok(Some(decision))
|
||||
}
|
||||
|
||||
@@ -41,6 +41,7 @@ pub(super) async fn resolve_local_video_create_candidate_payload_parts(
|
||||
let spec_metadata = local_video_create_spec_metadata(spec);
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
|
||||
let provider_family = provider_video_create_family(spec.family);
|
||||
let transport_unsupported_reason = video_create_transport_unsupported_reason(
|
||||
@@ -124,7 +125,7 @@ pub(super) async fn resolve_local_video_create_candidate_payload_parts(
|
||||
provider_family,
|
||||
&mapped_model,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
Some(&parts.headers),
|
||||
Some(effective_headers),
|
||||
) else {
|
||||
mark_skipped_local_video_candidate_with_failure_diagnostic(
|
||||
state,
|
||||
@@ -146,7 +147,7 @@ pub(super) async fn resolve_local_video_create_candidate_payload_parts(
|
||||
|
||||
let Some(provider_request_headers) =
|
||||
build_video_create_headers(ProviderVideoCreateHeadersInput {
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
auth_header: &auth_header,
|
||||
auth_value: &auth_value,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
|
||||
@@ -15,6 +15,7 @@ use crate::ai_serving::planner::candidate_metadata::{
|
||||
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
attach_routing_policy_to_local_requested_model_input,
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
@@ -41,19 +42,21 @@ pub(super) async fn resolve_local_video_create_decision_input(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalVideoCreateSpec,
|
||||
) -> Option<LocalVideoCreateDecisionInput> {
|
||||
) -> Result<Option<LocalVideoCreateDecisionInput>, GatewayError> {
|
||||
let spec_metadata = local_video_create_spec_metadata(spec);
|
||||
let Some(auth_context) = resolve_local_video_create_auth_context(decision, spec.family) else {
|
||||
return None;
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let requested_model = extract_requested_model_from_request(
|
||||
let Some(requested_model) = extract_requested_model_from_request(
|
||||
parts,
|
||||
body_json,
|
||||
spec_metadata
|
||||
.requested_model_family
|
||||
.expect("video specs should declare requested-model family"),
|
||||
)?;
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
@@ -64,7 +67,7 @@ pub(super) async fn resolve_local_video_create_decision_input(
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Ok(None) => return Ok(None),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
@@ -72,14 +75,31 @@ pub(super) async fn resolve_local_video_create_decision_input(
|
||||
error = ?err,
|
||||
"gateway local video decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
|
||||
input.request_auth_channel = decision.request_auth_channel.clone();
|
||||
input.client_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
|
||||
Some(input)
|
||||
if let Err(err) = attach_routing_policy_to_local_requested_model_input(
|
||||
state,
|
||||
parts,
|
||||
&mut input,
|
||||
body_json,
|
||||
spec_metadata.api_format,
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
decision_kind = spec_metadata.decision_kind,
|
||||
error = ?err,
|
||||
"gateway local video decision routing profile resolution failed"
|
||||
);
|
||||
return Err(err);
|
||||
}
|
||||
Ok(Some(input))
|
||||
}
|
||||
|
||||
fn resolve_local_video_create_auth_context(
|
||||
@@ -196,6 +216,7 @@ pub(super) async fn build_local_video_create_candidate_attempt_source<'a>(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
@@ -261,6 +282,7 @@ async fn materialize_local_video_create_candidate_attempts(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
|
||||
@@ -29,7 +29,7 @@ pub(crate) struct LocalStandardSyncAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
trace_id: &'a str,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_json: serde_json::Value,
|
||||
input: LocalStandardDecisionInput,
|
||||
spec: LocalStandardSpec,
|
||||
requested_model_family: RequestedModelFamily,
|
||||
@@ -40,7 +40,7 @@ pub(crate) struct LocalStandardStreamAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
trace_id: &'a str,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_json: serde_json::Value,
|
||||
input: LocalStandardDecisionInput,
|
||||
spec: LocalStandardSpec,
|
||||
requested_model_family: RequestedModelFamily,
|
||||
@@ -61,7 +61,7 @@ pub(crate) async fn build_local_sync_attempt_source<'a>(
|
||||
.expect("standard spec metadata should include requested-model family");
|
||||
let Some(input) =
|
||||
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
@@ -82,9 +82,15 @@ pub(crate) async fn build_local_sync_attempt_source<'a>(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (candidates, candidate_count) =
|
||||
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
let (candidates, candidate_count) = build_local_standard_candidate_attempt_source(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
&effective_body_json,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
if candidate_count == 0 {
|
||||
return Ok(None);
|
||||
@@ -95,7 +101,7 @@ pub(crate) async fn build_local_sync_attempt_source<'a>(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
body_json: effective_body_json,
|
||||
input,
|
||||
spec,
|
||||
requested_model_family,
|
||||
@@ -119,7 +125,7 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
|
||||
.expect("standard spec metadata should include requested-model family");
|
||||
let Some(input) =
|
||||
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
@@ -140,9 +146,15 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let (candidates, candidate_count) =
|
||||
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
let (candidates, candidate_count) = build_local_standard_candidate_attempt_source(
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
&effective_body_json,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
if candidate_count == 0 {
|
||||
return Ok(None);
|
||||
@@ -153,7 +165,7 @@ pub(crate) async fn build_local_stream_attempt_source<'a>(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
body_json: effective_body_json,
|
||||
input,
|
||||
spec,
|
||||
requested_model_family,
|
||||
@@ -228,19 +240,19 @@ impl LocalStandardSyncAttemptSource<'_> {
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
&self.input,
|
||||
attempt,
|
||||
self.spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
match build_sync_plan_from_requested_model_family(
|
||||
self.requested_model_family,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
payload,
|
||||
) {
|
||||
Ok(value) => Ok(value),
|
||||
@@ -265,19 +277,19 @@ impl LocalStandardStreamAttemptSource<'_> {
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
&self.input,
|
||||
attempt,
|
||||
self.spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
match build_stream_plan_from_requested_model_family(
|
||||
self.requested_model_family,
|
||||
self.parts,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
payload,
|
||||
) {
|
||||
Ok(value) => Ok(value),
|
||||
@@ -309,7 +321,7 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
|
||||
|
||||
let Some(input) =
|
||||
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -322,6 +334,7 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
let (mut source, candidate_count) =
|
||||
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
@@ -331,7 +344,7 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
|
||||
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
@@ -358,7 +371,7 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
|
||||
|
||||
let Some(input) =
|
||||
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -371,6 +384,7 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
let (mut source, candidate_count) =
|
||||
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
@@ -380,7 +394,7 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
|
||||
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
@@ -405,7 +419,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
.expect("standard spec metadata should include requested-model family");
|
||||
let Some(input) =
|
||||
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
@@ -426,6 +440,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
let (mut source, candidate_count) =
|
||||
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
@@ -438,7 +453,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
|
||||
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
@@ -479,7 +494,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
.expect("standard spec metadata should include requested-model family");
|
||||
let Some(input) =
|
||||
resolve_local_standard_decision_input(state, parts, trace_id, decision, body_json, spec)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
set_local_runtime_miss_diagnostic_reason(
|
||||
state,
|
||||
@@ -500,6 +515,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
let (mut source, candidate_count) =
|
||||
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
|
||||
.await?;
|
||||
@@ -512,7 +528,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
|
||||
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
@@ -16,6 +16,7 @@ use crate::ai_serving::planner::candidate_source::{
|
||||
};
|
||||
use crate::ai_serving::planner::common::extract_requested_model_from_request;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
attach_routing_policy_to_local_requested_model_input,
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
@@ -39,19 +40,21 @@ pub(super) async fn resolve_local_standard_decision_input(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Option<LocalStandardDecisionInput> {
|
||||
) -> Result<Option<LocalStandardDecisionInput>, GatewayError> {
|
||||
let spec_metadata = local_standard_spec_metadata(spec);
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
return None;
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let requested_model = extract_requested_model_from_request(
|
||||
let Some(requested_model) = extract_requested_model_from_request(
|
||||
parts,
|
||||
body_json,
|
||||
spec_metadata
|
||||
.requested_model_family
|
||||
.expect("standard specs should declare requested-model family"),
|
||||
)?;
|
||||
) else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
state,
|
||||
@@ -62,7 +65,7 @@ pub(super) async fn resolve_local_standard_decision_input(
|
||||
.await
|
||||
{
|
||||
Ok(Some(resolved_input)) => resolved_input,
|
||||
Ok(None) => return None,
|
||||
Ok(None) => return Ok(None),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
@@ -70,14 +73,31 @@ pub(super) async fn resolve_local_standard_decision_input(
|
||||
error = ?err,
|
||||
"gateway local standard decision auth snapshot read failed"
|
||||
);
|
||||
return None;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
|
||||
input.request_auth_channel = decision.request_auth_channel.clone();
|
||||
input.client_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
|
||||
Some(input)
|
||||
if let Err(err) = attach_routing_policy_to_local_requested_model_input(
|
||||
state,
|
||||
parts,
|
||||
&mut input,
|
||||
body_json,
|
||||
spec_metadata.api_format,
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
api_format = spec_metadata.api_format,
|
||||
error = ?err,
|
||||
"gateway local standard decision routing profile resolution failed"
|
||||
);
|
||||
return Err(err);
|
||||
}
|
||||
Ok(Some(input))
|
||||
}
|
||||
|
||||
pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
@@ -104,6 +124,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
spec_metadata.require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
&input.auth_snapshot,
|
||||
input.routing_policy.as_ref(),
|
||||
input.client_session_affinity.as_ref(),
|
||||
false,
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
|
||||
@@ -128,6 +149,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
@@ -228,6 +250,7 @@ pub(super) async fn build_local_standard_candidate_attempt_source<'a>(
|
||||
&input.auth_snapshot,
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
@@ -320,6 +343,7 @@ async fn maybe_append_gemini_image_openai_image_preselection(
|
||||
spec_metadata.require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
&input.auth_snapshot,
|
||||
input.routing_policy.as_ref(),
|
||||
input.client_session_affinity.as_ref(),
|
||||
false,
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
|
||||
|
||||
@@ -3,6 +3,7 @@ use crate::ai_serving::planner::candidate_materialization::{
|
||||
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
|
||||
mark_skipped_local_execution_candidate_with_failure_diagnostic,
|
||||
};
|
||||
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
|
||||
};
|
||||
@@ -24,7 +25,7 @@ use crate::ai_serving::{
|
||||
};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
|
||||
AiExecutionDecision, AppState,
|
||||
AiExecutionDecision, AppState, GatewayError,
|
||||
};
|
||||
|
||||
use super::request::resolve_local_standard_candidate_payload_parts;
|
||||
@@ -38,7 +39,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
input: &LocalStandardDecisionInput,
|
||||
attempt: LocalStandardCandidateAttempt,
|
||||
spec: LocalStandardSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let spec_metadata = local_standard_spec_metadata(spec);
|
||||
if api_format_alias_matches(
|
||||
&attempt.eligible.provider_api_format,
|
||||
@@ -70,10 +71,13 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
..
|
||||
} = &attempt;
|
||||
let candidate = &eligible.candidate;
|
||||
let resolved = resolve_local_standard_candidate_payload_parts(
|
||||
let Some(resolved) = resolve_local_standard_candidate_payload_parts(
|
||||
state, parts, trace_id, body_json, input, &attempt, spec,
|
||||
)
|
||||
.await?;
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let proxy = state
|
||||
.resolve_transport_proxy_snapshot_with_tunnel_affinity(&resolved.transport)
|
||||
.await;
|
||||
@@ -93,6 +97,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
spec_metadata.api_format,
|
||||
resolved.provider_api_format.as_str(),
|
||||
);
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let report_context = append_local_failover_policy_to_value(
|
||||
append_execution_contract_fields_to_value(
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
@@ -120,7 +125,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
body_rules: resolved.transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
original_headers: effective_headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
@@ -159,41 +164,41 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
||||
transport,
|
||||
} = resolved;
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.to_string(),
|
||||
provider_name: candidate.provider_name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts,
|
||||
upstream_is_stream,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.to_string(),
|
||||
provider_name: candidate.provider_name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key: None,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts,
|
||||
upstream_is_stream,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
});
|
||||
apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
|
||||
Ok(Some(decision))
|
||||
}
|
||||
|
||||
pub(super) async fn mark_skipped_local_standard_candidate(
|
||||
@@ -347,6 +352,9 @@ mod tests {
|
||||
required_capabilities: None,
|
||||
request_auth_channel: None,
|
||||
client_session_affinity: None,
|
||||
routing_policy: None,
|
||||
routing_trace_seed: None,
|
||||
routing_context: None,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -512,6 +520,7 @@ mod tests {
|
||||
claude_stream_spec(),
|
||||
)
|
||||
.await
|
||||
.expect("same-format candidate should not fail routing mutation")
|
||||
.expect("same-format candidate should build a standard-family payload");
|
||||
|
||||
assert_eq!(payload.endpoint_id.as_deref(), Some("endpoint-claude"));
|
||||
@@ -547,6 +556,7 @@ mod tests {
|
||||
claude_stream_spec(),
|
||||
)
|
||||
.await
|
||||
.expect("cross-format candidate should not fail routing mutation")
|
||||
.expect("cross-format candidate should still build after the same-format candidate");
|
||||
|
||||
assert_eq!(
|
||||
|
||||
@@ -65,6 +65,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let transport = &attempt.eligible.transport;
|
||||
let provider_api_format = attempt.eligible.provider_api_format.as_str();
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
if spec_metadata.api_format == "gemini:generate_content"
|
||||
&& provider_api_format == "openai:image"
|
||||
&& gemini_request_is_image_generation(body_json)
|
||||
@@ -209,7 +210,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
|
||||
transport.endpoint.body_rules.as_ref()
|
||||
},
|
||||
Some(input.auth_context.api_key_id.as_str()),
|
||||
Some(&parts.headers),
|
||||
Some(effective_headers),
|
||||
enable_model_directives,
|
||||
) {
|
||||
Some(body) => body,
|
||||
@@ -312,7 +313,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
|
||||
transport,
|
||||
provider_api_format,
|
||||
same_format: false,
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
auth_header: &prepared_candidate.auth_header,
|
||||
auth_value: &prepared_candidate.auth_value,
|
||||
extra_headers: &BTreeMap::new(),
|
||||
@@ -343,7 +344,7 @@ pub(crate) async fn resolve_local_standard_candidate_payload_parts(
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format,
|
||||
Some(trace_id),
|
||||
@@ -448,9 +449,10 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
|
||||
|
||||
let upstream_is_stream = true;
|
||||
let upstream_url = build_openai_image_upstream_url(transport, None);
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let Some(mut provider_request_headers) =
|
||||
build_openai_image_headers(ProviderOpenAiImageHeadersInput {
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
auth_header: &prepared_candidate.auth_header,
|
||||
auth_value: &prepared_candidate.auth_value,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
@@ -478,7 +480,7 @@ async fn resolve_local_gemini_image_to_openai_image_candidate_payload_parts(
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&converted.body_json,
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format,
|
||||
Some(trace_id),
|
||||
@@ -517,12 +519,13 @@ async fn build_kiro_cross_format_payload_parts(
|
||||
kiro_auth: &KiroRequestAuth,
|
||||
) -> Option<LocalStandardCandidatePayloadParts> {
|
||||
let candidate = &attempt.eligible.candidate;
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let provider_request_body = match build_kiro_provider_request_body(
|
||||
&claude_request_body,
|
||||
&mapped_model,
|
||||
&kiro_auth.auth_config,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
Some(&parts.headers),
|
||||
Some(effective_headers),
|
||||
) {
|
||||
Some(body) => body,
|
||||
None => {
|
||||
@@ -573,7 +576,7 @@ async fn build_kiro_cross_format_payload_parts(
|
||||
}
|
||||
};
|
||||
let provider_request_headers = match build_kiro_provider_headers(KiroProviderHeadersInput {
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: original_body_json,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
|
||||
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, insert_provider_stream_event_api_format,
|
||||
LocalExecutionReportContextParts,
|
||||
@@ -105,6 +106,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
} else {
|
||||
Some(body_json)
|
||||
};
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let report_context = append_local_failover_policy_to_value(
|
||||
append_execution_contract_fields_to_value(
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
@@ -132,7 +134,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
body_rules: transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
original_headers: effective_headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
@@ -160,39 +162,39 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
||||
&transport,
|
||||
);
|
||||
|
||||
Ok(Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream,
|
||||
decision_kind: decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: "openai:chat".to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts,
|
||||
upstream_is_stream,
|
||||
report_kind: Some(report_kind),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
)))
|
||||
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
|
||||
decision_is_stream,
|
||||
decision_kind: decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: "openai:chat".to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts,
|
||||
upstream_is_stream,
|
||||
report_kind: Some(report_kind),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
});
|
||||
apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
|
||||
Ok(Some(decision))
|
||||
}
|
||||
|
||||
@@ -190,6 +190,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
resolve_provider_chat_request_redaction(state, parts, body_json, input, candidate_id)
|
||||
.await?;
|
||||
let body_json = redaction.body_json.as_ref();
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
|
||||
if provider_api_format == "openai:chat" {
|
||||
if let Some(skip_reason) = local_openai_chat_transport_unsupported_reason(transport) {
|
||||
@@ -241,7 +242,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
upstream_is_stream,
|
||||
force_body_stream_field,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
enable_model_directives,
|
||||
) else {
|
||||
mark_skipped_local_openai_chat_candidate_with_extra_data(
|
||||
@@ -286,7 +287,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
transport,
|
||||
provider_api_format,
|
||||
same_format: true,
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
auth_header: &prepared_candidate.auth_header,
|
||||
auth_value: &prepared_candidate.auth_value,
|
||||
extra_headers: &BTreeMap::new(),
|
||||
@@ -317,7 +318,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
transport.endpoint.api_format.as_str(),
|
||||
Some(trace_id),
|
||||
@@ -480,7 +481,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
transport.endpoint.body_rules.as_ref()
|
||||
},
|
||||
Some(input.auth_context.api_key_id.as_str()),
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
enable_model_directives,
|
||||
) else {
|
||||
mark_skipped_local_openai_chat_candidate_with_extra_data(
|
||||
@@ -575,7 +576,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
transport,
|
||||
provider_api_format: provider_api_format.as_str(),
|
||||
same_format: false,
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
auth_header: &prepared_candidate.auth_header,
|
||||
auth_value: &prepared_candidate.auth_value,
|
||||
extra_headers: &BTreeMap::new(),
|
||||
@@ -606,7 +607,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format.as_str(),
|
||||
Some(trace_id),
|
||||
@@ -664,12 +665,13 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
|
||||
request_redacted: bool,
|
||||
) -> Option<LocalOpenAiChatCandidatePayloadParts> {
|
||||
let candidate = &eligible.candidate;
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let provider_request_body = match build_kiro_provider_request_body(
|
||||
&claude_request_body,
|
||||
&mapped_model,
|
||||
&kiro_auth.auth_config,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
Some(&parts.headers),
|
||||
Some(effective_headers),
|
||||
) {
|
||||
Some(body) => body,
|
||||
None => {
|
||||
@@ -720,7 +722,7 @@ async fn build_kiro_openai_chat_cross_format_payload_parts(
|
||||
}
|
||||
};
|
||||
let mut provider_request_headers = match build_kiro_provider_headers(KiroProviderHeadersInput {
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: original_body_json,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
|
||||
@@ -140,6 +140,7 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
@@ -220,6 +221,7 @@ pub(crate) async fn build_local_openai_chat_candidate_attempt_source<'a>(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
@@ -298,6 +300,7 @@ pub(crate) async fn build_lazy_local_openai_chat_candidate_attempt_source<'a>(
|
||||
&input.auth_snapshot,
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
|
||||
@@ -136,10 +136,11 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, parts, trace_id, decision, body_json, plan_kind, false,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
|
||||
let (mut source, _) = build_lazy_local_openai_chat_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, false,
|
||||
@@ -187,10 +188,11 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, parts, trace_id, decision, body_json, plan_kind, false,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
|
||||
let (mut source, _) = build_lazy_local_openai_chat_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, true,
|
||||
|
||||
@@ -26,6 +26,7 @@ pub(crate) async fn list_local_openai_chat_candidates(
|
||||
require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
&input.auth_snapshot,
|
||||
input.routing_policy.as_ref(),
|
||||
input.client_session_affinity.as_ref(),
|
||||
false,
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel,
|
||||
|
||||
@@ -4,11 +4,12 @@ use super::super::{GatewayControlDecision, LocalOpenAiChatDecisionInput};
|
||||
use super::diagnostic::set_local_openai_chat_miss_diagnostic;
|
||||
use crate::ai_serving::planner::common::extract_standard_requested_model;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
attach_routing_policy_to_local_requested_model_input,
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::resolve_local_decision_execution_runtime_auth_context;
|
||||
use crate::client_session_affinity::client_session_affinity_from_parts;
|
||||
use crate::AppState;
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
pub(crate) async fn resolve_local_openai_chat_decision_input(
|
||||
state: &AppState,
|
||||
@@ -18,7 +19,7 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
record_miss_diagnostic: bool,
|
||||
) -> Option<LocalOpenAiChatDecisionInput> {
|
||||
) -> Result<Option<LocalOpenAiChatDecisionInput>, GatewayError> {
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
@@ -37,7 +38,7 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
|
||||
"missing_auth_context",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(requested_model) = extract_standard_requested_model(body_json) else {
|
||||
@@ -55,7 +56,7 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
|
||||
"missing_requested_model",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
@@ -84,7 +85,7 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
|
||||
"auth_snapshot_missing",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
return Ok(None);
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
@@ -102,12 +103,28 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
|
||||
"auth_snapshot_read_failed",
|
||||
);
|
||||
}
|
||||
return None;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
|
||||
input.request_auth_channel = decision.request_auth_channel.clone();
|
||||
input.client_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
|
||||
Some(input)
|
||||
if let Err(err) = attach_routing_policy_to_local_requested_model_input(
|
||||
state,
|
||||
parts,
|
||||
&mut input,
|
||||
body_json,
|
||||
"openai:chat",
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai chat decision routing profile resolution failed"
|
||||
);
|
||||
return Err(err);
|
||||
}
|
||||
Ok(Some(input))
|
||||
}
|
||||
|
||||
@@ -23,7 +23,7 @@ pub(crate) struct LocalOpenAiChatStreamAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
trace_id: &'a str,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_json: serde_json::Value,
|
||||
input: LocalOpenAiChatDecisionInput,
|
||||
candidates: LocalOpenAiChatCandidateAttemptSource<'a>,
|
||||
}
|
||||
@@ -43,13 +43,18 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, parts, trace_id, decision, body_json, plan_kind, true,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
|
||||
let (candidates, candidate_count) = build_lazy_local_openai_chat_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, true,
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
&effective_body_json,
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
if candidate_count == 0 {
|
||||
@@ -77,7 +82,7 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
body_json: effective_body_json,
|
||||
input,
|
||||
candidates,
|
||||
},
|
||||
@@ -127,7 +132,7 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
&self.input,
|
||||
attempt,
|
||||
OPENAI_CHAT_STREAM_PLAN_KIND,
|
||||
@@ -139,7 +144,7 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
match build_openai_chat_stream_plan_from_decision(self.parts, self.body_json, payload) {
|
||||
match build_openai_chat_stream_plan_from_decision(self.parts, &self.body_json, payload) {
|
||||
Ok(value) => Ok(value),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
@@ -168,7 +173,7 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, parts, trace_id, decision, body_json, plan_kind, true,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
@@ -23,7 +23,7 @@ pub(crate) struct LocalOpenAiChatSyncAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
trace_id: &'a str,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_json: serde_json::Value,
|
||||
input: LocalOpenAiChatDecisionInput,
|
||||
candidates: LocalOpenAiChatCandidateAttemptSource<'a>,
|
||||
}
|
||||
@@ -43,13 +43,18 @@ pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, parts, trace_id, decision, body_json, plan_kind, true,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
|
||||
let (candidates, candidate_count) = build_lazy_local_openai_chat_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, false,
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
&effective_body_json,
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
if candidate_count == 0 {
|
||||
@@ -77,7 +82,7 @@ pub(crate) async fn build_local_openai_chat_sync_attempt_source<'a>(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
body_json: effective_body_json,
|
||||
input,
|
||||
candidates,
|
||||
},
|
||||
@@ -127,7 +132,7 @@ impl LocalOpenAiChatSyncAttemptSource<'_> {
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
&self.input,
|
||||
attempt,
|
||||
OPENAI_CHAT_SYNC_PLAN_KIND,
|
||||
@@ -139,7 +144,7 @@ impl LocalOpenAiChatSyncAttemptSource<'_> {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
match build_openai_chat_sync_plan_from_decision(self.parts, self.body_json, payload) {
|
||||
match build_openai_chat_sync_plan_from_decision(self.parts, &self.body_json, payload) {
|
||||
Ok(value) => Ok(value),
|
||||
Err(err) => {
|
||||
warn!(
|
||||
@@ -168,7 +173,7 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
|
||||
let Some(input) = resolve_local_openai_chat_decision_input(
|
||||
state, parts, trace_id, decision, body_json, plan_kind, true,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
@@ -2,6 +2,7 @@ use serde_json::json;
|
||||
use tracing::debug;
|
||||
|
||||
use crate::ai_serving::build_request_trace_proxy_value;
|
||||
use crate::ai_serving::planner::decision_input::apply_provider_request_routing_policy_to_decision;
|
||||
use crate::ai_serving::planner::report_context::{
|
||||
build_local_execution_report_context, insert_provider_stream_event_api_format,
|
||||
LocalExecutionReportContextParts,
|
||||
@@ -15,7 +16,7 @@ use crate::ai_serving::transport::{
|
||||
};
|
||||
use crate::{
|
||||
append_execution_contract_fields_to_value, append_local_failover_policy_to_value,
|
||||
AiExecutionDecision, AppState,
|
||||
AiExecutionDecision, AppState, GatewayError,
|
||||
};
|
||||
|
||||
use super::request::resolve_local_openai_responses_candidate_payload_parts;
|
||||
@@ -30,7 +31,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
|
||||
input: &LocalOpenAiResponsesDecisionInput,
|
||||
attempt: LocalOpenAiResponsesCandidateAttempt,
|
||||
spec: LocalOpenAiResponsesSpec,
|
||||
) -> Option<AiExecutionDecision> {
|
||||
) -> Result<Option<AiExecutionDecision>, GatewayError> {
|
||||
let spec_metadata = local_openai_responses_spec_metadata(spec);
|
||||
let attempt_identity = attempt.attempt_identity();
|
||||
let LocalOpenAiResponsesCandidateAttempt {
|
||||
@@ -39,7 +40,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
|
||||
candidate_id,
|
||||
..
|
||||
} = attempt;
|
||||
let resolved = resolve_local_openai_responses_candidate_payload_parts(
|
||||
let Some(resolved) = resolve_local_openai_responses_candidate_payload_parts(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
@@ -50,7 +51,10 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
|
||||
&candidate_id,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
.await
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let candidate = &eligible.candidate;
|
||||
|
||||
let prompt_cache_key = resolved
|
||||
@@ -78,6 +82,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
|
||||
&mut extra_fields,
|
||||
resolved.transport.provider.provider_type.as_str(),
|
||||
);
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let report_context = append_local_failover_policy_to_value(
|
||||
append_execution_contract_fields_to_value(
|
||||
build_local_execution_report_context(LocalExecutionReportContextParts {
|
||||
@@ -105,7 +110,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
|
||||
body_rules: resolved.transport.endpoint.body_rules.as_ref(),
|
||||
provider_request_method: Some(serde_json::Value::Null),
|
||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||
original_headers: &parts.headers,
|
||||
original_headers: effective_headers,
|
||||
request_path: Some(parts.uri.path()),
|
||||
request_query_string: parts.uri.query(),
|
||||
request_origin: Some(crate::ai_serving::request_origin_from_parts(parts)),
|
||||
@@ -172,39 +177,39 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
|
||||
transport,
|
||||
} = resolved;
|
||||
|
||||
Some(build_ai_execution_decision_response(
|
||||
AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts,
|
||||
upstream_is_stream,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
},
|
||||
))
|
||||
let mut decision = build_ai_execution_decision_response(AiExecutionDecisionResponseParts {
|
||||
decision_is_stream: spec_metadata.require_streaming,
|
||||
decision_kind: spec_metadata.decision_kind.to_string(),
|
||||
execution_strategy,
|
||||
conversion_mode,
|
||||
request_id: trace_id.to_string(),
|
||||
candidate_id: candidate_id.clone(),
|
||||
provider_name: transport.provider.name.clone(),
|
||||
provider_id: candidate.provider_id.clone(),
|
||||
endpoint_id: candidate.endpoint_id.clone(),
|
||||
key_id: candidate.key_id.clone(),
|
||||
upstream_base_url: transport.endpoint.base_url.clone(),
|
||||
upstream_url,
|
||||
provider_request_method: None,
|
||||
auth_header: Some(auth_header),
|
||||
auth_value: Some(auth_value),
|
||||
provider_api_format,
|
||||
client_api_format: spec_metadata.api_format.to_string(),
|
||||
model_name: input.requested_model.clone(),
|
||||
mapped_model,
|
||||
prompt_cache_key,
|
||||
provider_request_headers,
|
||||
provider_request_body: Some(provider_request_body),
|
||||
provider_request_body_base64: None,
|
||||
content_type: Some("application/json".to_string()),
|
||||
proxy,
|
||||
transport_profile,
|
||||
timeouts,
|
||||
upstream_is_stream,
|
||||
report_kind: spec_metadata.report_kind.map(ToOwned::to_owned),
|
||||
report_context: Some(report_context),
|
||||
auth_context: input.auth_context.clone(),
|
||||
});
|
||||
apply_provider_request_routing_policy_to_decision(input, &mut decision)?;
|
||||
Ok(Some(decision))
|
||||
}
|
||||
|
||||
@@ -236,6 +236,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
|
||||
);
|
||||
let force_body_stream_field =
|
||||
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let Some(mut base_provider_request_body) = (if needs_bidirectional_conversion {
|
||||
build_cross_format_openai_responses_request_body(
|
||||
body_json,
|
||||
@@ -251,7 +252,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
|
||||
transport.endpoint.body_rules.as_ref()
|
||||
},
|
||||
Some(input.auth_context.api_key_id.as_str()),
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
enable_model_directives,
|
||||
)
|
||||
} else {
|
||||
@@ -268,7 +269,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
|
||||
transport.endpoint.body_rules.as_ref()
|
||||
},
|
||||
Some(input.auth_context.api_key_id.as_str()),
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
enable_model_directives,
|
||||
)
|
||||
}) else {
|
||||
@@ -432,7 +433,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
|
||||
transport,
|
||||
provider_api_format,
|
||||
same_format,
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
auth_header: &auth_header,
|
||||
auth_value: &auth_value,
|
||||
extra_headers: &extra_headers,
|
||||
@@ -463,7 +464,7 @@ pub(crate) async fn resolve_local_openai_responses_candidate_payload_parts(
|
||||
apply_codex_openai_responses_special_headers(
|
||||
&mut provider_request_headers,
|
||||
&provider_request_body,
|
||||
&parts.headers,
|
||||
effective_headers,
|
||||
transport.provider.provider_type.as_str(),
|
||||
provider_api_format,
|
||||
Some(trace_id),
|
||||
@@ -545,12 +546,13 @@ async fn build_kiro_openai_responses_payload_parts(
|
||||
kiro_auth: &KiroRequestAuth,
|
||||
) -> Option<LocalOpenAiResponsesCandidatePayloadParts> {
|
||||
let candidate = &eligible.candidate;
|
||||
let effective_headers = input.effective_headers(&parts.headers);
|
||||
let provider_request_body = match build_kiro_provider_request_body(
|
||||
&claude_request_body,
|
||||
&mapped_model,
|
||||
&kiro_auth.auth_config,
|
||||
transport.endpoint.body_rules.as_ref(),
|
||||
Some(&parts.headers),
|
||||
Some(effective_headers),
|
||||
) {
|
||||
Some(body) => body,
|
||||
None => {
|
||||
@@ -601,7 +603,7 @@ async fn build_kiro_openai_responses_payload_parts(
|
||||
}
|
||||
};
|
||||
let provider_request_headers = match build_kiro_provider_headers(KiroProviderHeadersInput {
|
||||
headers: &parts.headers,
|
||||
headers: effective_headers,
|
||||
provider_request_body: &provider_request_body,
|
||||
original_request_body: original_body_json,
|
||||
header_rules: transport.endpoint.header_rules.as_ref(),
|
||||
|
||||
@@ -19,6 +19,7 @@ use crate::ai_serving::planner::candidate_source::{
|
||||
};
|
||||
use crate::ai_serving::planner::common::extract_standard_requested_model;
|
||||
use crate::ai_serving::planner::decision_input::{
|
||||
attach_routing_policy_to_local_requested_model_input,
|
||||
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
|
||||
};
|
||||
use crate::ai_serving::planner::materialization_policy::{
|
||||
@@ -48,7 +49,7 @@ pub(crate) async fn resolve_local_openai_responses_decision_input(
|
||||
decision: &GatewayControlDecision,
|
||||
body_json: &serde_json::Value,
|
||||
plan_kind: &str,
|
||||
) -> Option<LocalOpenAiResponsesDecisionInput> {
|
||||
) -> Result<Option<LocalOpenAiResponsesDecisionInput>, GatewayError> {
|
||||
let Some(auth_context) = resolve_local_decision_execution_runtime_auth_context(decision) else {
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
@@ -65,7 +66,7 @@ pub(crate) async fn resolve_local_openai_responses_decision_input(
|
||||
extract_standard_requested_model(body_json).as_deref(),
|
||||
"missing_auth_context",
|
||||
);
|
||||
return None;
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let Some(requested_model) = extract_standard_requested_model(body_json) else {
|
||||
@@ -81,7 +82,7 @@ pub(crate) async fn resolve_local_openai_responses_decision_input(
|
||||
None,
|
||||
"missing_requested_model",
|
||||
);
|
||||
return None;
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
let resolved_input = match resolve_local_authenticated_decision_input(
|
||||
@@ -108,7 +109,7 @@ pub(crate) async fn resolve_local_openai_responses_decision_input(
|
||||
Some(requested_model.as_str()),
|
||||
"auth_snapshot_missing",
|
||||
);
|
||||
return None;
|
||||
return Ok(None);
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
@@ -124,14 +125,30 @@ pub(crate) async fn resolve_local_openai_responses_decision_input(
|
||||
Some(requested_model.as_str()),
|
||||
"auth_snapshot_read_failed",
|
||||
);
|
||||
return None;
|
||||
return Err(err);
|
||||
}
|
||||
};
|
||||
|
||||
let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
|
||||
input.request_auth_channel = decision.request_auth_channel.clone();
|
||||
input.client_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
|
||||
Some(input)
|
||||
if let Err(err) = attach_routing_policy_to_local_requested_model_input(
|
||||
state,
|
||||
parts,
|
||||
&mut input,
|
||||
body_json,
|
||||
"openai:responses",
|
||||
)
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
error = ?err,
|
||||
"gateway local openai responses decision routing profile resolution failed"
|
||||
);
|
||||
return Err(err);
|
||||
}
|
||||
Ok(Some(input))
|
||||
}
|
||||
|
||||
pub(crate) async fn materialize_local_openai_responses_candidate_attempts(
|
||||
@@ -157,6 +174,7 @@ pub(crate) async fn materialize_local_openai_responses_candidate_attempts(
|
||||
spec_metadata.require_streaming,
|
||||
input.required_capabilities.as_ref(),
|
||||
&input.auth_snapshot,
|
||||
input.routing_policy.as_ref(),
|
||||
input.client_session_affinity.as_ref(),
|
||||
true,
|
||||
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
|
||||
@@ -170,6 +188,7 @@ pub(crate) async fn materialize_local_openai_responses_candidate_attempts(
|
||||
Some(&input.auth_snapshot),
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
@@ -256,6 +275,7 @@ pub(crate) async fn build_local_openai_responses_candidate_attempt_source<'a>(
|
||||
&input.auth_snapshot,
|
||||
input.client_session_affinity.as_ref(),
|
||||
input.required_capabilities.as_ref(),
|
||||
input.routing_policy.as_ref(),
|
||||
sticky_session_token.as_deref(),
|
||||
input.request_auth_channel.as_deref(),
|
||||
persistence_policy,
|
||||
|
||||
@@ -103,10 +103,11 @@ pub(crate) async fn maybe_build_sync_local_openai_responses_decision_payload(
|
||||
let Some(input) = resolve_local_openai_responses_decision_input(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
|
||||
let (mut source, _) = build_local_openai_responses_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
@@ -117,7 +118,7 @@ pub(crate) async fn maybe_build_sync_local_openai_responses_decision_payload(
|
||||
if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
@@ -141,10 +142,11 @@ pub(crate) async fn maybe_build_stream_local_openai_responses_decision_payload(
|
||||
let Some(input) = resolve_local_openai_responses_decision_input(
|
||||
state, parts, trace_id, decision, body_json, plan_kind,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
let body_json = input.effective_body_json(body_json);
|
||||
|
||||
let (mut source, _) = build_local_openai_responses_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
@@ -155,7 +157,7 @@ pub(crate) async fn maybe_build_stream_local_openai_responses_decision_payload(
|
||||
if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
{
|
||||
return Ok(Some(payload));
|
||||
}
|
||||
|
||||
@@ -29,7 +29,7 @@ pub(crate) struct LocalOpenAiResponsesSyncAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
trace_id: &'a str,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_json: serde_json::Value,
|
||||
input: LocalOpenAiResponsesDecisionInput,
|
||||
spec: LocalOpenAiResponsesSpec,
|
||||
candidates: LocalOpenAiResponsesCandidateAttemptSource<'a>,
|
||||
@@ -39,7 +39,7 @@ pub(crate) struct LocalOpenAiResponsesStreamAttemptSource<'a> {
|
||||
state: &'a AppState,
|
||||
parts: &'a http::request::Parts,
|
||||
trace_id: &'a str,
|
||||
body_json: &'a serde_json::Value,
|
||||
body_json: serde_json::Value,
|
||||
input: LocalOpenAiResponsesDecisionInput,
|
||||
spec: LocalOpenAiResponsesSpec,
|
||||
candidates: LocalOpenAiResponsesCandidateAttemptSource<'a>,
|
||||
@@ -62,7 +62,7 @@ pub(super) async fn build_local_sync_attempt_source<'a>(
|
||||
body_json,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -74,8 +74,13 @@ pub(super) async fn build_local_sync_attempt_source<'a>(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
let (candidates, candidate_count) = build_local_openai_responses_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
&effective_body_json,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
@@ -88,7 +93,7 @@ pub(super) async fn build_local_sync_attempt_source<'a>(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
body_json: effective_body_json,
|
||||
input,
|
||||
spec,
|
||||
candidates,
|
||||
@@ -114,7 +119,7 @@ pub(super) async fn build_local_stream_attempt_source<'a>(
|
||||
body_json,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
@@ -126,8 +131,13 @@ pub(super) async fn build_local_stream_attempt_source<'a>(
|
||||
Some(input.requested_model.as_str()),
|
||||
"candidate_evaluation_incomplete",
|
||||
);
|
||||
let effective_body_json = input.effective_body_json(body_json).clone();
|
||||
let (candidates, candidate_count) = build_local_openai_responses_candidate_attempt_source(
|
||||
state, trace_id, &input, body_json, spec,
|
||||
state,
|
||||
trace_id,
|
||||
&input,
|
||||
&effective_body_json,
|
||||
spec,
|
||||
)
|
||||
.await?;
|
||||
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
|
||||
@@ -140,7 +150,7 @@ pub(super) async fn build_local_stream_attempt_source<'a>(
|
||||
state,
|
||||
parts,
|
||||
trace_id,
|
||||
body_json,
|
||||
body_json: effective_body_json,
|
||||
input,
|
||||
spec,
|
||||
candidates,
|
||||
@@ -214,19 +224,19 @@ impl LocalOpenAiResponsesSyncAttemptSource<'_> {
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
&self.input,
|
||||
attempt,
|
||||
self.spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
match build_openai_responses_sync_plan_from_decision(
|
||||
self.parts,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
payload,
|
||||
self.spec.compact,
|
||||
) {
|
||||
@@ -252,19 +262,19 @@ impl LocalOpenAiResponsesStreamAttemptSource<'_> {
|
||||
self.state,
|
||||
self.parts,
|
||||
self.trace_id,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
&self.input,
|
||||
attempt,
|
||||
self.spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(None);
|
||||
};
|
||||
|
||||
match build_openai_responses_stream_plan_from_decision(
|
||||
self.parts,
|
||||
self.body_json,
|
||||
&self.body_json,
|
||||
payload,
|
||||
self.spec.compact,
|
||||
) {
|
||||
@@ -298,7 +308,7 @@ pub(super) async fn build_local_sync_plan_and_reports(
|
||||
body_json,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
@@ -325,7 +335,7 @@ pub(super) async fn build_local_sync_plan_and_reports(
|
||||
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
@@ -370,7 +380,7 @@ pub(super) async fn build_local_stream_plan_and_reports(
|
||||
body_json,
|
||||
spec_metadata.decision_kind,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
@@ -397,7 +407,7 @@ pub(super) async fn build_local_stream_plan_and_reports(
|
||||
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
|
||||
state, parts, trace_id, body_json, &input, attempt, spec,
|
||||
)
|
||||
.await
|
||||
.await?
|
||||
else {
|
||||
continue;
|
||||
};
|
||||
|
||||
@@ -137,6 +137,11 @@ const PERMISSION_GROUPS: &[PermissionGroup] = &[
|
||||
label: "代理节点",
|
||||
assignable: true,
|
||||
},
|
||||
PermissionGroup {
|
||||
scope: "routing_profiles",
|
||||
label: "调度分组",
|
||||
assignable: true,
|
||||
},
|
||||
PermissionGroup {
|
||||
scope: "security",
|
||||
label: "安全",
|
||||
@@ -446,6 +451,9 @@ fn permission_key(scope: &str, access: &str) -> &'static str {
|
||||
("proxy_nodes", "read") => "admin:proxy_nodes:read",
|
||||
("proxy_nodes", "write") => "admin:proxy_nodes:write",
|
||||
("proxy_nodes", "admin") => "admin:proxy_nodes:admin",
|
||||
("routing_profiles", "read") => "admin:routing_profiles:read",
|
||||
("routing_profiles", "write") => "admin:routing_profiles:write",
|
||||
("routing_profiles", "admin") => "admin:routing_profiles:admin",
|
||||
("security", "read") => "admin:security:read",
|
||||
("security", "write") => "admin:security:write",
|
||||
("security", "admin") => "admin:security:admin",
|
||||
|
||||
@@ -14,6 +14,8 @@ mod observability_families;
|
||||
mod operations_families;
|
||||
#[path = "admin/provider_ops_routes.rs"]
|
||||
mod provider_ops_routes;
|
||||
#[path = "admin/routing_families.rs"]
|
||||
mod routing_families;
|
||||
#[path = "admin/system_families.rs"]
|
||||
mod system_families;
|
||||
|
||||
@@ -23,6 +25,7 @@ use model_provider_families::classify_admin_model_provider_family_route;
|
||||
use observability_families::classify_admin_observability_family_route;
|
||||
use operations_families::classify_admin_operations_family_route;
|
||||
use provider_ops_routes::classify_admin_provider_ops_routes;
|
||||
use routing_families::classify_admin_routing_family_route;
|
||||
use system_families::classify_admin_system_family_route;
|
||||
|
||||
pub(super) fn classify_admin_route(
|
||||
@@ -67,6 +70,10 @@ pub(super) fn classify_admin_route(
|
||||
classify_admin_system_family_route(method, normalized_path, normalized_path_no_trailing)
|
||||
{
|
||||
Some(route)
|
||||
} else if let Some(route) =
|
||||
classify_admin_routing_family_route(method, normalized_path_no_trailing)
|
||||
{
|
||||
Some(route)
|
||||
} else if let Some(route) = classify_admin_provider_ops_routes(method, normalized_path) {
|
||||
Some(route)
|
||||
} else if let Some(route) = classify_admin_model_provider_family_route(method, normalized_path)
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
use axum::http;
|
||||
|
||||
use super::{classified, ClassifiedRoute};
|
||||
|
||||
pub(super) fn classify_admin_routing_family_route(
|
||||
method: &http::Method,
|
||||
normalized_path_no_trailing: &str,
|
||||
) -> Option<ClassifiedRoute> {
|
||||
let path = normalized_path_no_trailing;
|
||||
if method == http::Method::GET && path == "/api/admin/routing/groups" {
|
||||
Some(routing_route("list_groups"))
|
||||
} else if method == http::Method::POST && path == "/api/admin/routing/groups" {
|
||||
Some(routing_route("create_group"))
|
||||
} else if method == http::Method::GET
|
||||
&& path.starts_with("/api/admin/routing/groups/")
|
||||
&& path.ends_with("/versions")
|
||||
&& path.matches('/').count() == 6
|
||||
{
|
||||
Some(routing_route("list_group_versions"))
|
||||
} else if method == http::Method::POST
|
||||
&& path.starts_with("/api/admin/routing/groups/")
|
||||
&& path.ends_with("/publish")
|
||||
&& path.matches('/').count() == 6
|
||||
{
|
||||
Some(routing_route("publish_group"))
|
||||
} else if method == http::Method::POST
|
||||
&& path.starts_with("/api/admin/routing/groups/")
|
||||
&& path.ends_with("/dry-run")
|
||||
&& path.matches('/').count() == 6
|
||||
{
|
||||
Some(routing_route("dry_run_group"))
|
||||
} else if method == http::Method::GET
|
||||
&& path.starts_with("/api/admin/routing/groups/")
|
||||
&& path.matches('/').count() == 5
|
||||
{
|
||||
Some(routing_route("get_group"))
|
||||
} else if method == http::Method::PATCH
|
||||
&& path.starts_with("/api/admin/routing/groups/")
|
||||
&& path.matches('/').count() == 5
|
||||
{
|
||||
Some(routing_route("update_group"))
|
||||
} else if method == http::Method::DELETE
|
||||
&& path.starts_with("/api/admin/routing/groups/")
|
||||
&& path.matches('/').count() == 5
|
||||
{
|
||||
Some(routing_route("delete_group"))
|
||||
} else if method == http::Method::GET && path == "/api/admin/routing/bindings" {
|
||||
Some(routing_route("list_bindings"))
|
||||
} else if method == http::Method::POST && path == "/api/admin/routing/bindings" {
|
||||
Some(routing_route("create_binding"))
|
||||
} else if method == http::Method::PATCH
|
||||
&& path.starts_with("/api/admin/routing/bindings/")
|
||||
&& path.matches('/').count() == 5
|
||||
{
|
||||
Some(routing_route("update_binding"))
|
||||
} else if method == http::Method::DELETE
|
||||
&& path.starts_with("/api/admin/routing/bindings/")
|
||||
&& path.matches('/').count() == 5
|
||||
{
|
||||
Some(routing_route("delete_binding"))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn routing_route(route_kind: &'static str) -> ClassifiedRoute {
|
||||
classified(
|
||||
"admin_proxy",
|
||||
"routing_profiles_manage",
|
||||
route_kind,
|
||||
"admin:routing_profiles",
|
||||
false,
|
||||
)
|
||||
}
|
||||
92
apps/aether-gateway/src/control/tests/admin_routing.rs
Normal file
92
apps/aether-gateway/src/control/tests/admin_routing.rs
Normal file
@@ -0,0 +1,92 @@
|
||||
use http::Uri;
|
||||
|
||||
use crate::handlers::shared::local_proxy_route_requires_buffered_body;
|
||||
|
||||
use super::{classify_control_route, headers, GatewayPublicRequestContext};
|
||||
|
||||
#[test]
|
||||
fn classifies_admin_routing_group_routes_as_admin_proxy_route() {
|
||||
let headers = headers(&[]);
|
||||
|
||||
let list_uri: Uri = "/api/admin/routing/groups"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let list = classify_control_route(&http::Method::GET, &list_uri, &headers)
|
||||
.expect("route should classify");
|
||||
assert_eq!(list.route_class.as_deref(), Some("admin_proxy"));
|
||||
assert_eq!(
|
||||
list.route_family.as_deref(),
|
||||
Some("routing_profiles_manage")
|
||||
);
|
||||
assert_eq!(list.route_kind.as_deref(), Some("list_groups"));
|
||||
assert_eq!(
|
||||
list.auth_endpoint_signature.as_deref(),
|
||||
Some("admin:routing_profiles")
|
||||
);
|
||||
|
||||
let create_uri: Uri = "/api/admin/routing/groups"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let create = classify_control_route(&http::Method::POST, &create_uri, &headers)
|
||||
.expect("route should classify");
|
||||
assert_eq!(
|
||||
create.route_family.as_deref(),
|
||||
Some("routing_profiles_manage")
|
||||
);
|
||||
assert_eq!(create.route_kind.as_deref(), Some("create_group"));
|
||||
|
||||
let update_uri: Uri = "/api/admin/routing/groups/group-1"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let update = classify_control_route(&http::Method::PATCH, &update_uri, &headers)
|
||||
.expect("route should classify");
|
||||
assert_eq!(
|
||||
update.route_family.as_deref(),
|
||||
Some("routing_profiles_manage")
|
||||
);
|
||||
assert_eq!(update.route_kind.as_deref(), Some("update_group"));
|
||||
|
||||
let dry_run_uri: Uri = "/api/admin/routing/groups/group-1/dry-run"
|
||||
.parse()
|
||||
.expect("uri should parse");
|
||||
let dry_run = classify_control_route(&http::Method::POST, &dry_run_uri, &headers)
|
||||
.expect("route should classify");
|
||||
assert_eq!(
|
||||
dry_run.route_family.as_deref(),
|
||||
Some("routing_profiles_manage")
|
||||
);
|
||||
assert_eq!(dry_run.route_kind.as_deref(), Some("dry_run_group"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn admin_routing_write_routes_buffer_request_body() {
|
||||
let headers = headers(&[]);
|
||||
let routes = [
|
||||
(http::Method::POST, "/api/admin/routing/groups"),
|
||||
(http::Method::PATCH, "/api/admin/routing/groups/group-1"),
|
||||
(
|
||||
http::Method::POST,
|
||||
"/api/admin/routing/groups/group-1/dry-run",
|
||||
),
|
||||
(http::Method::POST, "/api/admin/routing/bindings"),
|
||||
(http::Method::PATCH, "/api/admin/routing/bindings/binding-1"),
|
||||
];
|
||||
|
||||
for (method, path) in routes {
|
||||
let uri: Uri = path.parse().expect("uri should parse");
|
||||
let decision =
|
||||
classify_control_route(&method, &uri, &headers).expect("route should classify");
|
||||
let context = GatewayPublicRequestContext::from_request_parts(
|
||||
"trace-routing-write",
|
||||
&method,
|
||||
&uri,
|
||||
&headers,
|
||||
Some(decision),
|
||||
);
|
||||
|
||||
assert!(
|
||||
local_proxy_route_requires_buffered_body(&context),
|
||||
"{method} {path} should buffer request body"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -86,6 +86,7 @@ mod admin_provider_query;
|
||||
mod admin_provider_strategy;
|
||||
mod admin_providers_models;
|
||||
mod admin_proxy_nodes;
|
||||
mod admin_routing;
|
||||
mod admin_security;
|
||||
mod admin_stats;
|
||||
mod admin_usage;
|
||||
|
||||
@@ -49,6 +49,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -92,6 +94,8 @@ impl GatewayDataState {
|
||||
let pool_score_writer = backends.write().pool_scores();
|
||||
let provider_quota_reader = backends.read().provider_quotas();
|
||||
let provider_quota_writer = backends.write().provider_quotas();
|
||||
let routing_group_reader = backends.read().routing_groups();
|
||||
let routing_group_writer = backends.write().routing_groups();
|
||||
let usage_reader = backends.read().usage();
|
||||
let usage_writer = backends.write().usage();
|
||||
let user_reader = backends.read().users();
|
||||
@@ -133,6 +137,8 @@ impl GatewayDataState {
|
||||
pool_score_writer,
|
||||
provider_quota_reader,
|
||||
provider_quota_writer,
|
||||
routing_group_reader,
|
||||
routing_group_writer,
|
||||
usage_reader,
|
||||
usage_writer,
|
||||
user_reader,
|
||||
@@ -261,6 +267,14 @@ impl GatewayDataState {
|
||||
self.request_candidate_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_routing_group_reader(&self) -> bool {
|
||||
self.routing_group_reader.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_routing_group_writer(&self) -> bool {
|
||||
self.routing_group_writer.is_some()
|
||||
}
|
||||
|
||||
pub(crate) fn has_provider_catalog_reader(&self) -> bool {
|
||||
self.provider_catalog_reader.is_some()
|
||||
}
|
||||
|
||||
@@ -125,6 +125,9 @@ use aether_data_contracts::repository::provider_catalog::{
|
||||
use aether_data_contracts::repository::quota::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
|
||||
};
|
||||
use aether_data_contracts::repository::routing_profiles::{
|
||||
RoutingGroupReadRepository, RoutingGroupWriteRepository,
|
||||
};
|
||||
use aether_data_contracts::repository::settlement::{
|
||||
SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput,
|
||||
};
|
||||
@@ -170,6 +173,8 @@ pub(crate) struct GatewayDataState {
|
||||
pool_score_writer: Option<Arc<dyn PoolMemberScoreWriteRepository>>,
|
||||
provider_quota_reader: Option<Arc<dyn ProviderQuotaReadRepository>>,
|
||||
provider_quota_writer: Option<Arc<dyn ProviderQuotaWriteRepository>>,
|
||||
routing_group_reader: Option<Arc<dyn RoutingGroupReadRepository>>,
|
||||
routing_group_writer: Option<Arc<dyn RoutingGroupWriteRepository>>,
|
||||
usage_reader: Option<Arc<dyn UsageReadRepository>>,
|
||||
usage_writer: Option<Arc<dyn UsageWriteRepository>>,
|
||||
user_reader: Option<Arc<dyn UserReadRepository>>,
|
||||
@@ -279,6 +284,14 @@ impl fmt::Debug for GatewayDataState {
|
||||
"has_provider_quota_writer",
|
||||
&self.provider_quota_writer.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_routing_group_reader",
|
||||
&self.routing_group_reader.is_some(),
|
||||
)
|
||||
.field(
|
||||
"has_routing_group_writer",
|
||||
&self.routing_group_writer.is_some(),
|
||||
)
|
||||
.field("has_usage_reader", &self.usage_reader.is_some())
|
||||
.field("has_usage_writer", &self.usage_writer.is_some())
|
||||
.field("has_user_preferences", &self.user_preferences.is_some())
|
||||
@@ -302,6 +315,7 @@ mod core;
|
||||
mod integrations;
|
||||
mod models;
|
||||
mod pool_scores;
|
||||
mod routing_profiles;
|
||||
mod runtime;
|
||||
#[cfg(test)]
|
||||
mod testing;
|
||||
|
||||
131
apps/aether-gateway/src/data/state/routing_profiles.rs
Normal file
131
apps/aether-gateway/src/data/state/routing_profiles.rs
Normal file
@@ -0,0 +1,131 @@
|
||||
use aether_data_contracts::repository::routing_profiles::{
|
||||
CreateRoutingGroupBindingRecord, CreateRoutingGroupRecord, CreateRoutingGroupVersionRecord,
|
||||
RoutingGroupBindingQuery, RoutingGroupLookupKey, RoutingGroupReadRepository,
|
||||
StoredRoutingGroup, StoredRoutingGroupBinding, StoredRoutingGroupVersion,
|
||||
UpdateRoutingGroupBindingRecord, UpdateRoutingGroupRecord,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::{DataLayerError, GatewayDataState};
|
||||
|
||||
impl GatewayDataState {
|
||||
pub(crate) fn routing_group_read_repository(
|
||||
&self,
|
||||
) -> Option<Arc<dyn RoutingGroupReadRepository>> {
|
||||
self.routing_group_reader.clone()
|
||||
}
|
||||
|
||||
pub(crate) async fn list_routing_groups(
|
||||
&self,
|
||||
) -> Result<Vec<StoredRoutingGroup>, DataLayerError> {
|
||||
match &self.routing_group_reader {
|
||||
Some(repository) => repository.list_routing_groups().await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn find_routing_group(
|
||||
&self,
|
||||
lookup: RoutingGroupLookupKey<'_>,
|
||||
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
match &self.routing_group_reader {
|
||||
Some(repository) => repository.find_routing_group(lookup).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_routing_group_bindings(
|
||||
&self,
|
||||
query: &RoutingGroupBindingQuery,
|
||||
) -> Result<Vec<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
match &self.routing_group_reader {
|
||||
Some(repository) => repository.list_routing_group_bindings(query).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn list_routing_group_versions(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Vec<StoredRoutingGroupVersion>, DataLayerError> {
|
||||
match &self.routing_group_reader {
|
||||
Some(repository) => repository.list_routing_group_versions(group_id).await,
|
||||
None => Ok(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_routing_group(
|
||||
&self,
|
||||
record: CreateRoutingGroupRecord,
|
||||
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
match &self.routing_group_writer {
|
||||
Some(repository) => repository.create_routing_group(record).await.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_routing_group(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupRecord,
|
||||
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
match &self.routing_group_writer {
|
||||
Some(repository) => repository.update_routing_group(id, patch).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_routing_group(&self, id: &str) -> Result<bool, DataLayerError> {
|
||||
match &self.routing_group_writer {
|
||||
Some(repository) => repository.delete_routing_group(id).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_routing_group_binding(
|
||||
&self,
|
||||
record: CreateRoutingGroupBindingRecord,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
match &self.routing_group_writer {
|
||||
Some(repository) => repository
|
||||
.create_routing_group_binding(record)
|
||||
.await
|
||||
.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn update_routing_group_binding(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupBindingRecord,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
match &self.routing_group_writer {
|
||||
Some(repository) => repository.update_routing_group_binding(id, patch).await,
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_routing_group_binding(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<bool, DataLayerError> {
|
||||
match &self.routing_group_writer {
|
||||
Some(repository) => repository.delete_routing_group_binding(id).await,
|
||||
None => Ok(false),
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) async fn create_routing_group_version(
|
||||
&self,
|
||||
record: CreateRoutingGroupVersionRecord,
|
||||
) -> Result<Option<StoredRoutingGroupVersion>, DataLayerError> {
|
||||
match &self.routing_group_writer {
|
||||
Some(repository) => repository
|
||||
.create_routing_group_version(record)
|
||||
.await
|
||||
.map(Some),
|
||||
None => Ok(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -38,6 +38,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -92,6 +94,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
|
||||
@@ -73,6 +73,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -126,6 +128,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -175,6 +179,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -312,6 +318,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -389,6 +397,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -447,6 +457,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: Some(provider_quota_reader),
|
||||
provider_quota_writer: Some(provider_quota_writer),
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -514,6 +526,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: Some(provider_quota_reader),
|
||||
provider_quota_writer: Some(provider_quota_writer),
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -563,6 +577,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -613,6 +629,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -674,6 +692,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: Some(usage_reader),
|
||||
usage_writer: Some(usage_writer),
|
||||
user_reader: None,
|
||||
@@ -737,6 +757,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -784,6 +806,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: Some(repository),
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -846,6 +870,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: Some(repository),
|
||||
@@ -901,6 +927,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: Some(user_repository),
|
||||
@@ -961,6 +989,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: Some(usage_reader),
|
||||
usage_writer: Some(usage_writer),
|
||||
user_reader: Some(user_repository),
|
||||
@@ -1022,6 +1052,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: Some(user_repository),
|
||||
@@ -1082,6 +1114,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: Some(usage_reader),
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1131,6 +1165,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1180,6 +1216,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1241,6 +1279,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1307,6 +1347,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1356,6 +1398,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1410,6 +1454,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1481,6 +1527,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1547,6 +1595,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1597,6 +1647,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1647,6 +1699,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1699,6 +1753,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: Some(usage_repository),
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1749,6 +1805,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1799,6 +1857,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1849,6 +1909,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1907,6 +1969,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: Some(provider_quota_reader),
|
||||
provider_quota_writer: Some(provider_quota_writer),
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -1966,6 +2030,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: Some(provider_quota_reader),
|
||||
provider_quota_writer: Some(provider_quota_writer),
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -2028,6 +2094,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: Some(provider_quota_reader),
|
||||
provider_quota_writer: Some(provider_quota_writer),
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -2096,6 +2164,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -2165,6 +2235,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -2238,6 +2310,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: Some(usage_reader),
|
||||
usage_writer: Some(usage_writer),
|
||||
user_reader: None,
|
||||
@@ -2318,6 +2392,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: Some(usage_reader),
|
||||
usage_writer: Some(usage_writer),
|
||||
user_reader: None,
|
||||
@@ -2380,6 +2456,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: Some(provider_quota_reader),
|
||||
provider_quota_writer: Some(provider_quota_writer),
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -2433,6 +2511,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: Some(usage_reader),
|
||||
usage_writer: Some(usage_writer),
|
||||
user_reader: None,
|
||||
@@ -2482,6 +2562,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -2537,6 +2619,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -2596,6 +2680,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: Some(usage_reader),
|
||||
usage_writer: Some(usage_writer),
|
||||
user_reader: None,
|
||||
@@ -2656,6 +2742,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: Some(usage_reader),
|
||||
usage_writer: Some(usage_writer),
|
||||
user_reader: None,
|
||||
@@ -2709,6 +2797,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: Some(provider_quota_reader),
|
||||
provider_quota_writer: Some(provider_quota_writer),
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
|
||||
@@ -42,6 +42,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -98,6 +100,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -151,6 +155,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -208,6 +214,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -269,6 +277,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
@@ -339,6 +349,8 @@ impl GatewayDataState {
|
||||
pool_score_writer: None,
|
||||
provider_quota_reader: None,
|
||||
provider_quota_writer: None,
|
||||
routing_group_reader: None,
|
||||
routing_group_writer: None,
|
||||
usage_reader: None,
|
||||
usage_writer: None,
|
||||
user_reader: None,
|
||||
|
||||
@@ -16,6 +16,7 @@ use aether_pool_core::{
|
||||
PoolMemberSignals, PoolRuntimeState, PoolSchedulingConfig, PoolSchedulingPreset,
|
||||
};
|
||||
use aether_provider_pool::ProviderPoolService;
|
||||
use aether_routing_core::{RankingOverlay, ResolvedRoutingPolicy};
|
||||
use tracing::warn;
|
||||
|
||||
use crate::ai_serving::{
|
||||
@@ -40,6 +41,7 @@ use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||
|
||||
static LOAD_BALANCE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
|
||||
const POOL_ACTIVE_PROBE_SEALED_SKIP_REASON: &str = "pool_active_probe_sealed";
|
||||
const ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON: &str = "routing_profile_disallowed_key";
|
||||
|
||||
type PoolCatalogKeyContext = PoolMemberSignals;
|
||||
|
||||
@@ -187,6 +189,7 @@ pub(crate) struct PoolKeyCursor<'a> {
|
||||
sticky_session_token: Option<String>,
|
||||
requested_model: Option<String>,
|
||||
request_auth_channel: Option<String>,
|
||||
routing_overlay: Option<RankingOverlay>,
|
||||
runtime_miss_trace_id: Option<String>,
|
||||
record_runtime_miss_diagnostic: bool,
|
||||
pool_key_order: StoredPoolKeyCandidateOrder,
|
||||
@@ -216,7 +219,26 @@ impl<'a> PoolKeyCursor<'a> {
|
||||
requested_model: Option<&str>,
|
||||
request_auth_channel: Option<&str>,
|
||||
) -> Self {
|
||||
let pool_key_order = pool_key_candidate_order_for_group(&group);
|
||||
Self::new_with_routing_policy(
|
||||
state,
|
||||
group,
|
||||
sticky_session_token,
|
||||
requested_model,
|
||||
request_auth_channel,
|
||||
None,
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn new_with_routing_policy(
|
||||
state: PlannerAppState<'a>,
|
||||
group: EligibleLocalExecutionCandidate,
|
||||
sticky_session_token: Option<&str>,
|
||||
requested_model: Option<&str>,
|
||||
request_auth_channel: Option<&str>,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
) -> Self {
|
||||
let pool_key_order = pool_key_candidate_order_for_group(&group, routing_policy);
|
||||
let routing_overlay = routing_policy.map(|policy| policy.ranking_overlay.clone());
|
||||
let pool_config = pool_config_for_candidate(&group);
|
||||
let score_top_n = pool_config
|
||||
.as_ref()
|
||||
@@ -236,6 +258,7 @@ impl<'a> PoolKeyCursor<'a> {
|
||||
sticky_session_token: sticky_session_token.map(str::to_string),
|
||||
requested_model: requested_model.map(str::to_string),
|
||||
request_auth_channel: request_auth_channel.map(str::to_string),
|
||||
routing_overlay,
|
||||
runtime_miss_trace_id: None,
|
||||
record_runtime_miss_diagnostic: false,
|
||||
pool_key_order,
|
||||
@@ -551,6 +574,9 @@ impl<'a> PoolKeyCursor<'a> {
|
||||
async fn next_queued_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
|
||||
while let Some(candidate) = self.queued_candidates.pop_front() {
|
||||
let mut candidate = candidate;
|
||||
if self.skip_candidate_if_routing_profile_disallowed(&candidate) {
|
||||
continue;
|
||||
}
|
||||
if self.skip_candidate_if_runtime_cooldown(&candidate).await {
|
||||
continue;
|
||||
}
|
||||
@@ -562,6 +588,28 @@ impl<'a> PoolKeyCursor<'a> {
|
||||
None
|
||||
}
|
||||
|
||||
fn skip_candidate_if_routing_profile_disallowed(
|
||||
&mut self,
|
||||
candidate: &EligibleLocalExecutionCandidate,
|
||||
) -> bool {
|
||||
let Some(overlay) = self.routing_overlay.as_ref() else {
|
||||
return false;
|
||||
};
|
||||
if overlay.key_allowed(candidate.candidate.key_id.as_str()) {
|
||||
return false;
|
||||
}
|
||||
self.record_skip_reason(ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON);
|
||||
self.skipped_candidates
|
||||
.push(SkippedLocalExecutionCandidate {
|
||||
candidate: candidate.candidate.clone(),
|
||||
skip_reason: ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON,
|
||||
transport: Some(candidate.transport.clone()),
|
||||
ranking: candidate.ranking.clone(),
|
||||
extra_data: None,
|
||||
});
|
||||
true
|
||||
}
|
||||
|
||||
async fn skip_candidate_if_runtime_cooldown(
|
||||
&mut self,
|
||||
candidate: &EligibleLocalExecutionCandidate,
|
||||
@@ -958,19 +1006,38 @@ fn should_trigger_active_probe_burst_for_request(
|
||||
|
||||
fn pool_key_candidate_order_for_group(
|
||||
group: &EligibleLocalExecutionCandidate,
|
||||
routing_policy: Option<&ResolvedRoutingPolicy>,
|
||||
) -> StoredPoolKeyCandidateOrder {
|
||||
let Some(pool_config) = pool_config_for_candidate(group) else {
|
||||
return StoredPoolKeyCandidateOrder::InternalPriority;
|
||||
};
|
||||
let presets = pool_config
|
||||
.scheduling_presets
|
||||
.iter()
|
||||
.map(|preset| PoolSchedulingPreset {
|
||||
preset: preset.preset.clone(),
|
||||
enabled: preset.enabled,
|
||||
mode: preset.mode.clone(),
|
||||
let override_presets = routing_policy
|
||||
.and_then(|policy| {
|
||||
policy
|
||||
.pool_policy_overrides
|
||||
.get(group.candidate.provider_id.as_str())
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
.filter(|override_policy| !override_policy.scheduling_presets.is_empty());
|
||||
let presets = match override_presets {
|
||||
Some(override_policy) => override_policy
|
||||
.scheduling_presets
|
||||
.iter()
|
||||
.map(|preset| PoolSchedulingPreset {
|
||||
preset: preset.preset.clone(),
|
||||
enabled: preset.enabled,
|
||||
mode: preset.mode.clone(),
|
||||
})
|
||||
.collect::<Vec<_>>(),
|
||||
None => pool_config
|
||||
.scheduling_presets
|
||||
.iter()
|
||||
.map(|preset| PoolSchedulingPreset {
|
||||
preset: preset.preset.clone(),
|
||||
enabled: preset.enabled,
|
||||
mode: preset.mode.clone(),
|
||||
})
|
||||
.collect::<Vec<_>>(),
|
||||
};
|
||||
let active_presets = ProviderPoolService::with_builtin_adapters()
|
||||
.normalize_scheduling_presets(group.transport.provider.provider_type.as_str(), &presets)
|
||||
.into_iter()
|
||||
@@ -1071,6 +1138,7 @@ mod tests {
|
||||
apply_local_execution_pool_scheduler_with_runtime_map, build_pool_catalog_key_context,
|
||||
pool_config_for_candidate, should_trigger_active_probe_burst_for_request,
|
||||
PoolCatalogKeyContext, PoolKeyCursor, POOL_ACTIVE_PROBE_SEALED_SKIP_REASON,
|
||||
ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON,
|
||||
};
|
||||
use crate::ai_serving::{
|
||||
apply_local_runtime_candidate_terminal_reason, EligibleLocalExecutionCandidate,
|
||||
@@ -1096,6 +1164,9 @@ mod tests {
|
||||
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
|
||||
GatewayProviderTransportProvider,
|
||||
};
|
||||
use aether_routing_core::{
|
||||
RankingOverlay, ResolvedRoutingPolicy, RoutingSchedulingMode, RoutingSetPriorityMode,
|
||||
};
|
||||
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
||||
use serde_json::json;
|
||||
use std::collections::{BTreeMap, BTreeSet, VecDeque};
|
||||
@@ -2103,6 +2174,59 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_key_cursor_filters_expanded_keys_by_routing_profile_allowed_keys() {
|
||||
let app = AppState::new().expect("state should build");
|
||||
let provider_config = Some(json!({ "pool_advanced": { "lru_enabled": true } }));
|
||||
let group = sample_eligible_candidate(
|
||||
"provider-pool",
|
||||
"endpoint-1",
|
||||
"pool-group",
|
||||
10,
|
||||
provider_config.clone(),
|
||||
);
|
||||
let routing_policy = routing_policy_with_allowed_keys(["key-b"]);
|
||||
let mut cursor = PoolKeyCursor::new_with_routing_policy(
|
||||
PlannerAppState::new(&app),
|
||||
group,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some(&routing_policy),
|
||||
);
|
||||
cursor.queued_candidates = VecDeque::from([
|
||||
sample_eligible_candidate(
|
||||
"provider-pool",
|
||||
"endpoint-1",
|
||||
"key-a",
|
||||
10,
|
||||
provider_config.clone(),
|
||||
),
|
||||
sample_eligible_candidate("provider-pool", "endpoint-1", "key-b", 10, provider_config),
|
||||
]);
|
||||
|
||||
let candidate = cursor
|
||||
.next_key()
|
||||
.await
|
||||
.expect("cursor should skip disallowed pool key and return allowed key");
|
||||
assert_eq!(candidate.candidate.key_id, "key-b");
|
||||
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
|
||||
assert_eq!(
|
||||
cursor
|
||||
.skip_reason_counts
|
||||
.get(ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON),
|
||||
Some(&1)
|
||||
);
|
||||
let skipped = cursor.take_skipped_candidates();
|
||||
assert_eq!(
|
||||
skipped
|
||||
.iter()
|
||||
.map(|item| (item.candidate.key_id.as_str(), item.skip_reason))
|
||||
.collect::<Vec<_>>(),
|
||||
vec![("key-a", ROUTING_PROFILE_DISALLOWED_KEY_SKIP_REASON)]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn pool_key_cursor_allows_parallel_requests_to_use_same_healthy_key() {
|
||||
let app = AppState::new().expect("state should build");
|
||||
@@ -2668,6 +2792,28 @@ mod tests {
|
||||
(provider, endpoint, keys, rows)
|
||||
}
|
||||
|
||||
fn routing_policy_with_allowed_keys<const N: usize>(
|
||||
key_ids: [&str; N],
|
||||
) -> ResolvedRoutingPolicy {
|
||||
ResolvedRoutingPolicy {
|
||||
group_id: Some("routing-group-1".to_string()),
|
||||
group_version: Some(1),
|
||||
selection_source: "test".to_string(),
|
||||
requested_model: "gpt-5".to_string(),
|
||||
resolved_model: "gpt-5".to_string(),
|
||||
priority_mode: RoutingSetPriorityMode::Provider,
|
||||
scheduling_mode: RoutingSchedulingMode::CacheAffinity,
|
||||
keep_priority_on_conversion: false,
|
||||
ranking_overlay: RankingOverlay {
|
||||
allowed_keys: key_ids.into_iter().map(str::to_string).collect(),
|
||||
..RankingOverlay::default()
|
||||
},
|
||||
mutation_plan: Default::default(),
|
||||
pool_policy_overrides: BTreeMap::new(),
|
||||
matched_rules: Vec::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_eligible_candidate(
|
||||
provider_id: &str,
|
||||
endpoint_id: &str,
|
||||
|
||||
@@ -6,6 +6,7 @@ pub(super) mod features;
|
||||
mod model;
|
||||
pub(super) mod observability;
|
||||
pub(super) mod provider;
|
||||
mod routing;
|
||||
mod system;
|
||||
mod users;
|
||||
|
||||
|
||||
@@ -64,6 +64,14 @@ impl<'a> AdminAppState<'a> {
|
||||
self.app.has_global_model_data_writer()
|
||||
}
|
||||
|
||||
pub(crate) fn has_routing_group_data_reader(&self) -> bool {
|
||||
self.app.has_routing_group_data_reader()
|
||||
}
|
||||
|
||||
pub(crate) fn has_routing_group_data_writer(&self) -> bool {
|
||||
self.app.has_routing_group_data_writer()
|
||||
}
|
||||
|
||||
pub(crate) fn has_usage_data_reader(&self) -> bool {
|
||||
self.app.has_usage_data_reader()
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ mod observability;
|
||||
mod provider;
|
||||
mod provider_oauth;
|
||||
mod route_request;
|
||||
mod routing_profiles;
|
||||
mod state;
|
||||
mod system;
|
||||
mod users;
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
use aether_data_contracts::repository::routing_profiles::{
|
||||
CreateRoutingGroupBindingRecord, CreateRoutingGroupRecord, CreateRoutingGroupVersionRecord,
|
||||
RoutingGroupBindingQuery, RoutingGroupLookupKey, StoredRoutingGroup, StoredRoutingGroupBinding,
|
||||
StoredRoutingGroupVersion, UpdateRoutingGroupBindingRecord, UpdateRoutingGroupRecord,
|
||||
};
|
||||
|
||||
use super::AdminAppState;
|
||||
use crate::GatewayError;
|
||||
|
||||
impl<'a> AdminAppState<'a> {
|
||||
pub(crate) async fn list_routing_groups(
|
||||
&self,
|
||||
) -> Result<Vec<StoredRoutingGroup>, GatewayError> {
|
||||
self.app.list_routing_groups().await
|
||||
}
|
||||
|
||||
pub(crate) async fn find_routing_group(
|
||||
&self,
|
||||
lookup: RoutingGroupLookupKey<'_>,
|
||||
) -> Result<Option<StoredRoutingGroup>, GatewayError> {
|
||||
self.app.find_routing_group(lookup).await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_routing_group_bindings(
|
||||
&self,
|
||||
query: &RoutingGroupBindingQuery,
|
||||
) -> Result<Vec<StoredRoutingGroupBinding>, GatewayError> {
|
||||
self.app.list_routing_group_bindings(query).await
|
||||
}
|
||||
|
||||
pub(crate) async fn list_routing_group_versions(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Vec<StoredRoutingGroupVersion>, GatewayError> {
|
||||
self.app.list_routing_group_versions(group_id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn create_routing_group(
|
||||
&self,
|
||||
record: CreateRoutingGroupRecord,
|
||||
) -> Result<Option<StoredRoutingGroup>, GatewayError> {
|
||||
self.app.create_routing_group(record).await
|
||||
}
|
||||
|
||||
pub(crate) async fn update_routing_group(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupRecord,
|
||||
) -> Result<Option<StoredRoutingGroup>, GatewayError> {
|
||||
self.app.update_routing_group(id, patch).await
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_routing_group(&self, id: &str) -> Result<bool, GatewayError> {
|
||||
self.app.delete_routing_group(id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn create_routing_group_binding(
|
||||
&self,
|
||||
record: CreateRoutingGroupBindingRecord,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, GatewayError> {
|
||||
self.app.create_routing_group_binding(record).await
|
||||
}
|
||||
|
||||
pub(crate) async fn update_routing_group_binding(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupBindingRecord,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, GatewayError> {
|
||||
self.app.update_routing_group_binding(id, patch).await
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_routing_group_binding(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<bool, GatewayError> {
|
||||
self.app.delete_routing_group_binding(id).await
|
||||
}
|
||||
|
||||
pub(crate) async fn create_routing_group_version(
|
||||
&self,
|
||||
record: CreateRoutingGroupVersionRecord,
|
||||
) -> Result<Option<StoredRoutingGroupVersion>, GatewayError> {
|
||||
self.app.create_routing_group_version(record).await
|
||||
}
|
||||
}
|
||||
@@ -1,6 +1,6 @@
|
||||
use super::{
|
||||
announcements, auth, billing, endpoint, features, model, observability, provider, request,
|
||||
system, users,
|
||||
routing, system, users,
|
||||
};
|
||||
|
||||
pub(crate) async fn maybe_build_local_admin_response(
|
||||
@@ -20,6 +20,10 @@ pub(crate) async fn maybe_build_local_admin_response(
|
||||
return Ok(Some(response));
|
||||
}
|
||||
|
||||
if let Some(response) = routing::maybe_build_local_admin_routing_response(request).await? {
|
||||
return Ok(Some(response));
|
||||
}
|
||||
|
||||
if let Some(response) = auth::maybe_build_local_admin_auth_response(request).await? {
|
||||
return Ok(Some(response));
|
||||
}
|
||||
|
||||
760
apps/aether-gateway/src/handlers/admin/routing/mod.rs
Normal file
760
apps/aether-gateway/src/handlers/admin/routing/mod.rs
Normal file
@@ -0,0 +1,760 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_data_contracts::repository::routing_profiles::{
|
||||
CreateRoutingGroupBindingRecord, CreateRoutingGroupRecord, CreateRoutingGroupVersionRecord,
|
||||
RoutingGroupBindingQuery, RoutingGroupBindingSubject, RoutingGroupLookupKey,
|
||||
StoredRoutingGroup, StoredRoutingGroupBinding, StoredRoutingGroupVersion,
|
||||
UpdateRoutingGroupBindingRecord, UpdateRoutingGroupRecord,
|
||||
};
|
||||
use aether_routing_core::{
|
||||
validate_routing_group_config, MutationPlan, RoutingGroupConfig, RoutingHeaderPatch,
|
||||
RoutingPatchSummary, RoutingRulePhase,
|
||||
};
|
||||
use axum::{
|
||||
body::{Body, Bytes},
|
||||
http::{self, HeaderMap, HeaderName, HeaderValue},
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Map, Value};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::clock::current_unix_secs;
|
||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||
use crate::handlers::admin::shared::{attach_admin_audit_response, query_param_value};
|
||||
use crate::routing::{
|
||||
apply_routing_mutation_plan, build_routing_trace_seed, resolve_gateway_routing_policy,
|
||||
GatewayRoutingPolicyInput,
|
||||
};
|
||||
use crate::GatewayError;
|
||||
|
||||
const ROUTING_GROUPS_ROOT: &str = "/api/admin/routing/groups";
|
||||
const ROUTING_BINDINGS_ROOT: &str = "/api/admin/routing/bindings";
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AdminRoutingGroupCreateRequest {
|
||||
#[serde(default)]
|
||||
id: Option<String>,
|
||||
name: String,
|
||||
#[serde(default)]
|
||||
description: Option<String>,
|
||||
#[serde(default = "default_true")]
|
||||
enabled: bool,
|
||||
#[serde(default)]
|
||||
is_system_default: bool,
|
||||
#[serde(default)]
|
||||
config_json: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AdminRoutingGroupBindingCreateRequest {
|
||||
#[serde(default)]
|
||||
id: Option<String>,
|
||||
group_id: String,
|
||||
subject_type: RoutingGroupBindingSubject,
|
||||
subject_id: String,
|
||||
#[serde(default)]
|
||||
is_default: bool,
|
||||
#[serde(default)]
|
||||
allow_explicit_select: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct AdminRoutingDryRunRequest {
|
||||
model: String,
|
||||
#[serde(default)]
|
||||
resolved_model: Option<String>,
|
||||
#[serde(default = "default_api_format")]
|
||||
api_format: String,
|
||||
#[serde(default)]
|
||||
user_id: Option<String>,
|
||||
#[serde(default)]
|
||||
api_key_id: Option<String>,
|
||||
#[serde(default)]
|
||||
headers: Option<Value>,
|
||||
#[serde(default)]
|
||||
body: Option<Value>,
|
||||
#[serde(default)]
|
||||
phase: Option<RoutingRulePhase>,
|
||||
}
|
||||
|
||||
pub(crate) async fn maybe_build_local_admin_routing_response(
|
||||
request: crate::handlers::admin::request::AdminRouteRequest<'_>,
|
||||
) -> crate::handlers::admin::request::AdminRouteResult {
|
||||
let state = request.state();
|
||||
let request_context = request.request_context();
|
||||
let request_body = request.request_body();
|
||||
|
||||
if request_context.route_family() != Some("routing_profiles_manage") {
|
||||
return Ok(None);
|
||||
}
|
||||
if !request_context.path().starts_with("/api/admin/routing/") {
|
||||
return Ok(None);
|
||||
}
|
||||
if !state.has_routing_group_data_reader() {
|
||||
return Ok(Some(data_unavailable_response()));
|
||||
}
|
||||
|
||||
let response = if request_context.path().starts_with(ROUTING_GROUPS_ROOT) {
|
||||
maybe_build_routing_groups_response(&state, &request_context, request_body).await?
|
||||
} else if request_context.path().starts_with(ROUTING_BINDINGS_ROOT) {
|
||||
maybe_build_routing_bindings_response(&state, &request_context, request_body).await?
|
||||
} else {
|
||||
None
|
||||
};
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
async fn maybe_build_routing_groups_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let path = normalized_admin_path(request_context.path());
|
||||
match (request_context.method(), path.as_str()) {
|
||||
(&http::Method::GET, ROUTING_GROUPS_ROOT) => {
|
||||
let groups = state.list_routing_groups().await?;
|
||||
Ok(Some(
|
||||
Json(json!({
|
||||
"items": groups.iter().map(routing_group_payload).collect::<Vec<_>>(),
|
||||
"total": groups.len(),
|
||||
}))
|
||||
.into_response(),
|
||||
))
|
||||
}
|
||||
(&http::Method::POST, ROUTING_GROUPS_ROOT) => {
|
||||
if !state.has_routing_group_data_writer() {
|
||||
return Ok(Some(data_unavailable_response()));
|
||||
}
|
||||
let payload = parse_json_body::<AdminRoutingGroupCreateRequest>(request_body)?;
|
||||
let config_json = payload.config_json.unwrap_or_else(|| json!({}));
|
||||
validate_config_json(&config_json)?;
|
||||
let now = current_unix_secs() as i64;
|
||||
let record = CreateRoutingGroupRecord {
|
||||
id: payload.id.unwrap_or_else(|| Uuid::new_v4().to_string()),
|
||||
name: payload.name,
|
||||
description: payload.description,
|
||||
enabled: payload.enabled,
|
||||
is_system_default: payload.is_system_default,
|
||||
config_json,
|
||||
version: 1,
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
published_at: None,
|
||||
};
|
||||
let Some(created) = state.create_routing_group(record).await? else {
|
||||
return Ok(Some(data_unavailable_response()));
|
||||
};
|
||||
Ok(Some(attach_admin_audit_response(
|
||||
Json(routing_group_payload(&created)).into_response(),
|
||||
"admin_routing_group_created",
|
||||
"create_routing_group",
|
||||
"routing_group",
|
||||
&created.id,
|
||||
)))
|
||||
}
|
||||
_ => {
|
||||
let Some((group_id, suffix)) = routing_group_path_parts(path.as_str()) else {
|
||||
return Ok(None);
|
||||
};
|
||||
match (request_context.method(), suffix.as_deref()) {
|
||||
(&http::Method::GET, None) => {
|
||||
let Some(group) = state
|
||||
.find_routing_group(RoutingGroupLookupKey::Id(&group_id))
|
||||
.await?
|
||||
else {
|
||||
return Ok(Some(not_found_response(format!(
|
||||
"routing group {group_id} not found"
|
||||
))));
|
||||
};
|
||||
Ok(Some(Json(routing_group_payload(&group)).into_response()))
|
||||
}
|
||||
(&http::Method::PATCH, None) => {
|
||||
if !state.has_routing_group_data_writer() {
|
||||
return Ok(Some(data_unavailable_response()));
|
||||
}
|
||||
let patch = build_routing_group_update_patch(request_body)?;
|
||||
let Some(updated) = state.update_routing_group(&group_id, patch).await? else {
|
||||
return Ok(Some(not_found_response(format!(
|
||||
"routing group {group_id} not found"
|
||||
))));
|
||||
};
|
||||
Ok(Some(attach_admin_audit_response(
|
||||
Json(routing_group_payload(&updated)).into_response(),
|
||||
"admin_routing_group_updated",
|
||||
"update_routing_group",
|
||||
"routing_group",
|
||||
&updated.id,
|
||||
)))
|
||||
}
|
||||
(&http::Method::DELETE, None) => {
|
||||
if !state.has_routing_group_data_writer() {
|
||||
return Ok(Some(data_unavailable_response()));
|
||||
}
|
||||
if !state.delete_routing_group(&group_id).await? {
|
||||
return Ok(Some(not_found_response(format!(
|
||||
"routing group {group_id} not found"
|
||||
))));
|
||||
}
|
||||
Ok(Some(attach_admin_audit_response(
|
||||
http::StatusCode::NO_CONTENT.into_response(),
|
||||
"admin_routing_group_deleted",
|
||||
"delete_routing_group",
|
||||
"routing_group",
|
||||
&group_id,
|
||||
)))
|
||||
}
|
||||
(&http::Method::POST, Some("publish")) => {
|
||||
publish_routing_group(state, &group_id).await
|
||||
}
|
||||
(&http::Method::GET, Some("versions")) => {
|
||||
let versions = state.list_routing_group_versions(&group_id).await?;
|
||||
Ok(Some(Json(json!({
|
||||
"items": versions.iter().map(routing_group_version_payload).collect::<Vec<_>>(),
|
||||
"total": versions.len(),
|
||||
})).into_response()))
|
||||
}
|
||||
(&http::Method::POST, Some("dry-run")) => {
|
||||
dry_run_routing_group(state, &group_id, request_body).await
|
||||
}
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn maybe_build_routing_bindings_response(
|
||||
state: &AdminAppState<'_>,
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let path = normalized_admin_path(request_context.path());
|
||||
match (request_context.method(), path.as_str()) {
|
||||
(&http::Method::GET, ROUTING_BINDINGS_ROOT) => {
|
||||
let query = routing_binding_query_from_request(request_context)?;
|
||||
let bindings = state.list_routing_group_bindings(&query).await?;
|
||||
Ok(Some(
|
||||
Json(json!({
|
||||
"items": bindings.iter().map(routing_group_binding_payload).collect::<Vec<_>>(),
|
||||
"total": bindings.len(),
|
||||
}))
|
||||
.into_response(),
|
||||
))
|
||||
}
|
||||
(&http::Method::POST, ROUTING_BINDINGS_ROOT) => {
|
||||
if !state.has_routing_group_data_writer() {
|
||||
return Ok(Some(data_unavailable_response()));
|
||||
}
|
||||
let payload = parse_json_body::<AdminRoutingGroupBindingCreateRequest>(request_body)?;
|
||||
let now = current_unix_secs() as i64;
|
||||
let record = CreateRoutingGroupBindingRecord {
|
||||
id: payload.id.unwrap_or_else(|| Uuid::new_v4().to_string()),
|
||||
group_id: payload.group_id,
|
||||
subject_type: payload.subject_type,
|
||||
subject_id: payload.subject_id,
|
||||
is_default: payload.is_default,
|
||||
allow_explicit_select: payload.allow_explicit_select,
|
||||
created_at: now,
|
||||
updated_at: now,
|
||||
};
|
||||
let Some(created) = state.create_routing_group_binding(record).await? else {
|
||||
return Ok(Some(data_unavailable_response()));
|
||||
};
|
||||
Ok(Some(attach_admin_audit_response(
|
||||
Json(routing_group_binding_payload(&created)).into_response(),
|
||||
"admin_routing_group_binding_created",
|
||||
"create_routing_group_binding",
|
||||
"routing_group_binding",
|
||||
&created.id,
|
||||
)))
|
||||
}
|
||||
_ => {
|
||||
let Some(binding_id) = routing_binding_id_from_path(path.as_str()) else {
|
||||
return Ok(None);
|
||||
};
|
||||
match *request_context.method() {
|
||||
http::Method::PATCH => {
|
||||
if !state.has_routing_group_data_writer() {
|
||||
return Ok(Some(data_unavailable_response()));
|
||||
}
|
||||
let patch = build_routing_binding_update_patch(request_body)?;
|
||||
let Some(updated) = state
|
||||
.update_routing_group_binding(&binding_id, patch)
|
||||
.await?
|
||||
else {
|
||||
return Ok(Some(not_found_response(format!(
|
||||
"routing group binding {binding_id} not found"
|
||||
))));
|
||||
};
|
||||
Ok(Some(attach_admin_audit_response(
|
||||
Json(routing_group_binding_payload(&updated)).into_response(),
|
||||
"admin_routing_group_binding_updated",
|
||||
"update_routing_group_binding",
|
||||
"routing_group_binding",
|
||||
&updated.id,
|
||||
)))
|
||||
}
|
||||
http::Method::DELETE => {
|
||||
if !state.has_routing_group_data_writer() {
|
||||
return Ok(Some(data_unavailable_response()));
|
||||
}
|
||||
if !state.delete_routing_group_binding(&binding_id).await? {
|
||||
return Ok(Some(not_found_response(format!(
|
||||
"routing group binding {binding_id} not found"
|
||||
))));
|
||||
}
|
||||
Ok(Some(attach_admin_audit_response(
|
||||
http::StatusCode::NO_CONTENT.into_response(),
|
||||
"admin_routing_group_binding_deleted",
|
||||
"delete_routing_group_binding",
|
||||
"routing_group_binding",
|
||||
&binding_id,
|
||||
)))
|
||||
}
|
||||
_ => Ok(None),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn publish_routing_group(
|
||||
state: &AdminAppState<'_>,
|
||||
group_id: &str,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
if !state.has_routing_group_data_writer() {
|
||||
return Ok(Some(data_unavailable_response()));
|
||||
}
|
||||
let Some(group) = state
|
||||
.find_routing_group(RoutingGroupLookupKey::Id(group_id))
|
||||
.await?
|
||||
else {
|
||||
return Ok(Some(not_found_response(format!(
|
||||
"routing group {group_id} not found"
|
||||
))));
|
||||
};
|
||||
validate_config_json(&group.config_json)?;
|
||||
let latest_version = state
|
||||
.list_routing_group_versions(group_id)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|version| version.version)
|
||||
.max()
|
||||
.unwrap_or(0);
|
||||
let next_version = group.version.max(latest_version.saturating_add(1));
|
||||
let now = current_unix_secs() as i64;
|
||||
let Some(updated) = state
|
||||
.update_routing_group(
|
||||
group_id,
|
||||
UpdateRoutingGroupRecord {
|
||||
version: Some(next_version),
|
||||
updated_at: now,
|
||||
published_at: Some(Some(now)),
|
||||
..UpdateRoutingGroupRecord::default()
|
||||
},
|
||||
)
|
||||
.await?
|
||||
else {
|
||||
return Ok(Some(not_found_response(format!(
|
||||
"routing group {group_id} not found"
|
||||
))));
|
||||
};
|
||||
let _ = state
|
||||
.create_routing_group_version(CreateRoutingGroupVersionRecord {
|
||||
id: Uuid::new_v4().to_string(),
|
||||
group_id: group_id.to_string(),
|
||||
version: next_version,
|
||||
config_json: updated.config_json.clone(),
|
||||
created_at: now,
|
||||
created_by: None,
|
||||
})
|
||||
.await?;
|
||||
|
||||
Ok(Some(attach_admin_audit_response(
|
||||
Json(routing_group_payload(&updated)).into_response(),
|
||||
"admin_routing_group_published",
|
||||
"publish_routing_group",
|
||||
"routing_group",
|
||||
group_id,
|
||||
)))
|
||||
}
|
||||
|
||||
async fn dry_run_routing_group(
|
||||
state: &AdminAppState<'_>,
|
||||
group_id: &str,
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<Option<Response<Body>>, GatewayError> {
|
||||
let Some(group) = state
|
||||
.find_routing_group(RoutingGroupLookupKey::Id(group_id))
|
||||
.await?
|
||||
else {
|
||||
return Ok(Some(not_found_response(format!(
|
||||
"routing group {group_id} not found"
|
||||
))));
|
||||
};
|
||||
let payload = parse_json_body::<AdminRoutingDryRunRequest>(request_body)?;
|
||||
let requested_model = payload.model.trim();
|
||||
if requested_model.is_empty() {
|
||||
return Ok(Some(bad_request_response("model must not be empty")));
|
||||
}
|
||||
let resolved_model = payload
|
||||
.resolved_model
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(requested_model);
|
||||
let api_format = payload.api_format.trim();
|
||||
let headers_json = payload.headers.unwrap_or_else(|| json!({}));
|
||||
let mut header_map = header_map_from_value(&headers_json)?;
|
||||
let mut body = payload.body.unwrap_or_else(|| json!({}));
|
||||
let policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput {
|
||||
group_id: Some(group.id.as_str()),
|
||||
group_version: Some(group.version),
|
||||
group_config_json: &group.config_json,
|
||||
selection_source: "admin_dry_run",
|
||||
requested_model,
|
||||
resolved_model,
|
||||
api_format,
|
||||
user_id: payload.user_id.as_deref(),
|
||||
api_key_id: payload.api_key_id.as_deref(),
|
||||
headers: &headers_json,
|
||||
body: &body,
|
||||
phase: payload.phase.unwrap_or(RoutingRulePhase::ClientRequest),
|
||||
})?;
|
||||
let patch_summary = patch_summary(&policy.mutation_plan);
|
||||
apply_routing_mutation_plan(&mut body, &mut header_map, &policy.mutation_plan)?;
|
||||
let mut trace = build_routing_trace_seed(&policy, api_format);
|
||||
trace.client_request_patch_summary = patch_summary.clone();
|
||||
|
||||
Ok(Some(Json(json!({
|
||||
"group": routing_group_payload(&group),
|
||||
"policy": policy,
|
||||
"trace_seed": trace,
|
||||
"patch_summary": patch_summary,
|
||||
"mutated_body": body,
|
||||
"mutated_headers": header_map_payload(&header_map),
|
||||
"candidate_preview": {
|
||||
"status": "policy_only",
|
||||
"ranking_overlay": policy.ranking_overlay,
|
||||
"note": "full candidate preview is produced by runtime materialization once provider/key catalogs are enumerated"
|
||||
}
|
||||
})).into_response()))
|
||||
}
|
||||
|
||||
fn build_routing_group_update_patch(
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<UpdateRoutingGroupRecord, GatewayError> {
|
||||
let raw = parse_json_value_body(request_body)?;
|
||||
let Some(object) = raw.as_object() else {
|
||||
return Err(bad_request_error("request body must be a JSON object"));
|
||||
};
|
||||
let mut patch = UpdateRoutingGroupRecord {
|
||||
updated_at: current_unix_secs() as i64,
|
||||
..UpdateRoutingGroupRecord::default()
|
||||
};
|
||||
if let Some(value) = object.get("name") {
|
||||
patch.name = Some(required_string(value, "name")?);
|
||||
}
|
||||
if let Some(value) = object.get("description") {
|
||||
patch.description = Some(optional_string(value, "description")?);
|
||||
}
|
||||
if let Some(value) = object.get("enabled") {
|
||||
patch.enabled = Some(required_bool(value, "enabled")?);
|
||||
}
|
||||
if let Some(value) = object.get("is_system_default") {
|
||||
patch.is_system_default = Some(required_bool(value, "is_system_default")?);
|
||||
}
|
||||
if let Some(value) = object.get("config_json") {
|
||||
validate_config_json(value)?;
|
||||
patch.config_json = Some(value.clone());
|
||||
patch.version = object
|
||||
.get("version")
|
||||
.and_then(Value::as_i64)
|
||||
.or(Some(current_unix_secs() as i64));
|
||||
} else if let Some(value) = object.get("version") {
|
||||
patch.version = Some(required_i64(value, "version")?.max(1));
|
||||
}
|
||||
if let Some(value) = object.get("published_at") {
|
||||
patch.published_at = Some(optional_i64(value, "published_at")?);
|
||||
}
|
||||
Ok(patch)
|
||||
}
|
||||
|
||||
fn build_routing_binding_update_patch(
|
||||
request_body: Option<&Bytes>,
|
||||
) -> Result<UpdateRoutingGroupBindingRecord, GatewayError> {
|
||||
let raw = parse_json_value_body(request_body)?;
|
||||
let Some(object) = raw.as_object() else {
|
||||
return Err(bad_request_error("request body must be a JSON object"));
|
||||
};
|
||||
let mut patch = UpdateRoutingGroupBindingRecord {
|
||||
updated_at: current_unix_secs() as i64,
|
||||
..UpdateRoutingGroupBindingRecord::default()
|
||||
};
|
||||
if let Some(value) = object.get("group_id") {
|
||||
patch.group_id = Some(required_string(value, "group_id")?);
|
||||
}
|
||||
if let Some(value) = object.get("subject_type") {
|
||||
patch.subject_type = Some(routing_subject_from_value(value)?);
|
||||
}
|
||||
if let Some(value) = object.get("subject_id") {
|
||||
patch.subject_id = Some(required_string(value, "subject_id")?);
|
||||
}
|
||||
if let Some(value) = object.get("is_default") {
|
||||
patch.is_default = Some(required_bool(value, "is_default")?);
|
||||
}
|
||||
if let Some(value) = object.get("allow_explicit_select") {
|
||||
patch.allow_explicit_select = Some(required_bool(value, "allow_explicit_select")?);
|
||||
}
|
||||
Ok(patch)
|
||||
}
|
||||
|
||||
fn routing_binding_query_from_request(
|
||||
request_context: &AdminRequestContext<'_>,
|
||||
) -> Result<RoutingGroupBindingQuery, GatewayError> {
|
||||
let subject_type = query_param_value(request_context.query_string(), "subject_type")
|
||||
.map(|value| routing_subject_from_str(&value))
|
||||
.transpose()?;
|
||||
Ok(RoutingGroupBindingQuery {
|
||||
group_id: query_param_value(request_context.query_string(), "group_id"),
|
||||
subject_type,
|
||||
subject_id: query_param_value(request_context.query_string(), "subject_id"),
|
||||
})
|
||||
}
|
||||
|
||||
fn validate_config_json(value: &Value) -> Result<(), GatewayError> {
|
||||
if !value.is_object() {
|
||||
return Err(bad_request_error("config_json must be a JSON object"));
|
||||
}
|
||||
let config = serde_json::from_value::<RoutingGroupConfig>(value.clone())
|
||||
.map_err(|err| bad_request_error(format!("config_json is invalid: {err}")))?;
|
||||
validate_routing_group_config(&config)
|
||||
.map_err(|err| bad_request_error(format!("config_json is invalid: {err}")))
|
||||
}
|
||||
|
||||
fn parse_json_body<T>(request_body: Option<&Bytes>) -> Result<T, GatewayError>
|
||||
where
|
||||
T: for<'de> Deserialize<'de>,
|
||||
{
|
||||
let raw = request_body.ok_or_else(|| bad_request_error("request body is required"))?;
|
||||
serde_json::from_slice(raw)
|
||||
.map_err(|err| bad_request_error(format!("request body must be valid JSON: {err}")))
|
||||
}
|
||||
|
||||
fn parse_json_value_body(request_body: Option<&Bytes>) -> Result<Value, GatewayError> {
|
||||
parse_json_body::<Value>(request_body)
|
||||
}
|
||||
|
||||
fn header_map_from_value(value: &Value) -> Result<HeaderMap, GatewayError> {
|
||||
let Some(object) = value.as_object() else {
|
||||
return Err(bad_request_error("headers must be a JSON object"));
|
||||
};
|
||||
let mut headers = HeaderMap::new();
|
||||
for (name, value) in object {
|
||||
let Some(value) = value.as_str() else {
|
||||
return Err(bad_request_error(format!(
|
||||
"header {name} must have a string value"
|
||||
)));
|
||||
};
|
||||
let header_name = HeaderName::from_bytes(name.as_bytes())
|
||||
.map_err(|_| bad_request_error(format!("header {name} has invalid name")))?;
|
||||
let header_value = HeaderValue::from_str(value)
|
||||
.map_err(|_| bad_request_error(format!("header {name} has invalid value")))?;
|
||||
headers.insert(header_name, header_value);
|
||||
}
|
||||
Ok(headers)
|
||||
}
|
||||
|
||||
fn header_map_payload(headers: &HeaderMap) -> BTreeMap<String, String> {
|
||||
headers
|
||||
.iter()
|
||||
.filter_map(|(name, value)| {
|
||||
value
|
||||
.to_str()
|
||||
.ok()
|
||||
.map(|value| (name.as_str().to_string(), value.to_string()))
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn patch_summary(plan: &MutationPlan) -> RoutingPatchSummary {
|
||||
RoutingPatchSummary {
|
||||
body_paths: plan
|
||||
.body_patch
|
||||
.iter()
|
||||
.map(|operation| operation.path().to_string())
|
||||
.collect(),
|
||||
header_names: plan
|
||||
.header_patch
|
||||
.iter()
|
||||
.map(|operation| match operation {
|
||||
RoutingHeaderPatch::Set { name, .. } | RoutingHeaderPatch::Remove { name } => {
|
||||
name.clone()
|
||||
}
|
||||
})
|
||||
.collect(),
|
||||
failed_action: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn routing_group_payload(group: &StoredRoutingGroup) -> Value {
|
||||
json!({
|
||||
"id": group.id,
|
||||
"name": group.name,
|
||||
"description": group.description,
|
||||
"enabled": group.enabled,
|
||||
"is_system_default": group.is_system_default,
|
||||
"config_json": group.config_json,
|
||||
"version": group.version,
|
||||
"created_at": group.created_at,
|
||||
"updated_at": group.updated_at,
|
||||
"published_at": group.published_at,
|
||||
})
|
||||
}
|
||||
|
||||
fn routing_group_binding_payload(binding: &StoredRoutingGroupBinding) -> Value {
|
||||
json!({
|
||||
"id": binding.id,
|
||||
"group_id": binding.group_id,
|
||||
"subject_type": binding.subject_type,
|
||||
"subject_id": binding.subject_id,
|
||||
"is_default": binding.is_default,
|
||||
"allow_explicit_select": binding.allow_explicit_select,
|
||||
"created_at": binding.created_at,
|
||||
"updated_at": binding.updated_at,
|
||||
})
|
||||
}
|
||||
|
||||
fn routing_group_version_payload(version: &StoredRoutingGroupVersion) -> Value {
|
||||
json!({
|
||||
"id": version.id,
|
||||
"group_id": version.group_id,
|
||||
"version": version.version,
|
||||
"config_json": version.config_json,
|
||||
"created_at": version.created_at,
|
||||
"created_by": version.created_by,
|
||||
})
|
||||
}
|
||||
|
||||
fn normalized_admin_path(path: &str) -> String {
|
||||
let trimmed = path.trim_end_matches('/');
|
||||
if trimmed.is_empty() {
|
||||
"/".to_string()
|
||||
} else {
|
||||
trimmed.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn routing_group_path_parts(path: &str) -> Option<(String, Option<String>)> {
|
||||
let suffix = path.strip_prefix(&(ROUTING_GROUPS_ROOT.to_string() + "/"))?;
|
||||
let mut parts = suffix.split('/');
|
||||
let group_id = parts.next()?.trim();
|
||||
if group_id.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let suffix = parts.next().map(str::to_string);
|
||||
if parts.next().is_some() {
|
||||
return None;
|
||||
}
|
||||
Some((group_id.to_string(), suffix))
|
||||
}
|
||||
|
||||
fn routing_binding_id_from_path(path: &str) -> Option<String> {
|
||||
let suffix = path.strip_prefix(&(ROUTING_BINDINGS_ROOT.to_string() + "/"))?;
|
||||
if suffix.trim().is_empty() || suffix.contains('/') {
|
||||
return None;
|
||||
}
|
||||
Some(suffix.to_string())
|
||||
}
|
||||
|
||||
fn routing_subject_from_value(value: &Value) -> Result<RoutingGroupBindingSubject, GatewayError> {
|
||||
let Some(value) = value.as_str() else {
|
||||
return Err(bad_request_error("subject_type must be a string"));
|
||||
};
|
||||
routing_subject_from_str(value)
|
||||
}
|
||||
|
||||
fn routing_subject_from_str(value: &str) -> Result<RoutingGroupBindingSubject, GatewayError> {
|
||||
match value.trim() {
|
||||
"user" => Ok(RoutingGroupBindingSubject::User),
|
||||
"api_key" => Ok(RoutingGroupBindingSubject::ApiKey),
|
||||
"user_group" => Ok(RoutingGroupBindingSubject::UserGroup),
|
||||
other => Err(bad_request_error(format!(
|
||||
"unsupported subject_type: {other}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn required_string(value: &Value, field: &str) -> Result<String, GatewayError> {
|
||||
value
|
||||
.as_str()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string)
|
||||
.ok_or_else(|| bad_request_error(format!("{field} must be a non-empty string")))
|
||||
}
|
||||
|
||||
fn optional_string(value: &Value, field: &str) -> Result<Option<String>, GatewayError> {
|
||||
if value.is_null() {
|
||||
return Ok(None);
|
||||
}
|
||||
required_string(value, field).map(Some)
|
||||
}
|
||||
|
||||
fn required_bool(value: &Value, field: &str) -> Result<bool, GatewayError> {
|
||||
value
|
||||
.as_bool()
|
||||
.ok_or_else(|| bad_request_error(format!("{field} must be a boolean")))
|
||||
}
|
||||
|
||||
fn required_i64(value: &Value, field: &str) -> Result<i64, GatewayError> {
|
||||
value
|
||||
.as_i64()
|
||||
.ok_or_else(|| bad_request_error(format!("{field} must be an integer")))
|
||||
}
|
||||
|
||||
fn optional_i64(value: &Value, field: &str) -> Result<Option<i64>, GatewayError> {
|
||||
if value.is_null() {
|
||||
return Ok(None);
|
||||
}
|
||||
required_i64(value, field).map(Some)
|
||||
}
|
||||
|
||||
fn default_true() -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
fn default_api_format() -> String {
|
||||
"openai:chat".to_string()
|
||||
}
|
||||
|
||||
fn bad_request_error(detail: impl Into<String>) -> GatewayError {
|
||||
GatewayError::Client {
|
||||
status: http::StatusCode::BAD_REQUEST,
|
||||
message: detail.into(),
|
||||
}
|
||||
}
|
||||
|
||||
fn bad_request_response(detail: impl Into<String>) -> Response<Body> {
|
||||
(
|
||||
http::StatusCode::BAD_REQUEST,
|
||||
Json(json!({ "detail": detail.into() })),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn not_found_response(detail: impl Into<String>) -> Response<Body> {
|
||||
(
|
||||
http::StatusCode::NOT_FOUND,
|
||||
Json(json!({ "detail": detail.into() })),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
fn data_unavailable_response() -> Response<Body> {
|
||||
(
|
||||
http::StatusCode::SERVICE_UNAVAILABLE,
|
||||
Json(json!({ "detail": "routing profile data backend is unavailable" })),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
@@ -300,6 +300,11 @@ pub(crate) fn admin_proxy_local_requires_buffered_body(
|
||||
http::Method::POST,
|
||||
Some("query_models" | "test_model" | "test_model_failover"),
|
||||
)
|
||||
| (Some("routing_profiles_manage"), http::Method::POST, Some("create_group"))
|
||||
| (Some("routing_profiles_manage"), http::Method::PATCH, Some("update_group"))
|
||||
| (Some("routing_profiles_manage"), http::Method::POST, Some("dry_run_group"))
|
||||
| (Some("routing_profiles_manage"), http::Method::POST, Some("create_binding"))
|
||||
| (Some("routing_profiles_manage"), http::Method::PATCH, Some("update_binding"))
|
||||
| (Some("billing_manage"), http::Method::POST, Some("apply_preset"))
|
||||
| (Some("billing_manage"), http::Method::POST, Some("create_rule"))
|
||||
| (Some("billing_manage"), http::Method::PUT, Some("update_rule"))
|
||||
|
||||
@@ -59,6 +59,7 @@ mod rate_limit;
|
||||
mod request_candidate_runtime;
|
||||
mod roles;
|
||||
mod router;
|
||||
mod routing;
|
||||
mod scheduler;
|
||||
mod state;
|
||||
mod system_features;
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
@@ -11,6 +11,7 @@ mod cors;
|
||||
mod integrations;
|
||||
mod oauth;
|
||||
mod proxy;
|
||||
mod routing_profiles;
|
||||
mod runtime;
|
||||
#[cfg(test)]
|
||||
mod testing;
|
||||
|
||||
163
apps/aether-gateway/src/state/routing_profiles.rs
Normal file
163
apps/aether-gateway/src/state/routing_profiles.rs
Normal file
@@ -0,0 +1,163 @@
|
||||
use aether_data_contracts::repository::routing_profiles::{
|
||||
CreateRoutingGroupBindingRecord, CreateRoutingGroupRecord, CreateRoutingGroupVersionRecord,
|
||||
RoutingGroupBindingQuery, RoutingGroupLookupKey, RoutingGroupReadRepository,
|
||||
StoredRoutingGroup, StoredRoutingGroupBinding, StoredRoutingGroupVersion,
|
||||
UpdateRoutingGroupBindingRecord, UpdateRoutingGroupRecord,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
use super::{AppState, GatewayError};
|
||||
|
||||
impl AppState {
|
||||
pub(crate) fn has_routing_group_data_reader(&self) -> bool {
|
||||
self.data.has_routing_group_reader()
|
||||
}
|
||||
|
||||
pub(crate) fn has_routing_group_data_writer(&self) -> bool {
|
||||
self.data.has_routing_group_writer()
|
||||
}
|
||||
|
||||
pub(crate) fn routing_group_read_repository(
|
||||
&self,
|
||||
) -> Option<Arc<dyn RoutingGroupReadRepository>> {
|
||||
self.data.routing_group_read_repository()
|
||||
}
|
||||
|
||||
pub(crate) async fn list_routing_groups(
|
||||
&self,
|
||||
) -> Result<Vec<StoredRoutingGroup>, GatewayError> {
|
||||
self.data
|
||||
.list_routing_groups()
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn find_routing_group(
|
||||
&self,
|
||||
lookup: RoutingGroupLookupKey<'_>,
|
||||
) -> Result<Option<StoredRoutingGroup>, GatewayError> {
|
||||
self.data
|
||||
.find_routing_group(lookup)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_routing_group_bindings(
|
||||
&self,
|
||||
query: &RoutingGroupBindingQuery,
|
||||
) -> Result<Vec<StoredRoutingGroupBinding>, GatewayError> {
|
||||
self.data
|
||||
.list_routing_group_bindings(query)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn list_routing_group_versions(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Vec<StoredRoutingGroupVersion>, GatewayError> {
|
||||
self.data
|
||||
.list_routing_group_versions(group_id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
|
||||
pub(crate) async fn create_routing_group(
|
||||
&self,
|
||||
record: CreateRoutingGroupRecord,
|
||||
) -> Result<Option<StoredRoutingGroup>, GatewayError> {
|
||||
let created = self
|
||||
.data
|
||||
.create_routing_group(record)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if created.is_some() {
|
||||
self.invalidate_provider_routing_caches();
|
||||
}
|
||||
Ok(created)
|
||||
}
|
||||
|
||||
pub(crate) async fn update_routing_group(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupRecord,
|
||||
) -> Result<Option<StoredRoutingGroup>, GatewayError> {
|
||||
let updated = self
|
||||
.data
|
||||
.update_routing_group(id, patch)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if updated.is_some() {
|
||||
self.invalidate_provider_routing_caches();
|
||||
}
|
||||
Ok(updated)
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_routing_group(&self, id: &str) -> Result<bool, GatewayError> {
|
||||
let deleted = self
|
||||
.data
|
||||
.delete_routing_group(id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if deleted {
|
||||
self.invalidate_provider_routing_caches();
|
||||
}
|
||||
Ok(deleted)
|
||||
}
|
||||
|
||||
pub(crate) async fn create_routing_group_binding(
|
||||
&self,
|
||||
record: CreateRoutingGroupBindingRecord,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, GatewayError> {
|
||||
let created = self
|
||||
.data
|
||||
.create_routing_group_binding(record)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if created.is_some() {
|
||||
self.invalidate_provider_routing_caches();
|
||||
}
|
||||
Ok(created)
|
||||
}
|
||||
|
||||
pub(crate) async fn update_routing_group_binding(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupBindingRecord,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, GatewayError> {
|
||||
let updated = self
|
||||
.data
|
||||
.update_routing_group_binding(id, patch)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if updated.is_some() {
|
||||
self.invalidate_provider_routing_caches();
|
||||
}
|
||||
Ok(updated)
|
||||
}
|
||||
|
||||
pub(crate) async fn delete_routing_group_binding(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<bool, GatewayError> {
|
||||
let deleted = self
|
||||
.data
|
||||
.delete_routing_group_binding(id)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if deleted {
|
||||
self.invalidate_provider_routing_caches();
|
||||
}
|
||||
Ok(deleted)
|
||||
}
|
||||
|
||||
pub(crate) async fn create_routing_group_version(
|
||||
&self,
|
||||
record: CreateRoutingGroupVersionRecord,
|
||||
) -> Result<Option<StoredRoutingGroupVersion>, GatewayError> {
|
||||
self.data
|
||||
.create_routing_group_version(record)
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||
}
|
||||
}
|
||||
@@ -3400,7 +3400,7 @@ fn ai_serving_decision_inputs_share_authenticated_input_helper() {
|
||||
),
|
||||
(
|
||||
"apps/aether-gateway/src/ai_serving/planner/specialized/files/support.rs",
|
||||
"LocalAuthenticatedDecisionInput as LocalGeminiFilesDecisionInput",
|
||||
"LocalRequestedModelDecisionInput as LocalGeminiFilesDecisionInput",
|
||||
),
|
||||
] {
|
||||
let source = read_workspace_file(path);
|
||||
@@ -3433,7 +3433,7 @@ fn ai_serving_decision_inputs_share_authenticated_input_helper() {
|
||||
),
|
||||
(
|
||||
"apps/aether-gateway/src/ai_serving/planner/specialized/files/support.rs",
|
||||
"build_local_authenticated_decision_input(",
|
||||
"build_local_requested_model_decision_input(",
|
||||
),
|
||||
] {
|
||||
let source = read_workspace_file(path);
|
||||
|
||||
@@ -8,6 +8,7 @@ description = "Shared data contracts and repository traits for Aether Rust servi
|
||||
|
||||
[dependencies]
|
||||
aether-ai-formats.workspace = true
|
||||
aether-routing-core.workspace = true
|
||||
async-trait.workspace = true
|
||||
chrono.workspace = true
|
||||
serde.workspace = true
|
||||
|
||||
@@ -6,6 +6,7 @@ pub mod global_models;
|
||||
pub mod pool_scores;
|
||||
pub mod provider_catalog;
|
||||
pub mod quota;
|
||||
pub mod routing_profiles;
|
||||
pub mod settlement;
|
||||
pub mod usage;
|
||||
pub mod video_tasks;
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
mod types;
|
||||
|
||||
pub use aether_routing_core::{RoutingGroupBindingSubject, RoutingGroupConfig, RoutingGroupRecord};
|
||||
pub use types::{
|
||||
CreateRoutingGroupBindingRecord, CreateRoutingGroupRecord, CreateRoutingGroupVersionRecord,
|
||||
RoutingGroupBindingQuery, RoutingGroupLookupKey, RoutingGroupReadRepository,
|
||||
RoutingGroupWriteRepository, StoredRoutingGroup, StoredRoutingGroupBinding,
|
||||
StoredRoutingGroupVersion, UpdateRoutingGroupBindingRecord, UpdateRoutingGroupRecord,
|
||||
};
|
||||
@@ -0,0 +1,238 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
use aether_routing_core::RoutingGroupBindingSubject;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct StoredRoutingGroup {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
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,
|
||||
pub published_at: Option<i64>,
|
||||
}
|
||||
|
||||
impl StoredRoutingGroup {
|
||||
pub fn new(record: CreateRoutingGroupRecord) -> Result<Self, crate::DataLayerError> {
|
||||
validate_non_empty(&record.id, "routing_groups.id")?;
|
||||
validate_non_empty(&record.name, "routing_groups.name")?;
|
||||
validate_config_object(&record.config_json, "routing_groups.config_json")?;
|
||||
Ok(Self {
|
||||
id: record.id,
|
||||
name: record.name,
|
||||
description: record.description,
|
||||
enabled: record.enabled,
|
||||
is_system_default: record.is_system_default,
|
||||
config_json: record.config_json,
|
||||
version: record.version.max(1),
|
||||
created_at: record.created_at,
|
||||
updated_at: record.updated_at,
|
||||
published_at: record.published_at,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CreateRoutingGroupRecord {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
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,
|
||||
pub published_at: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct UpdateRoutingGroupRecord {
|
||||
pub name: Option<String>,
|
||||
pub description: Option<Option<String>>,
|
||||
pub enabled: Option<bool>,
|
||||
pub is_system_default: Option<bool>,
|
||||
pub config_json: Option<Value>,
|
||||
pub version: Option<i64>,
|
||||
pub updated_at: i64,
|
||||
pub published_at: Option<Option<i64>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct UpdateRoutingGroupBindingRecord {
|
||||
pub group_id: Option<String>,
|
||||
pub subject_type: Option<RoutingGroupBindingSubject>,
|
||||
pub subject_id: Option<String>,
|
||||
pub is_default: Option<bool>,
|
||||
pub allow_explicit_select: Option<bool>,
|
||||
pub updated_at: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct StoredRoutingGroupBinding {
|
||||
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,
|
||||
}
|
||||
|
||||
impl StoredRoutingGroupBinding {
|
||||
pub fn new(record: CreateRoutingGroupBindingRecord) -> Result<Self, crate::DataLayerError> {
|
||||
validate_non_empty(&record.id, "routing_group_bindings.id")?;
|
||||
validate_non_empty(&record.group_id, "routing_group_bindings.group_id")?;
|
||||
validate_non_empty(&record.subject_id, "routing_group_bindings.subject_id")?;
|
||||
Ok(Self {
|
||||
id: record.id,
|
||||
group_id: record.group_id,
|
||||
subject_type: record.subject_type,
|
||||
subject_id: record.subject_id,
|
||||
is_default: record.is_default,
|
||||
allow_explicit_select: record.allow_explicit_select,
|
||||
created_at: record.created_at,
|
||||
updated_at: record.updated_at,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CreateRoutingGroupBindingRecord {
|
||||
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)]
|
||||
pub struct StoredRoutingGroupVersion {
|
||||
pub id: String,
|
||||
pub group_id: String,
|
||||
pub version: i64,
|
||||
pub config_json: Value,
|
||||
pub created_at: i64,
|
||||
pub created_by: Option<String>,
|
||||
}
|
||||
|
||||
impl StoredRoutingGroupVersion {
|
||||
pub fn new(record: CreateRoutingGroupVersionRecord) -> Result<Self, crate::DataLayerError> {
|
||||
validate_non_empty(&record.id, "routing_group_versions.id")?;
|
||||
validate_non_empty(&record.group_id, "routing_group_versions.group_id")?;
|
||||
validate_config_object(&record.config_json, "routing_group_versions.config_json")?;
|
||||
Ok(Self {
|
||||
id: record.id,
|
||||
group_id: record.group_id,
|
||||
version: record.version.max(1),
|
||||
config_json: record.config_json,
|
||||
created_at: record.created_at,
|
||||
created_by: record.created_by,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct CreateRoutingGroupVersionRecord {
|
||||
pub id: String,
|
||||
pub group_id: String,
|
||||
pub version: i64,
|
||||
pub config_json: Value,
|
||||
pub created_at: i64,
|
||||
pub created_by: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum RoutingGroupLookupKey<'a> {
|
||||
Id(&'a str),
|
||||
Name(&'a str),
|
||||
SystemDefault,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
pub struct RoutingGroupBindingQuery {
|
||||
pub group_id: Option<String>,
|
||||
pub subject_type: Option<RoutingGroupBindingSubject>,
|
||||
pub subject_id: Option<String>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait RoutingGroupReadRepository: Send + Sync {
|
||||
async fn list_routing_groups(&self) -> Result<Vec<StoredRoutingGroup>, crate::DataLayerError>;
|
||||
|
||||
async fn find_routing_group(
|
||||
&self,
|
||||
lookup: RoutingGroupLookupKey<'_>,
|
||||
) -> Result<Option<StoredRoutingGroup>, crate::DataLayerError>;
|
||||
|
||||
async fn list_routing_group_bindings(
|
||||
&self,
|
||||
query: &RoutingGroupBindingQuery,
|
||||
) -> Result<Vec<StoredRoutingGroupBinding>, crate::DataLayerError>;
|
||||
|
||||
async fn list_routing_group_versions(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Vec<StoredRoutingGroupVersion>, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait RoutingGroupWriteRepository: Send + Sync {
|
||||
async fn create_routing_group(
|
||||
&self,
|
||||
record: CreateRoutingGroupRecord,
|
||||
) -> Result<StoredRoutingGroup, crate::DataLayerError>;
|
||||
|
||||
async fn update_routing_group(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupRecord,
|
||||
) -> Result<Option<StoredRoutingGroup>, crate::DataLayerError>;
|
||||
|
||||
async fn delete_routing_group(&self, id: &str) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn create_routing_group_binding(
|
||||
&self,
|
||||
record: CreateRoutingGroupBindingRecord,
|
||||
) -> Result<StoredRoutingGroupBinding, crate::DataLayerError>;
|
||||
|
||||
async fn delete_routing_group_binding(&self, id: &str) -> Result<bool, crate::DataLayerError>;
|
||||
|
||||
async fn update_routing_group_binding(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupBindingRecord,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, crate::DataLayerError>;
|
||||
|
||||
async fn create_routing_group_version(
|
||||
&self,
|
||||
record: CreateRoutingGroupVersionRecord,
|
||||
) -> Result<StoredRoutingGroupVersion, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
fn validate_non_empty(value: &str, field: &str) -> Result<(), crate::DataLayerError> {
|
||||
if value.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(format!(
|
||||
"{field} is empty"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn validate_config_object(value: &Value, field: &str) -> Result<(), crate::DataLayerError> {
|
||||
if !value.is_object() {
|
||||
return Err(crate::DataLayerError::InvalidInput(format!(
|
||||
"{field} must be a JSON object"
|
||||
)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,41 @@
|
||||
CREATE TABLE IF NOT EXISTS routing_groups (
|
||||
`id` VARCHAR(64) NOT NULL,
|
||||
`name` VARCHAR(255) NOT NULL,
|
||||
`description` LONGTEXT,
|
||||
`enabled` TINYINT(1) NOT NULL DEFAULT 1,
|
||||
`is_system_default` TINYINT(1) NOT NULL DEFAULT 0,
|
||||
`config_json` JSON NOT NULL,
|
||||
`version` BIGINT NOT NULL DEFAULT 1,
|
||||
`created_at` BIGINT NOT NULL,
|
||||
`updated_at` BIGINT NOT NULL,
|
||||
`published_at` BIGINT,
|
||||
PRIMARY KEY (`id`),
|
||||
UNIQUE KEY routing_groups_name_key (`name`),
|
||||
KEY routing_groups_system_default_idx (`is_system_default`, `enabled`)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routing_group_bindings (
|
||||
`id` VARCHAR(64) NOT NULL,
|
||||
`group_id` VARCHAR(64) NOT NULL,
|
||||
`subject_type` VARCHAR(32) NOT NULL,
|
||||
`subject_id` VARCHAR(64) NOT NULL,
|
||||
`is_default` TINYINT(1) NOT NULL DEFAULT 0,
|
||||
`allow_explicit_select` TINYINT(1) NOT NULL DEFAULT 1,
|
||||
`created_at` BIGINT NOT NULL,
|
||||
`updated_at` BIGINT NOT NULL,
|
||||
PRIMARY KEY (`id`),
|
||||
KEY routing_group_bindings_group_id_idx (`group_id`),
|
||||
KEY routing_group_bindings_subject_idx (`subject_type`, `subject_id`)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routing_group_versions (
|
||||
`id` VARCHAR(64) NOT NULL,
|
||||
`group_id` VARCHAR(64) NOT NULL,
|
||||
`version` BIGINT NOT NULL,
|
||||
`config_json` JSON NOT NULL,
|
||||
`created_at` BIGINT NOT NULL,
|
||||
`created_by` VARCHAR(64),
|
||||
PRIMARY KEY (`id`),
|
||||
UNIQUE KEY routing_group_versions_group_version_key (`group_id`, `version`),
|
||||
KEY routing_group_versions_group_id_idx (`group_id`)
|
||||
);
|
||||
@@ -0,0 +1,82 @@
|
||||
CREATE TABLE IF NOT EXISTS public.routing_groups (
|
||||
id character varying(64) NOT NULL,
|
||||
name character varying(255) NOT NULL,
|
||||
description text,
|
||||
enabled boolean DEFAULT true NOT NULL,
|
||||
is_system_default boolean DEFAULT false NOT NULL,
|
||||
config_json jsonb NOT NULL,
|
||||
version bigint DEFAULT 1 NOT NULL,
|
||||
created_at bigint NOT NULL,
|
||||
updated_at bigint NOT NULL,
|
||||
published_at bigint
|
||||
);
|
||||
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint WHERE conname = 'routing_groups_pkey'
|
||||
) THEN
|
||||
ALTER TABLE ONLY public.routing_groups
|
||||
ADD CONSTRAINT routing_groups_pkey PRIMARY KEY (id);
|
||||
END IF;
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint WHERE conname = 'routing_groups_name_key'
|
||||
) THEN
|
||||
ALTER TABLE ONLY public.routing_groups
|
||||
ADD CONSTRAINT routing_groups_name_key UNIQUE (name);
|
||||
END IF;
|
||||
END $$;
|
||||
CREATE INDEX IF NOT EXISTS routing_groups_system_default_idx
|
||||
ON public.routing_groups USING btree (is_system_default, enabled);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.routing_group_bindings (
|
||||
id character varying(64) NOT NULL,
|
||||
group_id character varying(64) NOT NULL,
|
||||
subject_type character varying(32) NOT NULL,
|
||||
subject_id character varying(64) NOT NULL,
|
||||
is_default boolean DEFAULT false NOT NULL,
|
||||
allow_explicit_select boolean DEFAULT true NOT NULL,
|
||||
created_at bigint NOT NULL,
|
||||
updated_at bigint NOT NULL
|
||||
);
|
||||
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint WHERE conname = 'routing_group_bindings_pkey'
|
||||
) THEN
|
||||
ALTER TABLE ONLY public.routing_group_bindings
|
||||
ADD CONSTRAINT routing_group_bindings_pkey PRIMARY KEY (id);
|
||||
END IF;
|
||||
END $$;
|
||||
CREATE INDEX IF NOT EXISTS routing_group_bindings_group_id_idx
|
||||
ON public.routing_group_bindings USING btree (group_id);
|
||||
CREATE INDEX IF NOT EXISTS routing_group_bindings_subject_idx
|
||||
ON public.routing_group_bindings USING btree (subject_type, subject_id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.routing_group_versions (
|
||||
id character varying(64) NOT NULL,
|
||||
group_id character varying(64) NOT NULL,
|
||||
version bigint NOT NULL,
|
||||
config_json jsonb NOT NULL,
|
||||
created_at bigint NOT NULL,
|
||||
created_by character varying(64)
|
||||
);
|
||||
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint WHERE conname = 'routing_group_versions_pkey'
|
||||
) THEN
|
||||
ALTER TABLE ONLY public.routing_group_versions
|
||||
ADD CONSTRAINT routing_group_versions_pkey PRIMARY KEY (id);
|
||||
END IF;
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint WHERE conname = 'routing_group_versions_group_version_key'
|
||||
) THEN
|
||||
ALTER TABLE ONLY public.routing_group_versions
|
||||
ADD CONSTRAINT routing_group_versions_group_version_key UNIQUE (group_id, version);
|
||||
END IF;
|
||||
END $$;
|
||||
CREATE INDEX IF NOT EXISTS routing_group_versions_group_id_idx
|
||||
ON public.routing_group_versions USING btree (group_id);
|
||||
@@ -0,0 +1,45 @@
|
||||
CREATE TABLE IF NOT EXISTS routing_groups (
|
||||
id TEXT PRIMARY KEY NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
description TEXT,
|
||||
enabled INTEGER NOT NULL DEFAULT 1,
|
||||
is_system_default INTEGER NOT NULL DEFAULT 0,
|
||||
config_json TEXT NOT NULL,
|
||||
version INTEGER NOT NULL DEFAULT 1,
|
||||
created_at INTEGER NOT NULL,
|
||||
updated_at INTEGER NOT NULL,
|
||||
published_at INTEGER,
|
||||
UNIQUE (name)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS routing_groups_system_default_idx
|
||||
ON routing_groups (is_system_default, enabled);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routing_group_bindings (
|
||||
id TEXT PRIMARY KEY NOT NULL,
|
||||
group_id TEXT NOT NULL,
|
||||
subject_type TEXT NOT NULL,
|
||||
subject_id TEXT NOT NULL,
|
||||
is_default INTEGER NOT NULL DEFAULT 0,
|
||||
allow_explicit_select INTEGER NOT NULL DEFAULT 1,
|
||||
created_at INTEGER NOT NULL,
|
||||
updated_at INTEGER NOT NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS routing_group_bindings_group_id_idx
|
||||
ON routing_group_bindings (group_id);
|
||||
CREATE INDEX IF NOT EXISTS routing_group_bindings_subject_idx
|
||||
ON routing_group_bindings (subject_type, subject_id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routing_group_versions (
|
||||
id TEXT PRIMARY KEY NOT NULL,
|
||||
group_id TEXT NOT NULL,
|
||||
version INTEGER NOT NULL,
|
||||
config_json TEXT NOT NULL,
|
||||
created_at INTEGER NOT NULL,
|
||||
created_by TEXT,
|
||||
UNIQUE (group_id, version)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS routing_group_versions_group_id_idx
|
||||
ON routing_group_versions (group_id);
|
||||
@@ -383,3 +383,45 @@ CREATE TABLE IF NOT EXISTS global_models (
|
||||
UNIQUE KEY global_models_name_key (`name`)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routing_groups (
|
||||
`id` VARCHAR(64) NOT NULL,
|
||||
`name` VARCHAR(255) NOT NULL,
|
||||
`description` LONGTEXT,
|
||||
`enabled` TINYINT(1) NOT NULL DEFAULT 1,
|
||||
`is_system_default` TINYINT(1) NOT NULL DEFAULT 0,
|
||||
`config_json` JSON NOT NULL,
|
||||
`version` BIGINT NOT NULL DEFAULT 1,
|
||||
`created_at` BIGINT NOT NULL,
|
||||
`updated_at` BIGINT NOT NULL,
|
||||
`published_at` BIGINT,
|
||||
PRIMARY KEY (`id`),
|
||||
UNIQUE KEY routing_groups_name_key (`name`),
|
||||
KEY routing_groups_system_default_idx (`is_system_default`, `enabled`)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routing_group_bindings (
|
||||
`id` VARCHAR(64) NOT NULL,
|
||||
`group_id` VARCHAR(64) NOT NULL,
|
||||
`subject_type` VARCHAR(32) NOT NULL,
|
||||
`subject_id` VARCHAR(64) NOT NULL,
|
||||
`is_default` TINYINT(1) NOT NULL DEFAULT 0,
|
||||
`allow_explicit_select` TINYINT(1) NOT NULL DEFAULT 1,
|
||||
`created_at` BIGINT NOT NULL,
|
||||
`updated_at` BIGINT NOT NULL,
|
||||
PRIMARY KEY (`id`),
|
||||
KEY routing_group_bindings_group_id_idx (`group_id`),
|
||||
KEY routing_group_bindings_subject_idx (`subject_type`, `subject_id`)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routing_group_versions (
|
||||
`id` VARCHAR(64) NOT NULL,
|
||||
`group_id` VARCHAR(64) NOT NULL,
|
||||
`version` BIGINT NOT NULL,
|
||||
`config_json` JSON NOT NULL,
|
||||
`created_at` BIGINT NOT NULL,
|
||||
`created_by` VARCHAR(64),
|
||||
PRIMARY KEY (`id`),
|
||||
UNIQUE KEY routing_group_versions_group_version_key (`group_id`, `version`),
|
||||
KEY routing_group_versions_group_id_idx (`group_id`)
|
||||
);
|
||||
|
||||
|
||||
@@ -396,3 +396,48 @@ CREATE TABLE IF NOT EXISTS public.global_models (
|
||||
ALTER TABLE ONLY public.global_models ADD CONSTRAINT global_models_pkey PRIMARY KEY (id);
|
||||
ALTER TABLE ONLY public.global_models ADD CONSTRAINT global_models_name_key UNIQUE (name);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.routing_groups (
|
||||
id character varying(64) NOT NULL,
|
||||
name character varying(255) NOT NULL,
|
||||
description text,
|
||||
enabled boolean DEFAULT true NOT NULL,
|
||||
is_system_default boolean DEFAULT false NOT NULL,
|
||||
config_json jsonb NOT NULL,
|
||||
version bigint DEFAULT 1 NOT NULL,
|
||||
created_at bigint NOT NULL,
|
||||
updated_at bigint NOT NULL,
|
||||
published_at bigint
|
||||
);
|
||||
|
||||
ALTER TABLE ONLY public.routing_groups ADD CONSTRAINT routing_groups_pkey PRIMARY KEY (id);
|
||||
ALTER TABLE ONLY public.routing_groups ADD CONSTRAINT routing_groups_name_key UNIQUE (name);
|
||||
CREATE INDEX IF NOT EXISTS routing_groups_system_default_idx ON public.routing_groups USING btree (is_system_default, enabled);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.routing_group_bindings (
|
||||
id character varying(64) NOT NULL,
|
||||
group_id character varying(64) NOT NULL,
|
||||
subject_type character varying(32) NOT NULL,
|
||||
subject_id character varying(64) NOT NULL,
|
||||
is_default boolean DEFAULT false NOT NULL,
|
||||
allow_explicit_select boolean DEFAULT true NOT NULL,
|
||||
created_at bigint NOT NULL,
|
||||
updated_at bigint NOT NULL
|
||||
);
|
||||
|
||||
ALTER TABLE ONLY public.routing_group_bindings ADD CONSTRAINT routing_group_bindings_pkey PRIMARY KEY (id);
|
||||
CREATE INDEX IF NOT EXISTS routing_group_bindings_group_id_idx ON public.routing_group_bindings USING btree (group_id);
|
||||
CREATE INDEX IF NOT EXISTS routing_group_bindings_subject_idx ON public.routing_group_bindings USING btree (subject_type, subject_id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.routing_group_versions (
|
||||
id character varying(64) NOT NULL,
|
||||
group_id character varying(64) NOT NULL,
|
||||
version bigint NOT NULL,
|
||||
config_json jsonb NOT NULL,
|
||||
created_at bigint NOT NULL,
|
||||
created_by character varying(64)
|
||||
);
|
||||
|
||||
ALTER TABLE ONLY public.routing_group_versions ADD CONSTRAINT routing_group_versions_pkey PRIMARY KEY (id);
|
||||
ALTER TABLE ONLY public.routing_group_versions ADD CONSTRAINT routing_group_versions_group_version_key UNIQUE (group_id, version);
|
||||
CREATE INDEX IF NOT EXISTS routing_group_versions_group_id_idx ON public.routing_group_versions USING btree (group_id);
|
||||
|
||||
|
||||
@@ -370,3 +370,42 @@ CREATE TABLE IF NOT EXISTS global_models (
|
||||
UNIQUE (name)
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routing_groups (
|
||||
id TEXT PRIMARY KEY NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
description TEXT,
|
||||
enabled INTEGER NOT NULL DEFAULT 1,
|
||||
is_system_default INTEGER NOT NULL DEFAULT 0,
|
||||
config_json TEXT NOT NULL,
|
||||
version INTEGER NOT NULL DEFAULT 1,
|
||||
created_at INTEGER NOT NULL,
|
||||
updated_at INTEGER NOT NULL,
|
||||
published_at INTEGER,
|
||||
UNIQUE (name)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS routing_groups_system_default_idx ON routing_groups (is_system_default, enabled);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routing_group_bindings (
|
||||
id TEXT PRIMARY KEY NOT NULL,
|
||||
group_id TEXT NOT NULL,
|
||||
subject_type TEXT NOT NULL,
|
||||
subject_id TEXT NOT NULL,
|
||||
is_default INTEGER NOT NULL DEFAULT 0,
|
||||
allow_explicit_select INTEGER NOT NULL DEFAULT 1,
|
||||
created_at INTEGER NOT NULL,
|
||||
updated_at INTEGER NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS routing_group_bindings_group_id_idx ON routing_group_bindings (group_id);
|
||||
CREATE INDEX IF NOT EXISTS routing_group_bindings_subject_idx ON routing_group_bindings (subject_type, subject_id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS routing_group_versions (
|
||||
id TEXT PRIMARY KEY NOT NULL,
|
||||
group_id TEXT NOT NULL,
|
||||
version INTEGER NOT NULL,
|
||||
config_json TEXT NOT NULL,
|
||||
created_at INTEGER NOT NULL,
|
||||
created_by TEXT,
|
||||
UNIQUE (group_id, version)
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS routing_group_versions_group_id_idx ON routing_group_versions (group_id);
|
||||
|
||||
|
||||
@@ -1694,3 +1694,155 @@ type = "unix_seconds"
|
||||
[[table.global_models.uniques]]
|
||||
name = "global_models_name_key"
|
||||
columns = ["name"]
|
||||
|
||||
[table.routing_groups]
|
||||
domain = "provider_catalog"
|
||||
order = 110
|
||||
primary_key = ["id"]
|
||||
|
||||
[[table.routing_groups.columns]]
|
||||
name = "id"
|
||||
type = "text_id"
|
||||
length = 64
|
||||
|
||||
[[table.routing_groups.columns]]
|
||||
name = "name"
|
||||
type = "text"
|
||||
length = 255
|
||||
|
||||
[[table.routing_groups.columns]]
|
||||
name = "description"
|
||||
type = "long_text"
|
||||
nullable = true
|
||||
|
||||
[[table.routing_groups.columns]]
|
||||
name = "enabled"
|
||||
type = "bool"
|
||||
default = true
|
||||
|
||||
[[table.routing_groups.columns]]
|
||||
name = "is_system_default"
|
||||
type = "bool"
|
||||
default = false
|
||||
|
||||
[[table.routing_groups.columns]]
|
||||
name = "config_json"
|
||||
type = "json"
|
||||
|
||||
[[table.routing_groups.columns]]
|
||||
name = "version"
|
||||
type = "int64"
|
||||
default = 1
|
||||
|
||||
[[table.routing_groups.columns]]
|
||||
name = "created_at"
|
||||
type = "unix_seconds"
|
||||
|
||||
[[table.routing_groups.columns]]
|
||||
name = "updated_at"
|
||||
type = "unix_seconds"
|
||||
|
||||
[[table.routing_groups.columns]]
|
||||
name = "published_at"
|
||||
type = "unix_seconds"
|
||||
nullable = true
|
||||
|
||||
[[table.routing_groups.uniques]]
|
||||
name = "routing_groups_name_key"
|
||||
columns = ["name"]
|
||||
|
||||
[[table.routing_groups.indexes]]
|
||||
name = "routing_groups_system_default_idx"
|
||||
columns = ["is_system_default", "enabled"]
|
||||
|
||||
[table.routing_group_bindings]
|
||||
domain = "provider_catalog"
|
||||
order = 111
|
||||
primary_key = ["id"]
|
||||
|
||||
[[table.routing_group_bindings.columns]]
|
||||
name = "id"
|
||||
type = "text_id"
|
||||
length = 64
|
||||
|
||||
[[table.routing_group_bindings.columns]]
|
||||
name = "group_id"
|
||||
type = "text_id"
|
||||
length = 64
|
||||
|
||||
[[table.routing_group_bindings.columns]]
|
||||
name = "subject_type"
|
||||
type = "text"
|
||||
length = 32
|
||||
|
||||
[[table.routing_group_bindings.columns]]
|
||||
name = "subject_id"
|
||||
type = "text_id"
|
||||
length = 64
|
||||
|
||||
[[table.routing_group_bindings.columns]]
|
||||
name = "is_default"
|
||||
type = "bool"
|
||||
default = false
|
||||
|
||||
[[table.routing_group_bindings.columns]]
|
||||
name = "allow_explicit_select"
|
||||
type = "bool"
|
||||
default = true
|
||||
|
||||
[[table.routing_group_bindings.columns]]
|
||||
name = "created_at"
|
||||
type = "unix_seconds"
|
||||
|
||||
[[table.routing_group_bindings.columns]]
|
||||
name = "updated_at"
|
||||
type = "unix_seconds"
|
||||
|
||||
[[table.routing_group_bindings.indexes]]
|
||||
name = "routing_group_bindings_group_id_idx"
|
||||
columns = ["group_id"]
|
||||
|
||||
[[table.routing_group_bindings.indexes]]
|
||||
name = "routing_group_bindings_subject_idx"
|
||||
columns = ["subject_type", "subject_id"]
|
||||
|
||||
[table.routing_group_versions]
|
||||
domain = "provider_catalog"
|
||||
order = 112
|
||||
primary_key = ["id"]
|
||||
|
||||
[[table.routing_group_versions.columns]]
|
||||
name = "id"
|
||||
type = "text_id"
|
||||
length = 64
|
||||
|
||||
[[table.routing_group_versions.columns]]
|
||||
name = "group_id"
|
||||
type = "text_id"
|
||||
length = 64
|
||||
|
||||
[[table.routing_group_versions.columns]]
|
||||
name = "version"
|
||||
type = "int64"
|
||||
|
||||
[[table.routing_group_versions.columns]]
|
||||
name = "config_json"
|
||||
type = "json"
|
||||
|
||||
[[table.routing_group_versions.columns]]
|
||||
name = "created_at"
|
||||
type = "unix_seconds"
|
||||
|
||||
[[table.routing_group_versions.columns]]
|
||||
name = "created_by"
|
||||
type = "text_id"
|
||||
length = 64
|
||||
nullable = true
|
||||
|
||||
[[table.routing_group_versions.uniques]]
|
||||
name = "routing_group_versions_group_version_key"
|
||||
columns = ["group_id", "version"]
|
||||
|
||||
[[table.routing_group_versions.indexes]]
|
||||
name = "routing_group_versions_group_id_idx"
|
||||
columns = ["group_id"]
|
||||
|
||||
@@ -50,6 +50,9 @@ use crate::repository::proxy_nodes::{
|
||||
use crate::repository::quota::{
|
||||
MysqlProviderQuotaRepository, ProviderQuotaReadRepository, ProviderQuotaWriteRepository,
|
||||
};
|
||||
use crate::repository::routing_profiles::{
|
||||
MysqlRoutingGroupRepository, RoutingGroupReadRepository, RoutingGroupWriteRepository,
|
||||
};
|
||||
use crate::repository::settlement::{MysqlSettlementRepository, SettlementWriteRepository};
|
||||
use crate::repository::usage::{
|
||||
MysqlUsageReadRepository, MysqlUsageWriteRepository, UsageReadRepository, UsageWriteRepository,
|
||||
@@ -195,6 +198,14 @@ impl MysqlBackend {
|
||||
Arc::new(MysqlPoolMemberScoreRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn routing_group_read_repository(&self) -> Arc<dyn RoutingGroupReadRepository> {
|
||||
Arc::new(MysqlRoutingGroupRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn routing_group_write_repository(&self) -> Arc<dyn RoutingGroupWriteRepository> {
|
||||
Arc::new(MysqlRoutingGroupRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
|
||||
Arc::new(MysqlProxyNodeReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
@@ -52,6 +52,9 @@ use crate::repository::proxy_nodes::{
|
||||
use crate::repository::quota::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, SqlxProviderQuotaRepository,
|
||||
};
|
||||
use crate::repository::routing_profiles::{
|
||||
PostgresRoutingGroupRepository, RoutingGroupReadRepository, RoutingGroupWriteRepository,
|
||||
};
|
||||
use crate::repository::settlement::{SettlementWriteRepository, SqlxSettlementRepository};
|
||||
use crate::repository::usage::{
|
||||
SqlxUsageReadRepository, UsageReadRepository, UsageWriteRepository,
|
||||
@@ -206,6 +209,14 @@ impl PostgresBackend {
|
||||
Arc::new(PostgresPoolMemberScoreRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn routing_group_read_repository(&self) -> Arc<dyn RoutingGroupReadRepository> {
|
||||
Arc::new(PostgresRoutingGroupRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn routing_group_write_repository(&self) -> Arc<dyn RoutingGroupWriteRepository> {
|
||||
Arc::new(PostgresRoutingGroupRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn provider_quota_read_repository(&self) -> Arc<dyn ProviderQuotaReadRepository> {
|
||||
Arc::new(SqlxProviderQuotaRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
@@ -18,6 +18,7 @@ use crate::repository::pool_scores::PoolScoreReadRepository;
|
||||
use crate::repository::provider_catalog::ProviderCatalogReadRepository;
|
||||
use crate::repository::proxy_nodes::ProxyNodeReadRepository;
|
||||
use crate::repository::quota::ProviderQuotaReadRepository;
|
||||
use crate::repository::routing_profiles::RoutingGroupReadRepository;
|
||||
use crate::repository::usage::UsageReadRepository;
|
||||
use crate::repository::users::UserReadRepository;
|
||||
use crate::repository::video_tasks::VideoTaskReadRepository;
|
||||
@@ -41,6 +42,7 @@ pub struct DataReadRepositories {
|
||||
request_candidates: Option<Arc<dyn RequestCandidateReadRepository>>,
|
||||
provider_catalog: Option<Arc<dyn ProviderCatalogReadRepository>>,
|
||||
provider_quotas: Option<Arc<dyn ProviderQuotaReadRepository>>,
|
||||
routing_groups: Option<Arc<dyn RoutingGroupReadRepository>>,
|
||||
usage: Option<Arc<dyn UsageReadRepository>>,
|
||||
users: Option<Arc<dyn UserReadRepository>>,
|
||||
video_tasks: Option<Arc<dyn VideoTaskReadRepository>>,
|
||||
@@ -72,6 +74,7 @@ impl fmt::Debug for DataReadRepositories {
|
||||
.field("has_request_candidates", &self.request_candidates.is_some())
|
||||
.field("has_provider_catalog", &self.provider_catalog.is_some())
|
||||
.field("has_provider_quotas", &self.provider_quotas.is_some())
|
||||
.field("has_routing_groups", &self.routing_groups.is_some())
|
||||
.field("has_usage", &self.usage.is_some())
|
||||
.field("has_users", &self.users.is_some())
|
||||
.field("has_video_tasks", &self.video_tasks.is_some())
|
||||
@@ -151,6 +154,10 @@ impl DataReadRepositories {
|
||||
.map(PostgresBackend::provider_quota_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::provider_quota_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::provider_quota_read_repository)),
|
||||
routing_groups: postgres
|
||||
.map(PostgresBackend::routing_group_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::routing_group_read_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::routing_group_read_repository)),
|
||||
usage: postgres
|
||||
.map(PostgresBackend::usage_read_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::usage_read_repository))
|
||||
@@ -241,6 +248,10 @@ impl DataReadRepositories {
|
||||
self.provider_quotas.clone()
|
||||
}
|
||||
|
||||
pub fn routing_groups(&self) -> Option<Arc<dyn RoutingGroupReadRepository>> {
|
||||
self.routing_groups.clone()
|
||||
}
|
||||
|
||||
pub fn usage(&self) -> Option<Arc<dyn UsageReadRepository>> {
|
||||
self.usage.clone()
|
||||
}
|
||||
@@ -274,6 +285,7 @@ impl DataReadRepositories {
|
||||
|| self.request_candidates.is_some()
|
||||
|| self.provider_catalog.is_some()
|
||||
|| self.provider_quotas.is_some()
|
||||
|| self.routing_groups.is_some()
|
||||
|| self.usage.is_some()
|
||||
|| self.users.is_some()
|
||||
|| self.video_tasks.is_some()
|
||||
|
||||
@@ -50,6 +50,9 @@ use crate::repository::proxy_nodes::{
|
||||
use crate::repository::quota::{
|
||||
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, SqliteProviderQuotaRepository,
|
||||
};
|
||||
use crate::repository::routing_profiles::{
|
||||
RoutingGroupReadRepository, RoutingGroupWriteRepository, SqliteRoutingGroupRepository,
|
||||
};
|
||||
use crate::repository::settlement::{SettlementWriteRepository, SqliteSettlementRepository};
|
||||
use crate::repository::usage::{
|
||||
SqliteUsageReadRepository, SqliteUsageWriteRepository, UsageReadRepository,
|
||||
@@ -208,6 +211,14 @@ impl SqliteBackend {
|
||||
Arc::new(SqlitePoolMemberScoreRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn routing_group_read_repository(&self) -> Arc<dyn RoutingGroupReadRepository> {
|
||||
Arc::new(SqliteRoutingGroupRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn routing_group_write_repository(&self) -> Arc<dyn RoutingGroupWriteRepository> {
|
||||
Arc::new(SqliteRoutingGroupRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
pub fn proxy_node_read_repository(&self) -> Arc<dyn ProxyNodeReadRepository> {
|
||||
Arc::new(SqliteProxyNodeReadRepository::new(self.pool_clone()))
|
||||
}
|
||||
|
||||
@@ -15,6 +15,7 @@ use crate::repository::pool_scores::PoolMemberScoreWriteRepository;
|
||||
use crate::repository::provider_catalog::ProviderCatalogWriteRepository;
|
||||
use crate::repository::proxy_nodes::ProxyNodeWriteRepository;
|
||||
use crate::repository::quota::ProviderQuotaWriteRepository;
|
||||
use crate::repository::routing_profiles::RoutingGroupWriteRepository;
|
||||
use crate::repository::settlement::SettlementWriteRepository;
|
||||
use crate::repository::usage::UsageWriteRepository;
|
||||
use crate::repository::video_tasks::VideoTaskWriteRepository;
|
||||
@@ -35,6 +36,7 @@ pub struct DataWriteRepositories {
|
||||
proxy_nodes: Option<Arc<dyn ProxyNodeWriteRepository>>,
|
||||
provider_catalog: Option<Arc<dyn ProviderCatalogWriteRepository>>,
|
||||
provider_quotas: Option<Arc<dyn ProviderQuotaWriteRepository>>,
|
||||
routing_groups: Option<Arc<dyn RoutingGroupWriteRepository>>,
|
||||
settlement: Option<Arc<dyn SettlementWriteRepository>>,
|
||||
usage: Option<Arc<dyn UsageWriteRepository>>,
|
||||
video_tasks: Option<Arc<dyn VideoTaskWriteRepository>>,
|
||||
@@ -60,6 +62,7 @@ impl fmt::Debug for DataWriteRepositories {
|
||||
.field("has_proxy_nodes", &self.proxy_nodes.is_some())
|
||||
.field("has_provider_catalog", &self.provider_catalog.is_some())
|
||||
.field("has_provider_quotas", &self.provider_quotas.is_some())
|
||||
.field("has_routing_groups", &self.routing_groups.is_some())
|
||||
.field("has_settlement", &self.settlement.is_some())
|
||||
.field("has_usage", &self.usage.is_some())
|
||||
.field("has_video_tasks", &self.video_tasks.is_some())
|
||||
@@ -127,6 +130,10 @@ impl DataWriteRepositories {
|
||||
.map(PostgresBackend::provider_quota_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::provider_quota_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::provider_quota_write_repository)),
|
||||
routing_groups: postgres
|
||||
.map(PostgresBackend::routing_group_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::routing_group_write_repository))
|
||||
.or_else(|| sqlite.map(SqliteBackend::routing_group_write_repository)),
|
||||
settlement: postgres
|
||||
.map(PostgresBackend::settlement_write_repository)
|
||||
.or_else(|| mysql.map(MysqlBackend::settlement_write_repository))
|
||||
@@ -203,6 +210,10 @@ impl DataWriteRepositories {
|
||||
self.provider_quotas.clone()
|
||||
}
|
||||
|
||||
pub fn routing_groups(&self) -> Option<Arc<dyn RoutingGroupWriteRepository>> {
|
||||
self.routing_groups.clone()
|
||||
}
|
||||
|
||||
pub fn provider_catalog(&self) -> Option<Arc<dyn ProviderCatalogWriteRepository>> {
|
||||
self.provider_catalog.clone()
|
||||
}
|
||||
@@ -233,6 +244,7 @@ impl DataWriteRepositories {
|
||||
|| self.proxy_nodes.is_some()
|
||||
|| self.provider_catalog.is_some()
|
||||
|| self.provider_quotas.is_some()
|
||||
|| self.routing_groups.is_some()
|
||||
|| self.settlement.is_some()
|
||||
|| self.usage.is_some()
|
||||
|| self.video_tasks.is_some()
|
||||
|
||||
@@ -7,7 +7,7 @@ use tracing::info;
|
||||
// Generated by build.rs from schema/bootstrap/postgres.
|
||||
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
|
||||
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
|
||||
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260515000000;
|
||||
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260516000000;
|
||||
|
||||
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
|
||||
SELECT COUNT(*)::BIGINT
|
||||
@@ -29,6 +29,7 @@ WHERE table_schema = 'public'
|
||||
'oauth_providers',
|
||||
'provider_api_keys',
|
||||
'proxy_nodes',
|
||||
'routing_groups',
|
||||
'user_groups',
|
||||
'usage_routing_snapshots',
|
||||
'usage_settlement_snapshots'
|
||||
|
||||
@@ -305,6 +305,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
|
||||
20260512090000,
|
||||
20260512110000,
|
||||
20260515000000,
|
||||
20260516000000,
|
||||
]
|
||||
);
|
||||
}
|
||||
@@ -587,6 +588,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
|
||||
20260512000000,
|
||||
20260512090000,
|
||||
20260512110000,
|
||||
20260516000000,
|
||||
]
|
||||
);
|
||||
assert_eq!(
|
||||
@@ -603,6 +605,7 @@ fn mysql_and_sqlite_migrations_include_enabled_incrementals() {
|
||||
20260512000000,
|
||||
20260512090000,
|
||||
20260512110000,
|
||||
20260516000000,
|
||||
]
|
||||
);
|
||||
}
|
||||
@@ -1119,6 +1122,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
|
||||
20260512090000,
|
||||
20260512110000,
|
||||
20260515000000,
|
||||
20260516000000,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -21,6 +21,7 @@ pub mod provider_catalog;
|
||||
pub mod provider_oauth;
|
||||
pub mod proxy_nodes;
|
||||
pub mod quota;
|
||||
pub mod routing_profiles;
|
||||
pub mod settlement;
|
||||
pub mod system;
|
||||
pub mod usage;
|
||||
|
||||
342
crates/aether-data/src/repository/routing_profiles/memory.rs
Normal file
342
crates/aether-data/src/repository/routing_profiles/memory.rs
Normal file
@@ -0,0 +1,342 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::{
|
||||
CreateRoutingGroupBindingRecord, CreateRoutingGroupRecord, CreateRoutingGroupVersionRecord,
|
||||
RoutingGroupBindingQuery, RoutingGroupLookupKey, RoutingGroupReadRepository,
|
||||
RoutingGroupWriteRepository, StoredRoutingGroup, StoredRoutingGroupBinding,
|
||||
StoredRoutingGroupVersion, UpdateRoutingGroupBindingRecord, UpdateRoutingGroupRecord,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryRoutingGroupRepository {
|
||||
groups: RwLock<BTreeMap<String, StoredRoutingGroup>>,
|
||||
bindings: RwLock<BTreeMap<String, StoredRoutingGroupBinding>>,
|
||||
versions: RwLock<BTreeMap<String, StoredRoutingGroupVersion>>,
|
||||
}
|
||||
|
||||
impl InMemoryRoutingGroupRepository {
|
||||
pub fn seed<I, B, V>(groups: I, bindings: B, versions: V) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredRoutingGroup>,
|
||||
B: IntoIterator<Item = StoredRoutingGroupBinding>,
|
||||
V: IntoIterator<Item = StoredRoutingGroupVersion>,
|
||||
{
|
||||
Self {
|
||||
groups: RwLock::new(
|
||||
groups
|
||||
.into_iter()
|
||||
.map(|item| (item.id.clone(), item))
|
||||
.collect(),
|
||||
),
|
||||
bindings: RwLock::new(
|
||||
bindings
|
||||
.into_iter()
|
||||
.map(|item| (item.id.clone(), item))
|
||||
.collect(),
|
||||
),
|
||||
versions: RwLock::new(
|
||||
versions
|
||||
.into_iter()
|
||||
.map(|item| (item.id.clone(), item))
|
||||
.collect(),
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RoutingGroupReadRepository for InMemoryRoutingGroupRepository {
|
||||
async fn list_routing_groups(&self) -> Result<Vec<StoredRoutingGroup>, DataLayerError> {
|
||||
let mut groups = self
|
||||
.groups
|
||||
.read()
|
||||
.expect("routing group repository lock")
|
||||
.values()
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
groups.sort_by(|left, right| left.name.cmp(&right.name).then(left.id.cmp(&right.id)));
|
||||
Ok(groups)
|
||||
}
|
||||
|
||||
async fn find_routing_group(
|
||||
&self,
|
||||
lookup: RoutingGroupLookupKey<'_>,
|
||||
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
let groups = self.groups.read().expect("routing group repository lock");
|
||||
Ok(match lookup {
|
||||
RoutingGroupLookupKey::Id(id) => groups.get(id).cloned(),
|
||||
RoutingGroupLookupKey::Name(name) => {
|
||||
groups.values().find(|group| group.name == name).cloned()
|
||||
}
|
||||
RoutingGroupLookupKey::SystemDefault => groups
|
||||
.values()
|
||||
.find(|group| group.is_system_default && group.enabled)
|
||||
.cloned(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_routing_group_bindings(
|
||||
&self,
|
||||
query: &RoutingGroupBindingQuery,
|
||||
) -> Result<Vec<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
let mut rows = self
|
||||
.bindings
|
||||
.read()
|
||||
.expect("routing group binding repository lock")
|
||||
.values()
|
||||
.filter(|row| {
|
||||
query
|
||||
.group_id
|
||||
.as_ref()
|
||||
.is_none_or(|group_id| &row.group_id == group_id)
|
||||
&& query
|
||||
.subject_type
|
||||
.as_ref()
|
||||
.is_none_or(|subject_type| &row.subject_type == subject_type)
|
||||
&& query
|
||||
.subject_id
|
||||
.as_ref()
|
||||
.is_none_or(|subject_id| &row.subject_id == subject_id)
|
||||
})
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
rows.sort_by(|left, right| {
|
||||
left.created_at
|
||||
.cmp(&right.created_at)
|
||||
.then(left.id.cmp(&right.id))
|
||||
});
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
async fn list_routing_group_versions(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Vec<StoredRoutingGroupVersion>, DataLayerError> {
|
||||
let mut rows = self
|
||||
.versions
|
||||
.read()
|
||||
.expect("routing group version repository lock")
|
||||
.values()
|
||||
.filter(|row| row.group_id == group_id)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>();
|
||||
rows.sort_by(|left, right| {
|
||||
right
|
||||
.version
|
||||
.cmp(&left.version)
|
||||
.then(right.created_at.cmp(&left.created_at))
|
||||
});
|
||||
Ok(rows)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RoutingGroupWriteRepository for InMemoryRoutingGroupRepository {
|
||||
async fn create_routing_group(
|
||||
&self,
|
||||
record: CreateRoutingGroupRecord,
|
||||
) -> Result<StoredRoutingGroup, DataLayerError> {
|
||||
let group = StoredRoutingGroup::new(record)?;
|
||||
self.groups
|
||||
.write()
|
||||
.expect("routing group repository lock")
|
||||
.insert(group.id.clone(), group.clone());
|
||||
Ok(group)
|
||||
}
|
||||
|
||||
async fn update_routing_group(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupRecord,
|
||||
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
let mut groups = self.groups.write().expect("routing group repository lock");
|
||||
let Some(group) = groups.get_mut(id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if let Some(name) = patch.name {
|
||||
if name.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"routing_groups.name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
group.name = name;
|
||||
}
|
||||
if let Some(description) = patch.description {
|
||||
group.description = description;
|
||||
}
|
||||
if let Some(enabled) = patch.enabled {
|
||||
group.enabled = enabled;
|
||||
}
|
||||
if let Some(is_system_default) = patch.is_system_default {
|
||||
group.is_system_default = is_system_default;
|
||||
}
|
||||
if let Some(config_json) = patch.config_json {
|
||||
if !config_json.is_object() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"routing_groups.config_json must be a JSON object".to_string(),
|
||||
));
|
||||
}
|
||||
group.config_json = config_json;
|
||||
}
|
||||
if let Some(version) = patch.version {
|
||||
group.version = version.max(1);
|
||||
}
|
||||
if let Some(published_at) = patch.published_at {
|
||||
group.published_at = published_at;
|
||||
}
|
||||
group.updated_at = patch.updated_at;
|
||||
Ok(Some(group.clone()))
|
||||
}
|
||||
|
||||
async fn delete_routing_group(&self, id: &str) -> Result<bool, DataLayerError> {
|
||||
Ok(self
|
||||
.groups
|
||||
.write()
|
||||
.expect("routing group repository lock")
|
||||
.remove(id)
|
||||
.is_some())
|
||||
}
|
||||
|
||||
async fn create_routing_group_binding(
|
||||
&self,
|
||||
record: CreateRoutingGroupBindingRecord,
|
||||
) -> Result<StoredRoutingGroupBinding, DataLayerError> {
|
||||
let binding = StoredRoutingGroupBinding::new(record)?;
|
||||
self.bindings
|
||||
.write()
|
||||
.expect("routing group binding repository lock")
|
||||
.insert(binding.id.clone(), binding.clone());
|
||||
Ok(binding)
|
||||
}
|
||||
|
||||
async fn delete_routing_group_binding(&self, id: &str) -> Result<bool, DataLayerError> {
|
||||
Ok(self
|
||||
.bindings
|
||||
.write()
|
||||
.expect("routing group binding repository lock")
|
||||
.remove(id)
|
||||
.is_some())
|
||||
}
|
||||
|
||||
async fn update_routing_group_binding(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupBindingRecord,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
let mut bindings = self
|
||||
.bindings
|
||||
.write()
|
||||
.expect("routing group binding repository lock");
|
||||
let Some(binding) = bindings.get_mut(id) else {
|
||||
return Ok(None);
|
||||
};
|
||||
if let Some(group_id) = patch.group_id {
|
||||
if group_id.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"routing_group_bindings.group_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
binding.group_id = group_id;
|
||||
}
|
||||
if let Some(subject_type) = patch.subject_type {
|
||||
binding.subject_type = subject_type;
|
||||
}
|
||||
if let Some(subject_id) = patch.subject_id {
|
||||
if subject_id.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"routing_group_bindings.subject_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
binding.subject_id = subject_id;
|
||||
}
|
||||
if let Some(is_default) = patch.is_default {
|
||||
binding.is_default = is_default;
|
||||
}
|
||||
if let Some(allow_explicit_select) = patch.allow_explicit_select {
|
||||
binding.allow_explicit_select = allow_explicit_select;
|
||||
}
|
||||
binding.updated_at = patch.updated_at;
|
||||
Ok(Some(binding.clone()))
|
||||
}
|
||||
|
||||
async fn create_routing_group_version(
|
||||
&self,
|
||||
record: CreateRoutingGroupVersionRecord,
|
||||
) -> Result<StoredRoutingGroupVersion, DataLayerError> {
|
||||
let version = StoredRoutingGroupVersion::new(record)?;
|
||||
self.versions
|
||||
.write()
|
||||
.expect("routing group version repository lock")
|
||||
.insert(version.id.clone(), version.clone());
|
||||
Ok(version)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use aether_data_contracts::repository::routing_profiles::RoutingGroupBindingSubject;
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn stores_groups_bindings_and_versions() {
|
||||
let repository = InMemoryRoutingGroupRepository::default();
|
||||
let group = repository
|
||||
.create_routing_group(CreateRoutingGroupRecord {
|
||||
id: "group-1".to_string(),
|
||||
name: "default".to_string(),
|
||||
description: None,
|
||||
enabled: true,
|
||||
is_system_default: true,
|
||||
config_json: json!({}),
|
||||
version: 1,
|
||||
created_at: 1,
|
||||
updated_at: 1,
|
||||
published_at: None,
|
||||
})
|
||||
.await
|
||||
.expect("group should store");
|
||||
|
||||
assert_eq!(
|
||||
repository
|
||||
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
|
||||
.await
|
||||
.unwrap()
|
||||
.as_ref()
|
||||
.map(|group| group.id.as_str()),
|
||||
Some(group.id.as_str())
|
||||
);
|
||||
|
||||
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();
|
||||
|
||||
assert_eq!(
|
||||
repository
|
||||
.list_routing_group_bindings(&RoutingGroupBindingQuery {
|
||||
subject_type: Some(RoutingGroupBindingSubject::ApiKey),
|
||||
subject_id: Some("api-key-1".to_string()),
|
||||
group_id: None,
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
}
|
||||
}
|
||||
16
crates/aether-data/src/repository/routing_profiles/mod.rs
Normal file
16
crates/aether-data/src/repository/routing_profiles/mod.rs
Normal file
@@ -0,0 +1,16 @@
|
||||
mod memory;
|
||||
mod mysql;
|
||||
mod postgres;
|
||||
mod sqlite;
|
||||
|
||||
pub(crate) use aether_data_contracts::repository::routing_profiles::{
|
||||
CreateRoutingGroupBindingRecord, CreateRoutingGroupRecord, CreateRoutingGroupVersionRecord,
|
||||
RoutingGroupBindingQuery, RoutingGroupBindingSubject, RoutingGroupLookupKey,
|
||||
RoutingGroupReadRepository, RoutingGroupWriteRepository, StoredRoutingGroup,
|
||||
StoredRoutingGroupBinding, StoredRoutingGroupVersion, UpdateRoutingGroupBindingRecord,
|
||||
UpdateRoutingGroupRecord,
|
||||
};
|
||||
pub use memory::InMemoryRoutingGroupRepository;
|
||||
pub use mysql::MysqlRoutingGroupRepository;
|
||||
pub use postgres::PostgresRoutingGroupRepository;
|
||||
pub use sqlite::SqliteRoutingGroupRepository;
|
||||
417
crates/aether-data/src/repository/routing_profiles/mysql.rs
Normal file
417
crates/aether-data/src/repository/routing_profiles/mysql.rs
Normal file
@@ -0,0 +1,417 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
use sqlx::{mysql::MySqlRow, Row};
|
||||
|
||||
use super::postgres::{
|
||||
apply_binding_patch, apply_group_patch, binding_subject_from_database,
|
||||
binding_subject_to_database,
|
||||
};
|
||||
use super::*;
|
||||
use crate::driver::mysql::MysqlPool;
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
const ROUTING_GROUP_SELECT: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
name,
|
||||
description,
|
||||
enabled,
|
||||
is_system_default,
|
||||
config_json,
|
||||
version,
|
||||
created_at,
|
||||
updated_at,
|
||||
published_at
|
||||
FROM routing_groups
|
||||
"#;
|
||||
|
||||
const ROUTING_GROUP_BINDING_SELECT: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
group_id,
|
||||
subject_type,
|
||||
subject_id,
|
||||
is_default,
|
||||
allow_explicit_select,
|
||||
created_at,
|
||||
updated_at
|
||||
FROM routing_group_bindings
|
||||
"#;
|
||||
|
||||
const ROUTING_GROUP_VERSION_SELECT: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
group_id,
|
||||
version,
|
||||
config_json,
|
||||
created_at,
|
||||
created_by
|
||||
FROM routing_group_versions
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MysqlRoutingGroupRepository {
|
||||
pool: MysqlPool,
|
||||
}
|
||||
|
||||
impl MysqlRoutingGroupRepository {
|
||||
pub fn new(pool: MysqlPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
async fn reload_group(&self, id: &str) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
self.find_routing_group(RoutingGroupLookupKey::Id(id)).await
|
||||
}
|
||||
|
||||
async fn find_binding_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
let row = sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_BINDING_SELECT} WHERE id = ? LIMIT 1"
|
||||
))
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
row.as_ref().map(map_binding_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RoutingGroupReadRepository for MysqlRoutingGroupRepository {
|
||||
async fn list_routing_groups(&self) -> Result<Vec<StoredRoutingGroup>, DataLayerError> {
|
||||
let rows = sqlx::query(&format!("{ROUTING_GROUP_SELECT} ORDER BY name ASC, id ASC"))
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_group_row).collect()
|
||||
}
|
||||
|
||||
async fn find_routing_group(
|
||||
&self,
|
||||
lookup: RoutingGroupLookupKey<'_>,
|
||||
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
let row = match lookup {
|
||||
RoutingGroupLookupKey::Id(id) => sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_SELECT} WHERE id = ? LIMIT 1"
|
||||
))
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?,
|
||||
RoutingGroupLookupKey::Name(name) => sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_SELECT} WHERE name = ? LIMIT 1"
|
||||
))
|
||||
.bind(name)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?,
|
||||
RoutingGroupLookupKey::SystemDefault => sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_SELECT} WHERE is_system_default = 1 AND enabled = 1 ORDER BY updated_at DESC, id ASC LIMIT 1"
|
||||
))
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?,
|
||||
};
|
||||
row.as_ref().map(map_group_row).transpose()
|
||||
}
|
||||
|
||||
async fn list_routing_group_bindings(
|
||||
&self,
|
||||
query: &RoutingGroupBindingQuery,
|
||||
) -> Result<Vec<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
let rows = sqlx::query(&format!(
|
||||
r#"
|
||||
{ROUTING_GROUP_BINDING_SELECT}
|
||||
WHERE (? IS NULL OR group_id = ?)
|
||||
AND (? IS NULL OR subject_type = ?)
|
||||
AND (? IS NULL OR subject_id = ?)
|
||||
ORDER BY created_at ASC, id ASC
|
||||
"#
|
||||
))
|
||||
.bind(query.group_id.as_deref())
|
||||
.bind(query.group_id.as_deref())
|
||||
.bind(query.subject_type.map(binding_subject_to_database))
|
||||
.bind(query.subject_type.map(binding_subject_to_database))
|
||||
.bind(query.subject_id.as_deref())
|
||||
.bind(query.subject_id.as_deref())
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_binding_row).collect()
|
||||
}
|
||||
|
||||
async fn list_routing_group_versions(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Vec<StoredRoutingGroupVersion>, DataLayerError> {
|
||||
let rows = sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_VERSION_SELECT} WHERE group_id = ? ORDER BY version DESC, created_at DESC, id ASC"
|
||||
))
|
||||
.bind(group_id)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_version_row).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RoutingGroupWriteRepository for MysqlRoutingGroupRepository {
|
||||
async fn create_routing_group(
|
||||
&self,
|
||||
record: CreateRoutingGroupRecord,
|
||||
) -> Result<StoredRoutingGroup, DataLayerError> {
|
||||
let group = StoredRoutingGroup::new(record)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO routing_groups (
|
||||
id, name, description, enabled, is_system_default, config_json,
|
||||
version, created_at, updated_at, published_at
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(&group.id)
|
||||
.bind(&group.name)
|
||||
.bind(&group.description)
|
||||
.bind(group.enabled)
|
||||
.bind(group.is_system_default)
|
||||
.bind(json_to_string(
|
||||
&group.config_json,
|
||||
"routing_groups.config_json",
|
||||
)?)
|
||||
.bind(group.version)
|
||||
.bind(group.created_at)
|
||||
.bind(group.updated_at)
|
||||
.bind(group.published_at)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(group)
|
||||
}
|
||||
|
||||
async fn update_routing_group(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupRecord,
|
||||
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
let Some(mut group) = self.reload_group(id).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
apply_group_patch(&mut group, patch)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE routing_groups
|
||||
SET name = ?,
|
||||
description = ?,
|
||||
enabled = ?,
|
||||
is_system_default = ?,
|
||||
config_json = ?,
|
||||
version = ?,
|
||||
updated_at = ?,
|
||||
published_at = ?
|
||||
WHERE id = ?
|
||||
"#,
|
||||
)
|
||||
.bind(&group.name)
|
||||
.bind(&group.description)
|
||||
.bind(group.enabled)
|
||||
.bind(group.is_system_default)
|
||||
.bind(json_to_string(
|
||||
&group.config_json,
|
||||
"routing_groups.config_json",
|
||||
)?)
|
||||
.bind(group.version)
|
||||
.bind(group.updated_at)
|
||||
.bind(group.published_at)
|
||||
.bind(id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(Some(group))
|
||||
}
|
||||
|
||||
async fn delete_routing_group(&self, id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut tx = self.pool.begin().await.map_sql_err()?;
|
||||
sqlx::query("DELETE FROM routing_group_bindings WHERE group_id = ?")
|
||||
.bind(id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
sqlx::query("DELETE FROM routing_group_versions WHERE group_id = ?")
|
||||
.bind(id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let rows_affected = sqlx::query("DELETE FROM routing_groups WHERE id = ?")
|
||||
.bind(id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected();
|
||||
tx.commit().await.map_sql_err()?;
|
||||
Ok(rows_affected > 0)
|
||||
}
|
||||
|
||||
async fn create_routing_group_binding(
|
||||
&self,
|
||||
record: CreateRoutingGroupBindingRecord,
|
||||
) -> Result<StoredRoutingGroupBinding, DataLayerError> {
|
||||
let binding = StoredRoutingGroupBinding::new(record)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO routing_group_bindings (
|
||||
id, group_id, subject_type, subject_id, is_default,
|
||||
allow_explicit_select, created_at, updated_at
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(&binding.id)
|
||||
.bind(&binding.group_id)
|
||||
.bind(binding_subject_to_database(binding.subject_type))
|
||||
.bind(&binding.subject_id)
|
||||
.bind(binding.is_default)
|
||||
.bind(binding.allow_explicit_select)
|
||||
.bind(binding.created_at)
|
||||
.bind(binding.updated_at)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(binding)
|
||||
}
|
||||
|
||||
async fn delete_routing_group_binding(&self, id: &str) -> Result<bool, DataLayerError> {
|
||||
Ok(
|
||||
sqlx::query("DELETE FROM routing_group_bindings WHERE id = ?")
|
||||
.bind(id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected()
|
||||
> 0,
|
||||
)
|
||||
}
|
||||
|
||||
async fn update_routing_group_binding(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupBindingRecord,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
let Some(mut binding) = self.find_binding_by_id(id).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
apply_binding_patch(&mut binding, patch)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE routing_group_bindings
|
||||
SET group_id = ?,
|
||||
subject_type = ?,
|
||||
subject_id = ?,
|
||||
is_default = ?,
|
||||
allow_explicit_select = ?,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
"#,
|
||||
)
|
||||
.bind(&binding.group_id)
|
||||
.bind(binding_subject_to_database(binding.subject_type))
|
||||
.bind(&binding.subject_id)
|
||||
.bind(binding.is_default)
|
||||
.bind(binding.allow_explicit_select)
|
||||
.bind(binding.updated_at)
|
||||
.bind(id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(Some(binding))
|
||||
}
|
||||
|
||||
async fn create_routing_group_version(
|
||||
&self,
|
||||
record: CreateRoutingGroupVersionRecord,
|
||||
) -> Result<StoredRoutingGroupVersion, DataLayerError> {
|
||||
let version = StoredRoutingGroupVersion::new(record)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO routing_group_versions (
|
||||
id, group_id, version, config_json, created_at, created_by
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(&version.id)
|
||||
.bind(&version.group_id)
|
||||
.bind(version.version)
|
||||
.bind(json_to_string(
|
||||
&version.config_json,
|
||||
"routing_group_versions.config_json",
|
||||
)?)
|
||||
.bind(version.created_at)
|
||||
.bind(&version.created_by)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(version)
|
||||
}
|
||||
}
|
||||
|
||||
fn map_group_row(row: &MySqlRow) -> Result<StoredRoutingGroup, DataLayerError> {
|
||||
Ok(StoredRoutingGroup {
|
||||
id: row.try_get("id").map_sql_err()?,
|
||||
name: row.try_get("name").map_sql_err()?,
|
||||
description: row.try_get("description").map_sql_err()?,
|
||||
enabled: row.try_get("enabled").map_sql_err()?,
|
||||
is_system_default: row.try_get("is_system_default").map_sql_err()?,
|
||||
config_json: json_from_string(
|
||||
row.try_get("config_json").map_sql_err()?,
|
||||
"routing_groups.config_json",
|
||||
)?,
|
||||
version: row.try_get("version").map_sql_err()?,
|
||||
created_at: row.try_get("created_at").map_sql_err()?,
|
||||
updated_at: row.try_get("updated_at").map_sql_err()?,
|
||||
published_at: row.try_get("published_at").map_sql_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_binding_row(row: &MySqlRow) -> Result<StoredRoutingGroupBinding, DataLayerError> {
|
||||
Ok(StoredRoutingGroupBinding {
|
||||
id: row.try_get("id").map_sql_err()?,
|
||||
group_id: row.try_get("group_id").map_sql_err()?,
|
||||
subject_type: binding_subject_from_database(row.try_get("subject_type").map_sql_err()?)?,
|
||||
subject_id: row.try_get("subject_id").map_sql_err()?,
|
||||
is_default: row.try_get("is_default").map_sql_err()?,
|
||||
allow_explicit_select: row.try_get("allow_explicit_select").map_sql_err()?,
|
||||
created_at: row.try_get("created_at").map_sql_err()?,
|
||||
updated_at: row.try_get("updated_at").map_sql_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_version_row(row: &MySqlRow) -> Result<StoredRoutingGroupVersion, DataLayerError> {
|
||||
Ok(StoredRoutingGroupVersion {
|
||||
id: row.try_get("id").map_sql_err()?,
|
||||
group_id: row.try_get("group_id").map_sql_err()?,
|
||||
version: row.try_get("version").map_sql_err()?,
|
||||
config_json: json_from_string(
|
||||
row.try_get("config_json").map_sql_err()?,
|
||||
"routing_group_versions.config_json",
|
||||
)?,
|
||||
created_at: row.try_get("created_at").map_sql_err()?,
|
||||
created_by: row.try_get("created_by").map_sql_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn json_to_string(value: &Value, field_name: &str) -> Result<String, DataLayerError> {
|
||||
serde_json::to_string(value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("{field_name} contains unserializable JSON: {err}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn json_from_string(value: String, field_name: &str) -> Result<Value, DataLayerError> {
|
||||
serde_json::from_str(&value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("{field_name} contains invalid JSON: {err}"))
|
||||
})
|
||||
}
|
||||
483
crates/aether-data/src/repository/routing_profiles/postgres.rs
Normal file
483
crates/aether-data/src/repository/routing_profiles/postgres.rs
Normal file
@@ -0,0 +1,483 @@
|
||||
use async_trait::async_trait;
|
||||
use futures_util::TryStreamExt;
|
||||
use sqlx::{postgres::PgRow, PgPool, Row};
|
||||
|
||||
use super::*;
|
||||
use crate::error::SqlxResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
const ROUTING_GROUP_SELECT: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
name,
|
||||
description,
|
||||
enabled,
|
||||
is_system_default,
|
||||
config_json,
|
||||
version,
|
||||
created_at,
|
||||
updated_at,
|
||||
published_at
|
||||
FROM routing_groups
|
||||
"#;
|
||||
|
||||
const ROUTING_GROUP_BINDING_SELECT: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
group_id,
|
||||
subject_type,
|
||||
subject_id,
|
||||
is_default,
|
||||
allow_explicit_select,
|
||||
created_at,
|
||||
updated_at
|
||||
FROM routing_group_bindings
|
||||
"#;
|
||||
|
||||
const ROUTING_GROUP_VERSION_SELECT: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
group_id,
|
||||
version,
|
||||
config_json,
|
||||
created_at,
|
||||
created_by
|
||||
FROM routing_group_versions
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PostgresRoutingGroupRepository {
|
||||
pool: PgPool,
|
||||
}
|
||||
|
||||
impl PostgresRoutingGroupRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
async fn reload_group(&self, id: &str) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
self.find_routing_group(RoutingGroupLookupKey::Id(id)).await
|
||||
}
|
||||
|
||||
async fn find_binding_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
let row = sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_BINDING_SELECT} WHERE id = $1 LIMIT 1"
|
||||
))
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
row.as_ref().map(map_binding_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RoutingGroupReadRepository for PostgresRoutingGroupRepository {
|
||||
async fn list_routing_groups(&self) -> Result<Vec<StoredRoutingGroup>, DataLayerError> {
|
||||
let sql = format!("{ROUTING_GROUP_SELECT} ORDER BY name ASC, id ASC");
|
||||
let mut rows = sqlx::query(&sql).fetch(&self.pool);
|
||||
let mut groups = Vec::new();
|
||||
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||
groups.push(map_group_row(&row)?);
|
||||
}
|
||||
Ok(groups)
|
||||
}
|
||||
|
||||
async fn find_routing_group(
|
||||
&self,
|
||||
lookup: RoutingGroupLookupKey<'_>,
|
||||
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
let row = match lookup {
|
||||
RoutingGroupLookupKey::Id(id) => sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_SELECT} WHERE id = $1 LIMIT 1"
|
||||
))
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?,
|
||||
RoutingGroupLookupKey::Name(name) => sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_SELECT} WHERE name = $1 LIMIT 1"
|
||||
))
|
||||
.bind(name)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?,
|
||||
RoutingGroupLookupKey::SystemDefault => sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_SELECT} WHERE is_system_default = TRUE AND enabled = TRUE ORDER BY updated_at DESC, id ASC LIMIT 1"
|
||||
))
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?,
|
||||
};
|
||||
row.as_ref().map(map_group_row).transpose()
|
||||
}
|
||||
|
||||
async fn list_routing_group_bindings(
|
||||
&self,
|
||||
query: &RoutingGroupBindingQuery,
|
||||
) -> Result<Vec<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
let sql = format!(
|
||||
r#"
|
||||
{ROUTING_GROUP_BINDING_SELECT}
|
||||
WHERE ($1::text IS NULL OR group_id = $1)
|
||||
AND ($2::text IS NULL OR subject_type = $2)
|
||||
AND ($3::text IS NULL OR subject_id = $3)
|
||||
ORDER BY created_at ASC, id ASC
|
||||
"#
|
||||
);
|
||||
let mut rows = sqlx::query(&sql)
|
||||
.bind(query.group_id.as_deref())
|
||||
.bind(query.subject_type.map(binding_subject_to_database))
|
||||
.bind(query.subject_id.as_deref())
|
||||
.fetch(&self.pool);
|
||||
let mut bindings = Vec::new();
|
||||
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||
bindings.push(map_binding_row(&row)?);
|
||||
}
|
||||
Ok(bindings)
|
||||
}
|
||||
|
||||
async fn list_routing_group_versions(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Vec<StoredRoutingGroupVersion>, DataLayerError> {
|
||||
let sql = format!(
|
||||
"{ROUTING_GROUP_VERSION_SELECT} WHERE group_id = $1 ORDER BY version DESC, created_at DESC, id ASC"
|
||||
);
|
||||
let mut rows = sqlx::query(&sql).bind(group_id).fetch(&self.pool);
|
||||
let mut versions = Vec::new();
|
||||
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||
versions.push(map_version_row(&row)?);
|
||||
}
|
||||
Ok(versions)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RoutingGroupWriteRepository for PostgresRoutingGroupRepository {
|
||||
async fn create_routing_group(
|
||||
&self,
|
||||
record: CreateRoutingGroupRecord,
|
||||
) -> Result<StoredRoutingGroup, DataLayerError> {
|
||||
let group = StoredRoutingGroup::new(record)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO routing_groups (
|
||||
id, name, description, enabled, is_system_default, config_json,
|
||||
version, created_at, updated_at, published_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)
|
||||
"#,
|
||||
)
|
||||
.bind(&group.id)
|
||||
.bind(&group.name)
|
||||
.bind(&group.description)
|
||||
.bind(group.enabled)
|
||||
.bind(group.is_system_default)
|
||||
.bind(&group.config_json)
|
||||
.bind(group.version)
|
||||
.bind(group.created_at)
|
||||
.bind(group.updated_at)
|
||||
.bind(group.published_at)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(group)
|
||||
}
|
||||
|
||||
async fn update_routing_group(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupRecord,
|
||||
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
let Some(mut group) = self.reload_group(id).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
apply_group_patch(&mut group, patch)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE routing_groups
|
||||
SET name = $2,
|
||||
description = $3,
|
||||
enabled = $4,
|
||||
is_system_default = $5,
|
||||
config_json = $6,
|
||||
version = $7,
|
||||
updated_at = $8,
|
||||
published_at = $9
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(id)
|
||||
.bind(&group.name)
|
||||
.bind(&group.description)
|
||||
.bind(group.enabled)
|
||||
.bind(group.is_system_default)
|
||||
.bind(&group.config_json)
|
||||
.bind(group.version)
|
||||
.bind(group.updated_at)
|
||||
.bind(group.published_at)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(Some(group))
|
||||
}
|
||||
|
||||
async fn delete_routing_group(&self, id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut tx = self.pool.begin().await.map_postgres_err()?;
|
||||
sqlx::query("DELETE FROM routing_group_bindings WHERE group_id = $1")
|
||||
.bind(id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
sqlx::query("DELETE FROM routing_group_versions WHERE group_id = $1")
|
||||
.bind(id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
let rows_affected = sqlx::query("DELETE FROM routing_groups WHERE id = $1")
|
||||
.bind(id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected();
|
||||
tx.commit().await.map_postgres_err()?;
|
||||
Ok(rows_affected > 0)
|
||||
}
|
||||
|
||||
async fn create_routing_group_binding(
|
||||
&self,
|
||||
record: CreateRoutingGroupBindingRecord,
|
||||
) -> Result<StoredRoutingGroupBinding, DataLayerError> {
|
||||
let binding = StoredRoutingGroupBinding::new(record)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO routing_group_bindings (
|
||||
id, group_id, subject_type, subject_id, is_default,
|
||||
allow_explicit_select, created_at, updated_at
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
|
||||
"#,
|
||||
)
|
||||
.bind(&binding.id)
|
||||
.bind(&binding.group_id)
|
||||
.bind(binding_subject_to_database(binding.subject_type))
|
||||
.bind(&binding.subject_id)
|
||||
.bind(binding.is_default)
|
||||
.bind(binding.allow_explicit_select)
|
||||
.bind(binding.created_at)
|
||||
.bind(binding.updated_at)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(binding)
|
||||
}
|
||||
|
||||
async fn delete_routing_group_binding(&self, id: &str) -> Result<bool, DataLayerError> {
|
||||
Ok(
|
||||
sqlx::query("DELETE FROM routing_group_bindings WHERE id = $1")
|
||||
.bind(id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?
|
||||
.rows_affected()
|
||||
> 0,
|
||||
)
|
||||
}
|
||||
|
||||
async fn update_routing_group_binding(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupBindingRecord,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
let Some(mut binding) = self.find_binding_by_id(id).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
apply_binding_patch(&mut binding, patch)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE routing_group_bindings
|
||||
SET group_id = $2,
|
||||
subject_type = $3,
|
||||
subject_id = $4,
|
||||
is_default = $5,
|
||||
allow_explicit_select = $6,
|
||||
updated_at = $7
|
||||
WHERE id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(id)
|
||||
.bind(&binding.group_id)
|
||||
.bind(binding_subject_to_database(binding.subject_type))
|
||||
.bind(&binding.subject_id)
|
||||
.bind(binding.is_default)
|
||||
.bind(binding.allow_explicit_select)
|
||||
.bind(binding.updated_at)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(Some(binding))
|
||||
}
|
||||
|
||||
async fn create_routing_group_version(
|
||||
&self,
|
||||
record: CreateRoutingGroupVersionRecord,
|
||||
) -> Result<StoredRoutingGroupVersion, DataLayerError> {
|
||||
let version = StoredRoutingGroupVersion::new(record)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO routing_group_versions (
|
||||
id, group_id, version, config_json, created_at, created_by
|
||||
)
|
||||
VALUES ($1, $2, $3, $4, $5, $6)
|
||||
"#,
|
||||
)
|
||||
.bind(&version.id)
|
||||
.bind(&version.group_id)
|
||||
.bind(version.version)
|
||||
.bind(&version.config_json)
|
||||
.bind(version.created_at)
|
||||
.bind(&version.created_by)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(version)
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn apply_group_patch(
|
||||
group: &mut StoredRoutingGroup,
|
||||
patch: UpdateRoutingGroupRecord,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if let Some(name) = patch.name {
|
||||
if name.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"routing_groups.name is empty".to_string(),
|
||||
));
|
||||
}
|
||||
group.name = name;
|
||||
}
|
||||
if let Some(description) = patch.description {
|
||||
group.description = description;
|
||||
}
|
||||
if let Some(enabled) = patch.enabled {
|
||||
group.enabled = enabled;
|
||||
}
|
||||
if let Some(is_system_default) = patch.is_system_default {
|
||||
group.is_system_default = is_system_default;
|
||||
}
|
||||
if let Some(config_json) = patch.config_json {
|
||||
if !config_json.is_object() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"routing_groups.config_json must be a JSON object".to_string(),
|
||||
));
|
||||
}
|
||||
group.config_json = config_json;
|
||||
}
|
||||
if let Some(version) = patch.version {
|
||||
group.version = version.max(1);
|
||||
}
|
||||
if let Some(published_at) = patch.published_at {
|
||||
group.published_at = published_at;
|
||||
}
|
||||
group.updated_at = patch.updated_at;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) fn apply_binding_patch(
|
||||
binding: &mut StoredRoutingGroupBinding,
|
||||
patch: UpdateRoutingGroupBindingRecord,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if let Some(group_id) = patch.group_id {
|
||||
if group_id.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"routing_group_bindings.group_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
binding.group_id = group_id;
|
||||
}
|
||||
if let Some(subject_type) = patch.subject_type {
|
||||
binding.subject_type = subject_type;
|
||||
}
|
||||
if let Some(subject_id) = patch.subject_id {
|
||||
if subject_id.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidInput(
|
||||
"routing_group_bindings.subject_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
binding.subject_id = subject_id;
|
||||
}
|
||||
if let Some(is_default) = patch.is_default {
|
||||
binding.is_default = is_default;
|
||||
}
|
||||
if let Some(allow_explicit_select) = patch.allow_explicit_select {
|
||||
binding.allow_explicit_select = allow_explicit_select;
|
||||
}
|
||||
binding.updated_at = patch.updated_at;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub(super) fn binding_subject_to_database(subject: RoutingGroupBindingSubject) -> &'static str {
|
||||
match subject {
|
||||
RoutingGroupBindingSubject::User => "user",
|
||||
RoutingGroupBindingSubject::ApiKey => "api_key",
|
||||
RoutingGroupBindingSubject::UserGroup => "user_group",
|
||||
}
|
||||
}
|
||||
|
||||
pub(super) fn binding_subject_from_database(
|
||||
value: String,
|
||||
) -> Result<RoutingGroupBindingSubject, DataLayerError> {
|
||||
match value.as_str() {
|
||||
"user" => Ok(RoutingGroupBindingSubject::User),
|
||||
"api_key" => Ok(RoutingGroupBindingSubject::ApiKey),
|
||||
"user_group" => Ok(RoutingGroupBindingSubject::UserGroup),
|
||||
_ => Err(DataLayerError::UnexpectedValue(format!(
|
||||
"invalid routing_group_bindings.subject_type: {value}"
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
fn map_group_row(row: &PgRow) -> Result<StoredRoutingGroup, DataLayerError> {
|
||||
Ok(StoredRoutingGroup {
|
||||
id: row.try_get("id").map_postgres_err()?,
|
||||
name: row.try_get("name").map_postgres_err()?,
|
||||
description: row.try_get("description").map_postgres_err()?,
|
||||
enabled: row.try_get("enabled").map_postgres_err()?,
|
||||
is_system_default: row.try_get("is_system_default").map_postgres_err()?,
|
||||
config_json: row.try_get("config_json").map_postgres_err()?,
|
||||
version: row.try_get("version").map_postgres_err()?,
|
||||
created_at: row.try_get("created_at").map_postgres_err()?,
|
||||
updated_at: row.try_get("updated_at").map_postgres_err()?,
|
||||
published_at: row.try_get("published_at").map_postgres_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_binding_row(row: &PgRow) -> Result<StoredRoutingGroupBinding, DataLayerError> {
|
||||
Ok(StoredRoutingGroupBinding {
|
||||
id: row.try_get("id").map_postgres_err()?,
|
||||
group_id: row.try_get("group_id").map_postgres_err()?,
|
||||
subject_type: binding_subject_from_database(
|
||||
row.try_get("subject_type").map_postgres_err()?,
|
||||
)?,
|
||||
subject_id: row.try_get("subject_id").map_postgres_err()?,
|
||||
is_default: row.try_get("is_default").map_postgres_err()?,
|
||||
allow_explicit_select: row.try_get("allow_explicit_select").map_postgres_err()?,
|
||||
created_at: row.try_get("created_at").map_postgres_err()?,
|
||||
updated_at: row.try_get("updated_at").map_postgres_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_version_row(row: &PgRow) -> Result<StoredRoutingGroupVersion, DataLayerError> {
|
||||
Ok(StoredRoutingGroupVersion {
|
||||
id: row.try_get("id").map_postgres_err()?,
|
||||
group_id: row.try_get("group_id").map_postgres_err()?,
|
||||
version: row.try_get("version").map_postgres_err()?,
|
||||
config_json: row.try_get("config_json").map_postgres_err()?,
|
||||
created_at: row.try_get("created_at").map_postgres_err()?,
|
||||
created_by: row.try_get("created_by").map_postgres_err()?,
|
||||
})
|
||||
}
|
||||
524
crates/aether-data/src/repository/routing_profiles/sqlite.rs
Normal file
524
crates/aether-data/src/repository/routing_profiles/sqlite.rs
Normal file
@@ -0,0 +1,524 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
use sqlx::{sqlite::SqliteRow, Row};
|
||||
|
||||
use super::postgres::{
|
||||
apply_binding_patch, apply_group_patch, binding_subject_from_database,
|
||||
binding_subject_to_database,
|
||||
};
|
||||
use super::*;
|
||||
use crate::driver::sqlite::SqlitePool;
|
||||
use crate::error::SqlResultExt;
|
||||
use crate::DataLayerError;
|
||||
|
||||
const ROUTING_GROUP_SELECT: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
name,
|
||||
description,
|
||||
enabled,
|
||||
is_system_default,
|
||||
config_json,
|
||||
version,
|
||||
created_at,
|
||||
updated_at,
|
||||
published_at
|
||||
FROM routing_groups
|
||||
"#;
|
||||
|
||||
const ROUTING_GROUP_BINDING_SELECT: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
group_id,
|
||||
subject_type,
|
||||
subject_id,
|
||||
is_default,
|
||||
allow_explicit_select,
|
||||
created_at,
|
||||
updated_at
|
||||
FROM routing_group_bindings
|
||||
"#;
|
||||
|
||||
const ROUTING_GROUP_VERSION_SELECT: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
group_id,
|
||||
version,
|
||||
config_json,
|
||||
created_at,
|
||||
created_by
|
||||
FROM routing_group_versions
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqliteRoutingGroupRepository {
|
||||
pool: SqlitePool,
|
||||
}
|
||||
|
||||
impl SqliteRoutingGroupRepository {
|
||||
pub fn new(pool: SqlitePool) -> Self {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
async fn reload_group(&self, id: &str) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
self.find_routing_group(RoutingGroupLookupKey::Id(id)).await
|
||||
}
|
||||
|
||||
async fn find_binding_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
let row = sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_BINDING_SELECT} WHERE id = ? LIMIT 1"
|
||||
))
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
row.as_ref().map(map_binding_row).transpose()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RoutingGroupReadRepository for SqliteRoutingGroupRepository {
|
||||
async fn list_routing_groups(&self) -> Result<Vec<StoredRoutingGroup>, DataLayerError> {
|
||||
let rows = sqlx::query(&format!("{ROUTING_GROUP_SELECT} ORDER BY name ASC, id ASC"))
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_group_row).collect()
|
||||
}
|
||||
|
||||
async fn find_routing_group(
|
||||
&self,
|
||||
lookup: RoutingGroupLookupKey<'_>,
|
||||
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
let row = match lookup {
|
||||
RoutingGroupLookupKey::Id(id) => sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_SELECT} WHERE id = ? LIMIT 1"
|
||||
))
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?,
|
||||
RoutingGroupLookupKey::Name(name) => sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_SELECT} WHERE name = ? LIMIT 1"
|
||||
))
|
||||
.bind(name)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?,
|
||||
RoutingGroupLookupKey::SystemDefault => sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_SELECT} WHERE is_system_default = 1 AND enabled = 1 ORDER BY updated_at DESC, id ASC LIMIT 1"
|
||||
))
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?,
|
||||
};
|
||||
row.as_ref().map(map_group_row).transpose()
|
||||
}
|
||||
|
||||
async fn list_routing_group_bindings(
|
||||
&self,
|
||||
query: &RoutingGroupBindingQuery,
|
||||
) -> Result<Vec<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
let rows = sqlx::query(&format!(
|
||||
r#"
|
||||
{ROUTING_GROUP_BINDING_SELECT}
|
||||
WHERE (? IS NULL OR group_id = ?)
|
||||
AND (? IS NULL OR subject_type = ?)
|
||||
AND (? IS NULL OR subject_id = ?)
|
||||
ORDER BY created_at ASC, id ASC
|
||||
"#
|
||||
))
|
||||
.bind(query.group_id.as_deref())
|
||||
.bind(query.group_id.as_deref())
|
||||
.bind(query.subject_type.map(binding_subject_to_database))
|
||||
.bind(query.subject_type.map(binding_subject_to_database))
|
||||
.bind(query.subject_id.as_deref())
|
||||
.bind(query.subject_id.as_deref())
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_binding_row).collect()
|
||||
}
|
||||
|
||||
async fn list_routing_group_versions(
|
||||
&self,
|
||||
group_id: &str,
|
||||
) -> Result<Vec<StoredRoutingGroupVersion>, DataLayerError> {
|
||||
let rows = sqlx::query(&format!(
|
||||
"{ROUTING_GROUP_VERSION_SELECT} WHERE group_id = ? ORDER BY version DESC, created_at DESC, id ASC"
|
||||
))
|
||||
.bind(group_id)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
rows.iter().map(map_version_row).collect()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RoutingGroupWriteRepository for SqliteRoutingGroupRepository {
|
||||
async fn create_routing_group(
|
||||
&self,
|
||||
record: CreateRoutingGroupRecord,
|
||||
) -> Result<StoredRoutingGroup, DataLayerError> {
|
||||
let group = StoredRoutingGroup::new(record)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO routing_groups (
|
||||
id, name, description, enabled, is_system_default, config_json,
|
||||
version, created_at, updated_at, published_at
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(&group.id)
|
||||
.bind(&group.name)
|
||||
.bind(&group.description)
|
||||
.bind(group.enabled)
|
||||
.bind(group.is_system_default)
|
||||
.bind(json_to_string(
|
||||
&group.config_json,
|
||||
"routing_groups.config_json",
|
||||
)?)
|
||||
.bind(group.version)
|
||||
.bind(group.created_at)
|
||||
.bind(group.updated_at)
|
||||
.bind(group.published_at)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(group)
|
||||
}
|
||||
|
||||
async fn update_routing_group(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupRecord,
|
||||
) -> Result<Option<StoredRoutingGroup>, DataLayerError> {
|
||||
let Some(mut group) = self.reload_group(id).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
apply_group_patch(&mut group, patch)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE routing_groups
|
||||
SET name = ?,
|
||||
description = ?,
|
||||
enabled = ?,
|
||||
is_system_default = ?,
|
||||
config_json = ?,
|
||||
version = ?,
|
||||
updated_at = ?,
|
||||
published_at = ?
|
||||
WHERE id = ?
|
||||
"#,
|
||||
)
|
||||
.bind(&group.name)
|
||||
.bind(&group.description)
|
||||
.bind(group.enabled)
|
||||
.bind(group.is_system_default)
|
||||
.bind(json_to_string(
|
||||
&group.config_json,
|
||||
"routing_groups.config_json",
|
||||
)?)
|
||||
.bind(group.version)
|
||||
.bind(group.updated_at)
|
||||
.bind(group.published_at)
|
||||
.bind(id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(Some(group))
|
||||
}
|
||||
|
||||
async fn delete_routing_group(&self, id: &str) -> Result<bool, DataLayerError> {
|
||||
let mut tx = self.pool.begin().await.map_sql_err()?;
|
||||
sqlx::query("DELETE FROM routing_group_bindings WHERE group_id = ?")
|
||||
.bind(id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
sqlx::query("DELETE FROM routing_group_versions WHERE group_id = ?")
|
||||
.bind(id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
let rows_affected = sqlx::query("DELETE FROM routing_groups WHERE id = ?")
|
||||
.bind(id)
|
||||
.execute(&mut *tx)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected();
|
||||
tx.commit().await.map_sql_err()?;
|
||||
Ok(rows_affected > 0)
|
||||
}
|
||||
|
||||
async fn create_routing_group_binding(
|
||||
&self,
|
||||
record: CreateRoutingGroupBindingRecord,
|
||||
) -> Result<StoredRoutingGroupBinding, DataLayerError> {
|
||||
let binding = StoredRoutingGroupBinding::new(record)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO routing_group_bindings (
|
||||
id, group_id, subject_type, subject_id, is_default,
|
||||
allow_explicit_select, created_at, updated_at
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(&binding.id)
|
||||
.bind(&binding.group_id)
|
||||
.bind(binding_subject_to_database(binding.subject_type))
|
||||
.bind(&binding.subject_id)
|
||||
.bind(binding.is_default)
|
||||
.bind(binding.allow_explicit_select)
|
||||
.bind(binding.created_at)
|
||||
.bind(binding.updated_at)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(binding)
|
||||
}
|
||||
|
||||
async fn delete_routing_group_binding(&self, id: &str) -> Result<bool, DataLayerError> {
|
||||
Ok(
|
||||
sqlx::query("DELETE FROM routing_group_bindings WHERE id = ?")
|
||||
.bind(id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
.rows_affected()
|
||||
> 0,
|
||||
)
|
||||
}
|
||||
|
||||
async fn update_routing_group_binding(
|
||||
&self,
|
||||
id: &str,
|
||||
patch: UpdateRoutingGroupBindingRecord,
|
||||
) -> Result<Option<StoredRoutingGroupBinding>, DataLayerError> {
|
||||
let Some(mut binding) = self.find_binding_by_id(id).await? else {
|
||||
return Ok(None);
|
||||
};
|
||||
apply_binding_patch(&mut binding, patch)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE routing_group_bindings
|
||||
SET group_id = ?,
|
||||
subject_type = ?,
|
||||
subject_id = ?,
|
||||
is_default = ?,
|
||||
allow_explicit_select = ?,
|
||||
updated_at = ?
|
||||
WHERE id = ?
|
||||
"#,
|
||||
)
|
||||
.bind(&binding.group_id)
|
||||
.bind(binding_subject_to_database(binding.subject_type))
|
||||
.bind(&binding.subject_id)
|
||||
.bind(binding.is_default)
|
||||
.bind(binding.allow_explicit_select)
|
||||
.bind(binding.updated_at)
|
||||
.bind(id)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(Some(binding))
|
||||
}
|
||||
|
||||
async fn create_routing_group_version(
|
||||
&self,
|
||||
record: CreateRoutingGroupVersionRecord,
|
||||
) -> Result<StoredRoutingGroupVersion, DataLayerError> {
|
||||
let version = StoredRoutingGroupVersion::new(record)?;
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO routing_group_versions (
|
||||
id, group_id, version, config_json, created_at, created_by
|
||||
)
|
||||
VALUES (?, ?, ?, ?, ?, ?)
|
||||
"#,
|
||||
)
|
||||
.bind(&version.id)
|
||||
.bind(&version.group_id)
|
||||
.bind(version.version)
|
||||
.bind(json_to_string(
|
||||
&version.config_json,
|
||||
"routing_group_versions.config_json",
|
||||
)?)
|
||||
.bind(version.created_at)
|
||||
.bind(&version.created_by)
|
||||
.execute(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
Ok(version)
|
||||
}
|
||||
}
|
||||
|
||||
fn map_group_row(row: &SqliteRow) -> Result<StoredRoutingGroup, DataLayerError> {
|
||||
Ok(StoredRoutingGroup {
|
||||
id: row.try_get("id").map_sql_err()?,
|
||||
name: row.try_get("name").map_sql_err()?,
|
||||
description: row.try_get("description").map_sql_err()?,
|
||||
enabled: row.try_get("enabled").map_sql_err()?,
|
||||
is_system_default: row.try_get("is_system_default").map_sql_err()?,
|
||||
config_json: json_from_string(
|
||||
row.try_get("config_json").map_sql_err()?,
|
||||
"routing_groups.config_json",
|
||||
)?,
|
||||
version: row.try_get("version").map_sql_err()?,
|
||||
created_at: row.try_get("created_at").map_sql_err()?,
|
||||
updated_at: row.try_get("updated_at").map_sql_err()?,
|
||||
published_at: row.try_get("published_at").map_sql_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_binding_row(row: &SqliteRow) -> Result<StoredRoutingGroupBinding, DataLayerError> {
|
||||
Ok(StoredRoutingGroupBinding {
|
||||
id: row.try_get("id").map_sql_err()?,
|
||||
group_id: row.try_get("group_id").map_sql_err()?,
|
||||
subject_type: binding_subject_from_database(row.try_get("subject_type").map_sql_err()?)?,
|
||||
subject_id: row.try_get("subject_id").map_sql_err()?,
|
||||
is_default: row.try_get("is_default").map_sql_err()?,
|
||||
allow_explicit_select: row.try_get("allow_explicit_select").map_sql_err()?,
|
||||
created_at: row.try_get("created_at").map_sql_err()?,
|
||||
updated_at: row.try_get("updated_at").map_sql_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn map_version_row(row: &SqliteRow) -> Result<StoredRoutingGroupVersion, DataLayerError> {
|
||||
Ok(StoredRoutingGroupVersion {
|
||||
id: row.try_get("id").map_sql_err()?,
|
||||
group_id: row.try_get("group_id").map_sql_err()?,
|
||||
version: row.try_get("version").map_sql_err()?,
|
||||
config_json: json_from_string(
|
||||
row.try_get("config_json").map_sql_err()?,
|
||||
"routing_group_versions.config_json",
|
||||
)?,
|
||||
created_at: row.try_get("created_at").map_sql_err()?,
|
||||
created_by: row.try_get("created_by").map_sql_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn json_to_string(value: &Value, field_name: &str) -> Result<String, DataLayerError> {
|
||||
serde_json::to_string(value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("{field_name} contains unserializable JSON: {err}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn json_from_string(value: String, field_name: &str) -> Result<Value, DataLayerError> {
|
||||
serde_json::from_str(&value).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!("{field_name} contains invalid JSON: {err}"))
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::json;
|
||||
|
||||
use super::*;
|
||||
use crate::lifecycle::migrate::run_sqlite_migrations;
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_routing_group_repository_round_trips() {
|
||||
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.expect("sqlite pool should connect");
|
||||
run_sqlite_migrations(&pool)
|
||||
.await
|
||||
.expect("sqlite migrations should run");
|
||||
|
||||
let repository = SqliteRoutingGroupRepository::new(pool);
|
||||
repository
|
||||
.create_routing_group(CreateRoutingGroupRecord {
|
||||
id: "routing-group-1".to_string(),
|
||||
name: "default".to_string(),
|
||||
description: Some("initial".to_string()),
|
||||
enabled: true,
|
||||
is_system_default: true,
|
||||
config_json: json!({"allowed_models": ["gpt-*"]}),
|
||||
version: 1,
|
||||
created_at: 10,
|
||||
updated_at: 10,
|
||||
published_at: None,
|
||||
})
|
||||
.await
|
||||
.expect("group should create");
|
||||
|
||||
let system_default = repository
|
||||
.find_routing_group(RoutingGroupLookupKey::SystemDefault)
|
||||
.await
|
||||
.expect("group lookup should succeed")
|
||||
.expect("system default should exist");
|
||||
assert_eq!(system_default.id, "routing-group-1");
|
||||
|
||||
repository
|
||||
.update_routing_group(
|
||||
"routing-group-1",
|
||||
UpdateRoutingGroupRecord {
|
||||
description: Some(None),
|
||||
version: Some(2),
|
||||
updated_at: 20,
|
||||
published_at: Some(Some(20)),
|
||||
..UpdateRoutingGroupRecord::default()
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect("group should update");
|
||||
|
||||
let binding = repository
|
||||
.create_routing_group_binding(CreateRoutingGroupBindingRecord {
|
||||
id: "binding-1".to_string(),
|
||||
group_id: "routing-group-1".to_string(),
|
||||
subject_type: RoutingGroupBindingSubject::ApiKey,
|
||||
subject_id: "api-key-1".to_string(),
|
||||
is_default: true,
|
||||
allow_explicit_select: true,
|
||||
created_at: 10,
|
||||
updated_at: 10,
|
||||
})
|
||||
.await
|
||||
.expect("binding should create");
|
||||
|
||||
assert_eq!(binding.subject_type, RoutingGroupBindingSubject::ApiKey);
|
||||
assert_eq!(
|
||||
repository
|
||||
.list_routing_group_bindings(&RoutingGroupBindingQuery {
|
||||
group_id: Some("routing-group-1".to_string()),
|
||||
subject_type: Some(RoutingGroupBindingSubject::ApiKey),
|
||||
subject_id: Some("api-key-1".to_string()),
|
||||
})
|
||||
.await
|
||||
.expect("bindings should list")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
|
||||
repository
|
||||
.create_routing_group_version(CreateRoutingGroupVersionRecord {
|
||||
id: "version-1".to_string(),
|
||||
group_id: "routing-group-1".to_string(),
|
||||
version: 2,
|
||||
config_json: json!({"allowed_models": ["gpt-*"]}),
|
||||
created_at: 20,
|
||||
created_by: Some("admin".to_string()),
|
||||
})
|
||||
.await
|
||||
.expect("version should create");
|
||||
|
||||
assert_eq!(
|
||||
repository
|
||||
.list_routing_group_versions("routing-group-1")
|
||||
.await
|
||||
.expect("versions should list")
|
||||
.len(),
|
||||
1
|
||||
);
|
||||
}
|
||||
}
|
||||
13
crates/aether-routing-core/Cargo.toml
Normal file
13
crates/aether-routing-core/Cargo.toml
Normal file
@@ -0,0 +1,13 @@
|
||||
[package]
|
||||
name = "aether-routing-core"
|
||||
version = "0.1.0"
|
||||
edition.workspace = true
|
||||
license.workspace = true
|
||||
repository.workspace = true
|
||||
description = "Pure routing profile policy, mutation, ranking overlay, and trace primitives for Aether"
|
||||
|
||||
[dependencies]
|
||||
regex.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
thiserror.workspace = true
|
||||
91
crates/aether-routing-core/src/actions.rs
Normal file
91
crates/aether-routing-core/src/actions.rs
Normal 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>,
|
||||
},
|
||||
}
|
||||
232
crates/aether-routing-core/src/conditions.rs
Normal file
232
crates/aether-routing-core/src/conditions.rs
Normal 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"}))));
|
||||
}
|
||||
}
|
||||
36
crates/aether-routing-core/src/lib.rs
Normal file
36
crates/aether-routing-core/src/lib.rs
Normal 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};
|
||||
131
crates/aether-routing-core/src/model.rs
Normal file
131
crates/aether-routing-core/src/model.rs
Normal 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
|
||||
}
|
||||
236
crates/aether-routing-core/src/mutations.rs
Normal file
236
crates/aether-routing-core/src/mutations.rs
Normal 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()))
|
||||
);
|
||||
}
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user