mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-04 16:37:46 +08:00
892 lines
32 KiB
Rust
892 lines
32 KiB
Rust
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::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 {
|
|
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
|
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
|
pub(crate) required_capabilities: Option<serde_json::Value>,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub(crate) struct LocalRequestedModelDecisionInput {
|
|
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
|
pub(crate) requested_model: String,
|
|
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
|
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)]
|
|
pub(crate) struct LocalAuthenticatedDecisionInput {
|
|
pub(crate) auth_context: ExecutionRuntimeAuthContext,
|
|
pub(crate) auth_snapshot: GatewayAuthApiKeySnapshot,
|
|
pub(crate) required_capabilities: Option<serde_json::Value>,
|
|
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,
|
|
}
|
|
|
|
#[async_trait]
|
|
impl AiAuthenticatedDecisionInputPort for GatewayAuthenticatedDecisionInputPort<'_> {
|
|
type AuthContext = ExecutionRuntimeAuthContext;
|
|
type AuthSnapshot = GatewayAuthApiKeySnapshot;
|
|
type RequiredCapabilities = serde_json::Value;
|
|
type ResolvedInput = ResolvedLocalDecisionAuthInput;
|
|
type Error = GatewayError;
|
|
|
|
async fn read_auth_snapshot(
|
|
&self,
|
|
auth_context: &Self::AuthContext,
|
|
) -> Result<Option<Self::AuthSnapshot>, Self::Error> {
|
|
self.state
|
|
.read_auth_api_key_snapshot(
|
|
&auth_context.user_id,
|
|
&auth_context.api_key_id,
|
|
self.now_unix_secs,
|
|
)
|
|
.await
|
|
}
|
|
|
|
async fn resolve_required_capabilities(
|
|
&self,
|
|
auth_context: &Self::AuthContext,
|
|
requested_model: Option<&str>,
|
|
explicit_required_capabilities: Option<&Self::RequiredCapabilities>,
|
|
) -> Result<Option<Self::RequiredCapabilities>, Self::Error> {
|
|
Ok(self
|
|
.state
|
|
.resolve_request_candidate_required_capabilities(
|
|
&auth_context.user_id,
|
|
&auth_context.api_key_id,
|
|
requested_model,
|
|
explicit_required_capabilities,
|
|
)
|
|
.await)
|
|
}
|
|
|
|
fn build_resolved_input(
|
|
&self,
|
|
auth_context: Self::AuthContext,
|
|
auth_snapshot: Self::AuthSnapshot,
|
|
required_capabilities: Option<Self::RequiredCapabilities>,
|
|
) -> Self::ResolvedInput {
|
|
ResolvedLocalDecisionAuthInput {
|
|
auth_context,
|
|
auth_snapshot,
|
|
required_capabilities,
|
|
}
|
|
}
|
|
}
|
|
|
|
pub(crate) fn build_local_requested_model_decision_input(
|
|
resolved_input: ResolvedLocalDecisionAuthInput,
|
|
requested_model: String,
|
|
) -> LocalRequestedModelDecisionInput {
|
|
LocalRequestedModelDecisionInput {
|
|
auth_context: resolved_input.auth_context,
|
|
requested_model,
|
|
auth_snapshot: resolved_input.auth_snapshot,
|
|
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 {
|
|
LocalAuthenticatedDecisionInput {
|
|
auth_context: resolved_input.auth_context,
|
|
auth_snapshot: resolved_input.auth_snapshot,
|
|
required_capabilities: resolved_input.required_capabilities,
|
|
client_session_affinity: None,
|
|
}
|
|
}
|
|
|
|
pub(crate) async fn resolve_local_authenticated_decision_input(
|
|
state: &AppState,
|
|
auth_context: ExecutionRuntimeAuthContext,
|
|
requested_model: Option<&str>,
|
|
explicit_required_capabilities: Option<&serde_json::Value>,
|
|
) -> Result<Option<ResolvedLocalDecisionAuthInput>, GatewayError> {
|
|
let port = GatewayAuthenticatedDecisionInputPort {
|
|
state: PlannerAppState::new(state),
|
|
now_unix_secs: current_unix_secs(),
|
|
};
|
|
|
|
run_ai_authenticated_decision_input(
|
|
&port,
|
|
auth_context,
|
|
requested_model,
|
|
explicit_required_capabilities,
|
|
)
|
|
.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,
|
|
api_key_ip_rules: 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()),
|
|
content_encoding: None,
|
|
request_gzip: None,
|
|
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)
|
|
);
|
|
}
|
|
}
|