Merge remote-tracking branch 'upstream/main' into feat/356-usage-record-columns

This commit is contained in:
RWDai
2026-05-18 17:49:35 +08:00
187 changed files with 15067 additions and 1394 deletions

View File

@@ -58,6 +58,11 @@ ADMIN_USERNAME=admin
# AETHER_GATEWAY_AUTO_PREPARE_DATABASE=true
# PostgreSQL 连接池配置(默认适合单实例/小型部署;高并发可按需调大)
# 推荐计算方式(单实例):
# MAX = CPU 核数 × 10AI 网关偏 IO 等待,可激进些;纯 OLTP 用 × 4
# MIN = MAX × 0.2(保留常驻连接应对突发流量,避免冷启动握手开销)
# 多实例部署时请按 实例数 × MAX 控制总和PG 端 max_connections 至少为该总和 + 20 余量
# AETHER_GATEWAY_DATA_POSTGRES_MIN_CONNECTIONS=1
# AETHER_GATEWAY_DATA_POSTGRES_MAX_CONNECTIONS=20
# AETHER_GATEWAY_DATA_POSTGRES_IDLE_TIMEOUT_MS=30000
# AETHER_GATEWAY_DATA_POSTGRES_STATEMENT_CACHE_CAPACITY=100
# AETHER_GATEWAY_DATA_POSTGRES_ACQUIRE_TIMEOUT_MS=3000

View File

@@ -247,8 +247,10 @@ jobs:
set -euo pipefail
if [[ "${GITHUB_REF_TYPE}" == "tag" ]]; then
VERSION="${GITHUB_REF_NAME}"
SOURCE_REF="${GITHUB_REF_NAME}"
else
VERSION="snapshot-${GITHUB_SHA::7}"
SOURCE_REF="${GITHUB_SHA}"
fi
mkdir -p package release-assets
@@ -262,9 +264,13 @@ jobs:
install -m 0755 "artifacts/aether-gateway-${platform}-${arch}/aether-gateway" "${root}/bin/aether-gateway"
cp -R artifacts/frontend-dist/. "${root}/frontend/"
sed "s/^VERSION=\"\${AETHER_VERSION:-}\"/VERSION=\"\${AETHER_VERSION:-${VERSION}}\"/" install.sh > "${root}/install.sh"
sed \
-e "s/^SOURCE_REF=\"\${AETHER_SOURCE_REF:-main}\"/SOURCE_REF=\"\${AETHER_SOURCE_REF:-${SOURCE_REF}}\"/" \
-e "s/^VERSION=\"\${AETHER_VERSION:-}\"/VERSION=\"\${AETHER_VERSION:-${VERSION}}\"/" \
install.sh > "${root}/install.sh"
chmod 0755 "${root}/install.sh"
install -m 0644 docker-compose.yml "${root}/docker-compose.yml"
install -m 0644 docker-compose.sqlite.yml "${root}/docker-compose.sqlite.yml"
install -m 0644 .env.example "${root}/.env.example"
install -m 0755 generate_keys.sh "${root}/generate_keys.sh"
install -m 0644 README.md "${root}/README.md"
@@ -274,7 +280,10 @@ jobs:
done
done
sed "s/^VERSION=\"\${AETHER_VERSION:-}\"/VERSION=\"\${AETHER_VERSION:-${VERSION}}\"/" install.sh > release-assets/install.sh
sed \
-e "s/^SOURCE_REF=\"\${AETHER_SOURCE_REF:-main}\"/SOURCE_REF=\"\${AETHER_SOURCE_REF:-${SOURCE_REF}}\"/" \
-e "s/^VERSION=\"\${AETHER_VERSION:-}\"/VERSION=\"\${AETHER_VERSION:-${VERSION}}\"/" \
install.sh > release-assets/install.sh
chmod +x release-assets/install.sh
(cd release-assets && sha256sum *.tar.gz > SHA256SUMS)

12
Cargo.lock generated
View File

@@ -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"

View File

@@ -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" }

View File

@@ -1,11 +1,14 @@
# syntax=docker/dockerfile:1
# Aether 运行镜像Rust gateway 直接服务 API + 前端静态文件(国内镜像源版本)
# 构建命令: docker build -f Dockerfile.app.local -t aether-app:latest .
# 构建命令: docker build --build-arg AETHER_BUILD_VERSION=v0.7.2 -f Dockerfile.app.local -t aether-app:latest .
ARG RUST_VERSION=1.95.0
# ==================== 前端构建 ====================
FROM node:22-slim AS frontend-builder
ARG AETHER_BUILD_VERSION
ENV AETHER_BUILD_VERSION=${AETHER_BUILD_VERSION} \
AETHER_VERSION=${AETHER_BUILD_VERSION}
WORKDIR /app/frontend
COPY frontend/package*.json ./
RUN --mount=type=cache,id=aether-npm-cache,target=/root/.npm,sharing=locked \
@@ -44,6 +47,9 @@ COPY crates/ ./crates/
RUN cargo chef prepare --recipe-path recipe.json
FROM gateway-base AS gateway-builder
ARG AETHER_BUILD_VERSION
ENV AETHER_BUILD_VERSION=${AETHER_BUILD_VERSION} \
AETHER_VERSION=${AETHER_BUILD_VERSION}
COPY --from=gateway-planner /build/recipe.json ./recipe.json
RUN --mount=type=cache,id=aether-cargo-registry,target=/usr/local/cargo/registry,sharing=locked \
--mount=type=cache,id=aether-cargo-git,target=/usr/local/cargo/git,sharing=locked \

View File

@@ -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

View File

@@ -2,14 +2,20 @@ use std::env;
use std::process::Command;
fn main() {
println!("cargo:rerun-if-env-changed=AETHER_BUILD_VERSION");
println!("cargo:rerun-if-env-changed=AETHER_VERSION");
println!("cargo:rerun-if-env-changed=GITHUB_REF_NAME");
println!("cargo:rerun-if-changed=../../.git/HEAD");
let package_version = env::var("CARGO_PKG_VERSION").unwrap_or_else(|_| "unknown".to_string());
let version = env::var("AETHER_VERSION")
let version = env::var("AETHER_BUILD_VERSION")
.ok()
.filter(|value| !value.trim().is_empty())
.or_else(|| {
env::var("AETHER_VERSION")
.ok()
.filter(|value| !value.trim().is_empty())
})
.or_else(|| {
env::var("GITHUB_REF_NAME")
.ok()

View File

@@ -3,10 +3,9 @@ pub(crate) use crate::handlers::admin::{
build_internal_control_error_response, create_provider_oauth_catalog_key,
find_duplicate_provider_oauth_key, maybe_build_local_admin_pool_response,
maybe_build_local_admin_response, persist_provider_quota_refresh_state,
provider_account_self_check_endpoint_for_provider,
provider_oauth_maintenance_endpoint_for_provider, provider_oauth_runtime_endpoint_for_provider,
provider_quota_refresh_endpoint_for_provider, provider_type_supports_account_self_check,
provider_type_supports_quota_refresh, reconcile_admin_fixed_provider_template_endpoints,
provider_quota_refresh_endpoint_for_provider, provider_type_supports_quota_refresh,
reconcile_admin_fixed_provider_template_endpoints,
refresh_provider_oauth_account_state_after_update, refresh_provider_pool_quota_locally,
update_existing_provider_oauth_catalog_key, AdminAppState,
AdminGatewayProviderTransportSnapshot, AdminLocalOAuthRefreshError, AdminRequestContext,

View File

@@ -49,10 +49,10 @@ pub(crate) use aether_ai_formats::api::{
resolve_claude_sync_spec, resolve_gemini_stream_spec, resolve_gemini_sync_spec,
resolve_local_image_stream_spec, resolve_local_image_sync_spec,
resolve_local_same_format_stream_spec, resolve_local_same_format_sync_spec,
AiControlPlanRequest, ExecutionRuntimeAuthContext, LocalCoreSyncErrorKind,
LocalOpenAiImageSpec, LocalSameFormatProviderFamily, LocalSameFormatProviderSpec,
LocalStandardSourceFamily, LocalStandardSourceMode, LocalStandardSpec,
StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
sanitize_request_path_and_query, AiControlPlanRequest, ExecutionRuntimeAuthContext,
LocalCoreSyncErrorKind, LocalOpenAiImageSpec, LocalSameFormatProviderFamily,
LocalSameFormatProviderSpec, LocalStandardSourceFamily, LocalStandardSourceMode,
LocalStandardSpec, StreamingStandardTerminalObserver, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION, GEMINI_FILES_DOWNLOAD_PLAN_KIND,
GEMINI_VIDEO_CANCEL_SYNC_PLAN_KIND, OPENAI_EMBEDDING_SYNC_PLAN_KIND,
OPENAI_IMAGE_STREAM_PLAN_KIND, OPENAI_IMAGE_SYNC_FINALIZE_REPORT_KIND,

View File

@@ -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,

View File

@@ -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;

View File

@@ -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,

View File

@@ -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,
)

View File

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

View File

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

View File

@@ -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,

View File

@@ -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(

View File

@@ -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(),

View File

@@ -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;
};

View File

@@ -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;
};

View File

@@ -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))
}

View File

@@ -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(),

View File

@@ -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(),

View File

@@ -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;
};

View File

@@ -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))
}

View File

@@ -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(),

View File

@@ -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,

View File

@@ -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;
};

View File

@@ -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))
}

View File

@@ -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(),

View File

@@ -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,

View File

@@ -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;
};

View File

@@ -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,

View File

@@ -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!(

View File

@@ -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(),

View File

@@ -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))
}

View File

@@ -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(),

View File

@@ -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,

View File

@@ -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,

View File

@@ -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,

View File

@@ -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))
}

View File

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

View File

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

View File

@@ -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))
}

View File

@@ -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(),

View File

@@ -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,

View File

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

View File

@@ -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;
};

View File

@@ -127,15 +127,12 @@ async fn balance_capacity_rejection(
let wallet_is_unlimited = wallet
.as_ref()
.is_some_and(|wallet| wallet.limit_mode.eq_ignore_ascii_case("unlimited"));
let (available_usd, require_cost_estimate) = match quota.as_ref() {
Some(quota) if !quota.allow_wallet_overage => (Some(quota.remaining_usd.max(0.0)), true),
Some(_) if wallet_is_unlimited => (None, false),
Some(quota) => (
Some(quota.remaining_usd.max(0.0) + wallet_available_usd.unwrap_or(0.0)),
true,
),
None if wallet_is_unlimited => (None, false),
None => (wallet_available_usd, false),
let available_usd = match quota.as_ref() {
Some(quota) if !quota.allow_wallet_overage => Some(quota.remaining_usd.max(0.0)),
Some(_) if wallet_is_unlimited => None,
Some(quota) => Some(quota.remaining_usd.max(0.0) + wallet_available_usd.unwrap_or(0.0)),
None if wallet_is_unlimited => None,
None => wallet_available_usd,
};
let Some(available_usd) = available_usd else {
return Ok(None);
@@ -146,24 +143,12 @@ async fn balance_capacity_rejection(
}));
}
let Some(requested_model) = requested_model else {
return if require_cost_estimate {
Ok(Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: Some(available_usd),
}))
} else {
Ok(None)
};
return Ok(None);
};
let Some(estimated_cost_usd) =
estimate_request_cost_upper_bound_usd(state, decision, requested_model, body).await?
else {
return if require_cost_estimate {
Ok(Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: Some(available_usd),
}))
} else {
Ok(None)
};
return Ok(None);
};
if estimated_cost_usd > available_usd + DAILY_QUOTA_EPSILON_USD {
return Ok(Some(GatewayLocalAuthRejection::BalanceDenied {
@@ -478,10 +463,15 @@ mod tests {
use std::sync::Arc;
use aether_data::repository::candidate_selection::InMemoryMinimalCandidateSelectionReadRepository;
use aether_data_contracts::repository::billing::StoredBillingModelContext;
use aether_data::repository::wallet::StoredWalletSnapshot;
use aether_data_contracts::repository::billing::{
BillingReadRepository, StoredBillingModelContext, UserDailyQuotaAvailabilityRecord,
};
use aether_data_contracts::repository::candidate_selection::{
StoredMinimalCandidateSelectionRow, StoredProviderModelMapping,
};
use aether_data_contracts::DataLayerError;
use async_trait::async_trait;
use axum::body::Bytes;
use axum::http::{HeaderMap, Uri};
use serde_json::json;
@@ -588,6 +578,25 @@ mod tests {
.with_data_state_for_tests(data)
}
fn state_with_quota_and_wallet(
quota: UserDailyQuotaAvailabilityRecord,
context: StoredBillingModelContext,
) -> AppState {
let candidate_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_row(),
]));
let billing_repository = Arc::new(FixedBillingReadRepository { quota, context });
let data = GatewayDataState::with_minimal_candidate_selection_and_billing_for_tests(
candidate_repository,
billing_repository,
);
AppState::new()
.expect("state should build")
.with_data_state_for_tests(data)
.with_auth_wallets_for_tests(vec![sample_wallet("user-1", 30.0)])
}
fn state_with_model_mapping() -> AppState {
state_with_rows(vec![sample_row()])
}
@@ -629,6 +638,72 @@ mod tests {
.expect("billing context should build")
}
fn sample_wallet(user_id: &str, balance: f64) -> StoredWalletSnapshot {
StoredWalletSnapshot::new(
format!("wallet-{user_id}"),
Some(user_id.to_string()),
None,
balance,
0.0,
"finite".to_string(),
"USD".to_string(),
"active".to_string(),
balance,
0.0,
0.0,
0.0,
100,
)
.expect("wallet should build")
}
fn quota_availability(
remaining_usd: f64,
allow_wallet_overage: bool,
) -> UserDailyQuotaAvailabilityRecord {
UserDailyQuotaAvailabilityRecord {
has_active_daily_quota: true,
total_quota_usd: remaining_usd,
used_usd: 0.0,
remaining_usd,
allow_wallet_overage,
}
}
#[derive(Debug)]
struct FixedBillingReadRepository {
quota: UserDailyQuotaAvailabilityRecord,
context: StoredBillingModelContext,
}
#[async_trait]
impl BillingReadRepository for FixedBillingReadRepository {
async fn find_model_context(
&self,
_provider_id: &str,
_provider_api_key_id: Option<&str>,
_global_model_name: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
Ok(Some(self.context.clone()))
}
async fn find_model_context_by_model_id(
&self,
_provider_id: &str,
_provider_api_key_id: Option<&str>,
_model_id: &str,
) -> Result<Option<StoredBillingModelContext>, DataLayerError> {
Ok(Some(self.context.clone()))
}
async fn find_user_daily_quota_availability(
&self,
_user_id: &str,
) -> Result<Option<UserDailyQuotaAvailabilityRecord>, DataLayerError> {
Ok(Some(self.quota.clone()))
}
}
#[tokio::test]
async fn model_rejection_allows_requested_model_that_resolves_to_allowed_global_model() {
let state = state_with_model_mapping();
@@ -702,6 +777,108 @@ mod tests {
);
}
#[tokio::test]
async fn positive_balance_allows_unbounded_output_request_without_cost_estimate() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 1.0,
"output_price_per_1m": 2.0
}]
})),
None,
None,
None,
);
for allow_wallet_overage in [false, true] {
let state = state_with_quota_and_wallet(
quota_availability(50.0, allow_wallet_overage),
context.clone(),
);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let body = Bytes::from_static(
br#"{"model":"gpt-5","messages":[{"role":"user","content":"hi"}],"stream":true}"#,
);
let rejection = request_model_local_rejection(
&state,
Some(&decision),
&uri,
&json_headers(),
&body,
)
.await
.expect("quota rejection should resolve");
assert_eq!(rejection, None);
}
}
#[tokio::test]
async fn positive_balance_still_denies_known_cost_above_available_capacity() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 0.0,
"output_price_per_1m": 60.0
}]
})),
None,
None,
None,
);
let state = state_with_quota_and_wallet(quota_availability(50.0, false), context);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let body = Bytes::from_static(
br#"{"model":"gpt-5","messages":[{"role":"user","content":"hi"}],"max_tokens":1000000}"#,
);
let rejection =
request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body)
.await
.expect("quota rejection should resolve");
assert_eq!(
rejection,
Some(GatewayLocalAuthRejection::BalanceDenied {
remaining: Some(50.0),
})
);
}
#[tokio::test]
async fn wallet_overage_policy_extends_known_cost_capacity_when_enabled() {
let context = billing_context_with_pricing(
Some(json!({
"tiers": [{
"up_to": null,
"input_price_per_1m": 0.0,
"output_price_per_1m": 70.0
}]
})),
None,
None,
None,
);
let state = state_with_quota_and_wallet(quota_availability(50.0, true), context);
let decision = decision_with_allowed_models(vec!["gpt-5".to_string()]);
let uri: Uri = "/v1/chat/completions".parse().expect("uri should parse");
let body = Bytes::from_static(
br#"{"model":"gpt-5","messages":[{"role":"user","content":"hi"}],"max_tokens":1000000}"#,
);
let rejection =
request_model_local_rejection(&state, Some(&decision), &uri, &json_headers(), &body)
.await
.expect("quota rejection should resolve");
assert_eq!(rejection, None);
}
#[test]
fn daily_quota_estimate_falls_back_to_default_tiers_when_model_tiers_empty() {
let context = billing_context_with_pricing(

View File

@@ -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",

View File

@@ -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)

View File

@@ -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,
)
}

View 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"
);
}
}

View File

@@ -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;

View File

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

View File

@@ -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;

View 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),
}
}
}

View File

@@ -37,7 +37,8 @@ use super::{
use aether_data_contracts::repository::usage::{
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest,
StoredProviderApiKeyWindowUsageSummary, StoredUsageDailySummary, UsageAuditListQuery,
UsageCleanupSummary, UsageCleanupWindow, UsageDailyHeatmapQuery,
UsageCleanupExecutionMode, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow,
UsageDailyHeatmapQuery,
};
use aether_runtime_state::RuntimeQueueStore;
use aether_video_tasks_core::read_data_backed_video_task_response;
@@ -980,11 +981,13 @@ impl GatewayDataState {
window: &UsageCleanupWindow,
batch_size: usize,
auto_delete_expired_keys: bool,
targets: UsageCleanupTargets,
mode: UsageCleanupExecutionMode,
) -> Result<UsageCleanupSummary, DataLayerError> {
match &self.usage_writer {
Some(repository) => {
repository
.cleanup_usage(window, batch_size, auto_delete_expired_keys)
.cleanup_usage(window, batch_size, auto_delete_expired_keys, targets, mode)
.await
}
None => Ok(UsageCleanupSummary::default()),
@@ -994,10 +997,16 @@ impl GatewayDataState {
pub(crate) async fn preview_usage_cleanup(
&self,
window: &UsageCleanupWindow,
targets: UsageCleanupTargets,
mode: UsageCleanupExecutionMode,
) -> Result<aether_data_contracts::repository::usage::UsageCleanupPreviewCounts, DataLayerError>
{
match &self.usage_writer {
Some(repository) => repository.preview_usage_cleanup(window).await,
Some(repository) => {
repository
.preview_usage_cleanup(window, targets, mode)
.await
}
None => {
Ok(aether_data_contracts::repository::usage::UsageCleanupPreviewCounts::default())
}

View File

@@ -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,

View File

@@ -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,60 @@ 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,
user_preferences: None,
usage_worker_queue: None,
video_task_reader: None,
video_task_writer: None,
background_task_reader: None,
background_task_writer: None,
wallet_reader: None,
wallet_writer: None,
settlement_writer: None,
system_config_values: None,
}
}
#[cfg(test)]
pub(crate) fn with_minimal_candidate_selection_and_billing_for_tests(
candidate_selection_repository: Arc<dyn MinimalCandidateSelectionReadRepository>,
billing_repository: Arc<dyn BillingReadRepository>,
) -> Self {
Self {
config: GatewayDataConfig::disabled(),
backends: None,
auth_api_key_reader: None,
auth_api_key_writer: None,
auth_module_reader: None,
auth_module_writer: None,
announcement_reader: None,
announcement_writer: None,
management_token_reader: None,
management_token_writer: None,
oauth_provider_reader: None,
oauth_provider_writer: None,
proxy_node_reader: None,
proxy_node_writer: None,
billing_reader: Some(billing_repository),
gemini_file_mapping_reader: None,
gemini_file_mapping_writer: None,
global_model_reader: None,
global_model_writer: None,
minimal_candidate_selection_reader: Some(candidate_selection_repository),
request_candidate_reader: None,
request_candidate_writer: None,
provider_catalog_reader: None,
provider_catalog_writer: None,
pool_score_reader: None,
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 +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,
@@ -1857,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,
@@ -1916,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,
@@ -1978,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,
@@ -2046,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,
@@ -2115,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,
@@ -2188,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,
@@ -2268,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,
@@ -2330,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,
@@ -2383,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,
@@ -2432,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,
@@ -2487,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,
@@ -2546,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,
@@ -2606,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,
@@ -2659,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,

View File

@@ -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,

View File

@@ -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,

View File

@@ -1004,6 +1004,24 @@ where
read_next_frame(lines).await
}
async fn next_stream_frame_until_downstream_closed<R>(
buffered_frames: &mut VecDeque<StreamFrame>,
lines: &mut FramedRead<R, LinesCodec>,
tx: &mpsc::Sender<Result<Bytes, IoError>>,
) -> Result<Option<StreamFrame>, GatewayError>
where
R: tokio::io::AsyncRead + Unpin,
{
if let Some(frame) = buffered_frames.pop_front() {
return Ok(Some(frame));
}
tokio::select! {
frame = read_next_frame(lines) => frame,
() = tx.closed() => Ok(None),
}
}
fn should_refresh_stream_usage_telemetry(
previous: Option<&ExecutionTelemetry>,
next: &ExecutionTelemetry,
@@ -2170,7 +2188,15 @@ async fn execute_stream_from_frame_stream(
image_stream_total_timeout.as_mut()
{
tokio::select! {
result = next_stream_frame(&mut buffered_frames, &mut lines) => result,
result = next_stream_frame_until_downstream_closed(
&mut buffered_frames,
&mut lines,
&tx,
) => result,
() = tx.closed() => {
downstream_dropped = true;
break;
}
_ = timeout_sleep.as_mut() => {
let timeout_ms = openai_image_stream_total_timeout_ms
.unwrap_or(OPENAI_IMAGE_STREAM_DEFAULT_TOTAL_TIMEOUT_MS);
@@ -2221,7 +2247,8 @@ async fn execute_stream_from_frame_stream(
}
}
} else {
next_stream_frame(&mut buffered_frames, &mut lines).await
next_stream_frame_until_downstream_closed(&mut buffered_frames, &mut lines, &tx)
.await
};
let next_frame = match next_frame_result {
Ok(frame) => frame,
@@ -2244,6 +2271,9 @@ async fn execute_stream_from_frame_stream(
}
};
let Some(frame) = next_frame else {
if tx.is_closed() {
downstream_dropped = true;
}
break;
};
let frame_elapsed_ms = stream_started_at_for_report
@@ -2461,6 +2491,7 @@ async fn execute_stream_from_frame_stream(
}
if downstream_dropped {
drop(lines);
debug!(
event_name = "execution_runtime_stream_flush_skipped",
log_type = "debug",
@@ -2906,7 +2937,9 @@ mod tests {
};
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data_contracts::repository::candidates::RequestCandidateReadRepository;
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus,
};
use aether_data_contracts::repository::usage::UsageReadRepository;
use aether_usage_runtime::UsageRuntimeConfig;
use async_stream::stream;
@@ -3174,6 +3207,140 @@ mod tests {
assert!(text.contains("\"type\":\"image_stream_total_timeout\""));
}
#[tokio::test]
async fn execute_stream_from_frame_stream_stops_upstream_when_client_drops_body() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let plan = ExecutionPlan {
request_id: "req-client-drop-cancels-upstream".into(),
candidate_id: Some("cand-client-drop-cancels-upstream".into()),
provider_name: Some("openai".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: "https://example.com/v1/chat/completions".into(),
headers: BTreeMap::from([
("content-type".into(), "application/json".into()),
("accept".into(), "text/event-stream".into()),
]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-5.4",
"messages": [],
"stream": true
})),
stream: true,
client_api_format: "openai:chat".into(),
provider_api_format: "openai:chat".into(),
model_name: Some("gpt-5.4".into()),
proxy: None,
transport_profile: None,
timeouts: None,
};
let frame_stream_dropped = Arc::new(Notify::new());
let frame_stream_dropped_for_stream = Arc::clone(&frame_stream_dropped);
let frame_stream = stream! {
struct NotifyOnDrop(Arc<Notify>);
impl Drop for NotifyOnDrop {
fn drop(&mut self) {
self.0.notify_waiters();
}
}
let _drop_guard = NotifyOnDrop(frame_stream_dropped_for_stream);
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
));
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"first\\\"}\\n\\n\"}}\n",
));
std::future::pending::<()>().await;
}
.boxed();
let response = execute_stream_from_frame_stream(
&state,
plan,
"trace-client-drop-cancels-upstream",
&test_decision(),
"openai_chat_stream",
None,
Some(json!({
"request_id": "req-client-drop-cancels-upstream",
"candidate_id": "cand-client-drop-cancels-upstream",
"candidate_index": 0,
"retry_index": 0,
"provider_api_format": "openai:chat",
"client_api_format": "openai:chat"
})),
crate::clock::current_unix_ms(),
Instant::now(),
frame_stream,
None,
)
.await
.expect("execution should succeed")
.expect("execution should return a client response");
let mut body_stream = response.into_body().into_data_stream();
let first = tokio::time::timeout(Duration::from_secs(1), async {
loop {
let chunk = body_stream
.next()
.await
.expect("body should yield first chunk")
.expect("first chunk should be ok");
if chunk.as_ref() != b": aether-keepalive\n\n" {
break chunk;
}
}
})
.await
.expect("first business chunk should arrive");
assert_eq!(first.as_ref(), b"data: {\"id\":\"first\"}\n\n");
drop(body_stream);
tokio::time::timeout(Duration::from_secs(1), frame_stream_dropped.notified())
.await
.expect("upstream frame stream should be dropped after client disconnect");
let candidates = tokio::time::timeout(Duration::from_secs(1), async {
loop {
let candidates = request_candidate_repository
.list_by_request_id("req-client-drop-cancels-upstream")
.await
.expect("request candidates should read");
if candidates
.first()
.is_some_and(|candidate| candidate.status == RequestCandidateStatus::Cancelled)
{
break candidates;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("candidate should be marked cancelled");
assert_eq!(candidates[0].status_code, Some(499));
assert_eq!(
candidates[0].error_type.as_deref(),
Some("downstream_disconnect")
);
}
#[tokio::test]
async fn execute_execution_runtime_stream_records_first_data_as_streaming_before_terminal_telemetry(
) {

View File

@@ -6,6 +6,7 @@ pub(super) mod features;
mod model;
pub(super) mod observability;
pub(super) mod provider;
mod routing;
mod system;
mod users;
@@ -27,8 +28,7 @@ pub(crate) use self::provider::oauth::provisioning::{
};
pub(crate) use self::provider::oauth::quota::dispatch::refresh_provider_pool_quota_locally;
pub(crate) use self::provider::oauth::quota::shared::{
persist_provider_quota_refresh_state, provider_account_self_check_endpoint_for_provider,
provider_quota_refresh_endpoint_for_provider, provider_type_supports_account_self_check,
persist_provider_quota_refresh_state, provider_quota_refresh_endpoint_for_provider,
provider_type_supports_quota_refresh,
};
pub(crate) use self::provider::oauth::runtime::{

View File

@@ -8,7 +8,8 @@ use self::invalid::{
codex_structured_invalid_reason,
};
use self::parse::{
parse_codex_backend_me_response, parse_codex_usage_headers, parse_codex_wham_usage_response,
build_codex_quota_exhausted_fallback_metadata, parse_codex_usage_headers,
parse_codex_wham_usage_response,
};
use self::plan::{build_codex_quota_request_spec, execute_codex_quota_plan};
use super::shared::{
@@ -110,7 +111,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
"key_id": key.id,
"key_name": key.name,
"status": "error",
"message": format!("backend-api/me 请求执行失败: {detail}"),
"message": format!("wham/usage 请求执行失败: {detail}"),
"status_code": 502,
}));
continue;
@@ -137,9 +138,7 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
.as_ref()
.and_then(|body| body.json_body.as_ref())
{
if let Some(parsed) = parse_codex_backend_me_response(body_json, now_unix_secs)
.or_else(|| parse_codex_wham_usage_response(body_json, now_unix_secs))
{
if let Some(parsed) = parse_codex_wham_usage_response(body_json, now_unix_secs) {
metadata_update = Some(json!({
"codex": merge_codex_quota_metadata(header_metadata.as_ref(), &parsed)
}));
@@ -152,21 +151,21 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
status = "success".to_string();
} else {
status = "no_metadata".to_string();
message = Some("backend-api/me 响应中未包含账号信息".to_string());
message = Some("响应中未包含限额信息".to_string());
}
} else {
message = Some("无法解析 backend-api/me API 响应".to_string());
message = Some("无法解析 wham/usage API 响应".to_string());
}
} else {
let err_msg = extract_execution_error_message(&result);
message = Some(match err_msg.as_deref() {
Some(detail) if !detail.is_empty() => {
format!(
"backend-api/me API 返回状态码 {}: {}",
"wham/usage API 返回状态码 {}: {}",
result.status_code, detail
)
}
_ => format!("backend-api/me API 返回状态码 {}", result.status_code),
_ => format!("wham/usage API 返回状态码 {}", result.status_code),
});
match result.status_code {
@@ -223,14 +222,26 @@ pub(crate) async fn refresh_codex_provider_quota_locally(
oauth_invalid_reason = reason;
status = "workspace_deactivated".to_string();
} else {
let (at, reason) = codex_build_invalid_state(
&key,
codex_structured_invalid_reason(402, err_msg.as_deref()),
now_unix_secs,
);
oauth_invalid_at_unix_secs = at;
oauth_invalid_reason = reason;
status = "payment_required".to_string();
let plan_type = transport
.key
.decrypted_auth_config
.as_deref()
.and_then(|raw| serde_json::from_str::<serde_json::Value>(raw).ok())
.and_then(|value| {
value
.get("plan_type")
.and_then(serde_json::Value::as_str)
.map(ToOwned::to_owned)
});
metadata_update = Some(json!({
"codex": build_codex_quota_exhausted_fallback_metadata(
plan_type.as_deref(),
now_unix_secs,
)
}));
(oauth_invalid_at_unix_secs, oauth_invalid_reason) =
quota_refresh_success_invalid_state(&key);
status = "quota_exhausted".to_string();
}
}
403 => {

View File

@@ -22,13 +22,6 @@ pub(super) fn parse_codex_wham_usage_response(
admin_provider_quota_pure::parse_codex_wham_usage_response(value, updated_at_unix_secs)
}
pub(super) fn parse_codex_backend_me_response(
value: &serde_json::Value,
updated_at_unix_secs: u64,
) -> Option<serde_json::Value> {
admin_provider_quota_pure::parse_codex_backend_me_response(value, updated_at_unix_secs)
}
pub(super) fn parse_codex_usage_headers(
headers: &BTreeMap<String, String>,
updated_at_unix_secs: u64,

View File

@@ -84,22 +84,6 @@ pub(crate) fn provider_quota_refresh_endpoint_for_provider(
)
}
pub(crate) fn provider_type_supports_account_self_check(provider_type: &str) -> bool {
ProviderPoolService::with_builtin_adapters().supports_account_self_check(provider_type)
}
pub(crate) fn provider_account_self_check_endpoint_for_provider(
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
include_inactive: bool,
) -> Option<StoredProviderCatalogEndpoint> {
ProviderPoolService::with_builtin_adapters().account_self_check_endpoint_for_provider(
provider_type,
endpoints,
include_inactive,
)
}
pub(crate) fn provider_quota_refresh_missing_endpoint_message(provider_type: &str) -> String {
ProviderPoolService::with_builtin_adapters()
.quota_refresh_missing_endpoint_message(provider_type)

View File

@@ -226,7 +226,7 @@ pub(crate) struct AdminProviderUpdateRequest {
pub(crate) type AdminProviderUpdatePatch = AdminTypedObjectPatch<AdminProviderUpdateRequest>;
pub(crate) const CODEX_WHAM_USAGE_URL: &str = "https://chatgpt.com/backend-api/me";
pub(crate) const CODEX_WHAM_USAGE_URL: &str = "https://chatgpt.com/backend-api/wham/usage";
pub(crate) const KIRO_USAGE_LIMITS_PATH: &str = "/getUsageLimits";
pub(crate) const KIRO_USAGE_SDK_VERSION: &str = "1.0.0";
pub(crate) const ANTIGRAVITY_FETCH_AVAILABLE_MODELS_PATH: &str = "/v1internal:fetchAvailableModels";

View File

@@ -1,8 +1,7 @@
pub(crate) use self::{
create::build_admin_create_provider_key_record,
payload::{build_admin_provider_keys_page_payload, build_admin_provider_keys_payload},
update::build_admin_update_provider_key_record,
};
pub(crate) use self::create::build_admin_create_provider_key_record;
pub(crate) use self::payload::build_admin_provider_keys_page_payload;
pub(crate) use self::payload::build_admin_provider_keys_payload;
pub(crate) use self::update::build_admin_update_provider_key_record;
mod create;
mod payload;

View File

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

View File

@@ -9,6 +9,7 @@ mod observability;
mod provider;
mod provider_oauth;
mod route_request;
mod routing_profiles;
mod state;
mod system;
mod users;

View File

@@ -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
}
}

View File

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

View 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()
}

View File

@@ -17,7 +17,9 @@ use crate::handlers::admin::system::shared::settings::{
build_admin_system_stats_payload, current_aether_version, fetch_latest_admin_system_release,
};
use crate::handlers::admin::system::shared::smtp::build_admin_smtp_test_payload;
use crate::maintenance::{ManualUsageCleanupMode, ManualUsageCleanupOptions};
use crate::GatewayError;
use aether_data_contracts::repository::usage::UsageCleanupTargets;
use axum::{
body::{Body, Bytes},
http,
@@ -26,6 +28,7 @@ use axum::{
};
use serde_json::json;
use std::time::Instant;
use url::form_urlencoded;
pub(super) async fn maybe_build_local_admin_core_system_response(
state: &AdminAppState<'_>,
@@ -263,7 +266,7 @@ pub(super) async fn maybe_build_local_admin_core_system_response(
&& request_path == "/api/admin/system/cleanup/usage/manual"
{
return Ok(Some(
build_manual_usage_cleanup_response(state, request_body).await?,
build_manual_usage_cleanup_response(state, request_context, request_body).await?,
));
}
@@ -625,50 +628,36 @@ async fn build_admin_system_cleanup_payload(
async fn build_manual_usage_cleanup_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
request_body: Option<&Bytes>,
) -> Result<Response<Body>, GatewayError> {
let older_than_days = match parse_manual_usage_cleanup_request(request_body) {
let options = match parse_manual_usage_cleanup_request(request_body) {
Ok(value) => value,
Err(response) => return Ok(response),
};
let actor_user_id = request_context
.decision()
.and_then(|decision| decision.admin_principal.as_ref())
.map(|principal| principal.user_id.clone());
match crate::maintenance::run_manual_usage_cleanup_once(
&state.app().data,
older_than_days,
None,
match crate::maintenance::start_manual_usage_cleanup_task(
std::sync::Arc::clone(&state.app().data),
options,
actor_user_id,
)
.await
{
Ok(summary) => {
let total = summary
.body_externalized
.saturating_add(summary.legacy_body_refs_migrated)
.saturating_add(summary.body_cleaned)
.saturating_add(summary.header_cleaned)
.saturating_add(summary.keys_cleaned)
.saturating_add(summary.records_deleted);
let message = match older_than_days {
Some(days) => {
format!("请求记录手动清理完成,清理 {days} 天前的记录,影响 {total}")
}
None => format!("请求记录手动清理完成(按当前策略),影响 {total}"),
};
Ok(task) => {
let payload = json!({
"message": message,
"requested_older_than_days": older_than_days,
"summary": {
"body_externalized": summary.body_externalized,
"legacy_body_refs_migrated": summary.legacy_body_refs_migrated,
"body_cleaned": summary.body_cleaned,
"header_cleaned": summary.header_cleaned,
"keys_cleaned": summary.keys_cleaned,
"records_deleted": summary.records_deleted,
},
"total_affected": total,
"message": task.message,
"mode": options.mode.as_str(),
"requested_older_than_days": options.requested_older_than_days,
"targets": options.targets,
"task": task,
});
Ok(attach_admin_audit_response(
Json(payload).into_response(),
"admin_system_usage_cleanup_completed",
"admin_system_usage_cleanup_started",
"manual_usage_cleanup",
"usage_cleanup",
"global",
@@ -696,12 +685,20 @@ async fn build_manual_usage_cleanup_preview_response(
Ok(value) => value,
Err(response) => return Ok(response),
};
let preview =
crate::maintenance::preview_manual_usage_cleanup(&state.app().data, older_than_days)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let options = match parse_manual_usage_cleanup_query_options(
request_context.query_string(),
older_than_days,
) {
Ok(value) => value,
Err(response) => return Ok(response),
};
let preview = crate::maintenance::preview_manual_usage_cleanup(&state.app().data, options)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
Ok(Json(json!({
"mode": preview.mode.as_str(),
"requested_older_than_days": preview.requested_older_than_days,
"targets": preview.targets,
"effective_cutoffs": {
"detail": preview.detail_cutoff,
"compressed": preview.compressed_cutoff,
@@ -720,12 +717,12 @@ async fn build_manual_usage_cleanup_preview_response(
fn parse_manual_usage_cleanup_request(
request_body: Option<&Bytes>,
) -> Result<Option<u32>, Response<Body>> {
) -> Result<ManualUsageCleanupOptions, Response<Body>> {
let Some(body) = request_body else {
return Ok(None);
return Ok(ManualUsageCleanupOptions::policy());
};
if body.is_empty() {
return Ok(None);
return Ok(ManualUsageCleanupOptions::policy());
}
let parsed: serde_json::Value = match serde_json::from_slice(body) {
Ok(value) => value,
@@ -738,45 +735,208 @@ fn parse_manual_usage_cleanup_request(
}
};
let Some(object) = parsed.as_object() else {
return Ok(None);
return Err(bad_manual_cleanup_request("请求体必须为 JSON 对象"));
};
match object.get("older_than_days") {
None | Some(serde_json::Value::Null) => Ok(None),
Some(value) => value
.as_u64()
.and_then(|value| u32::try_from(value).ok())
.filter(|days| *days >= 1)
.map(Some)
.ok_or_else(|| {
(
http::StatusCode::BAD_REQUEST,
Json(json!({
"detail": "older_than_days 必须为正整数",
})),
)
.into_response()
}),
parse_manual_usage_cleanup_options(
object.get("mode").and_then(serde_json::Value::as_str),
object.get("older_than_days"),
object.get("targets"),
)
}
fn parse_manual_usage_cleanup_query_options(
query_string: Option<&str>,
older_than_days: Option<u32>,
) -> Result<ManualUsageCleanupOptions, Response<Body>> {
let mode = query_param(query_string, "mode");
let targets = query_param(query_string, "targets").map(serde_json::Value::String);
let older_value =
older_than_days.map(|days| serde_json::Value::Number(serde_json::Number::from(days)));
parse_manual_usage_cleanup_options(mode.as_deref(), older_value.as_ref(), targets.as_ref())
}
fn parse_manual_usage_cleanup_options(
raw_mode: Option<&str>,
older_than_days: Option<&serde_json::Value>,
targets_value: Option<&serde_json::Value>,
) -> Result<ManualUsageCleanupOptions, Response<Body>> {
let (requested_older_than_days, requested_before_now) =
parse_manual_cleanup_older_than_days(older_than_days)?;
let mode =
parse_manual_cleanup_mode(raw_mode, requested_older_than_days, requested_before_now)?;
if mode == ManualUsageCleanupMode::OlderThanDays && requested_older_than_days.is_none() {
return Err(bad_manual_cleanup_request(
"older_than_days 模式必须提供正整数天数",
));
}
if mode == ManualUsageCleanupMode::BeforeNow && requested_older_than_days.is_some() {
return Err(bad_manual_cleanup_request(
"before_now 模式不能同时提供 older_than_days",
));
}
if raw_mode.is_some() && requested_before_now && mode != ManualUsageCleanupMode::BeforeNow {
return Err(bad_manual_cleanup_request(
"older_than_days 为 0 时必须使用 before_now 模式",
));
}
if raw_mode.is_some()
&& mode == ManualUsageCleanupMode::Policy
&& requested_older_than_days.is_some()
{
return Err(bad_manual_cleanup_request(
"policy 模式不能同时提供 older_than_days",
));
}
let targets = parse_manual_cleanup_targets(targets_value, mode)?;
if !targets.any_selected() {
return Err(bad_manual_cleanup_request("至少选择一个清理范围"));
}
if mode == ManualUsageCleanupMode::BeforeNow
&& (targets.headers || targets.records || targets.expired_keys)
{
return Err(bad_manual_cleanup_request(
"清理当前时刻之前只允许选择详细请求体和压缩请求体",
));
}
Ok(ManualUsageCleanupOptions {
mode,
requested_older_than_days,
targets,
})
}
fn parse_manual_cleanup_mode(
raw_mode: Option<&str>,
requested_older_than_days: Option<u32>,
requested_before_now: bool,
) -> Result<ManualUsageCleanupMode, Response<Body>> {
let Some(raw) = raw_mode.map(str::trim).filter(|value| !value.is_empty()) else {
return Ok(if requested_before_now {
ManualUsageCleanupMode::BeforeNow
} else if requested_older_than_days.is_some() {
ManualUsageCleanupMode::OlderThanDays
} else {
ManualUsageCleanupMode::Policy
});
};
match raw {
"policy" => Ok(ManualUsageCleanupMode::Policy),
"older_than_days" => Ok(ManualUsageCleanupMode::OlderThanDays),
"before_now" => Ok(ManualUsageCleanupMode::BeforeNow),
_ => Err(bad_manual_cleanup_request(
"mode 必须为 policy、older_than_days 或 before_now",
)),
}
}
fn parse_older_than_days_query(query_string: Option<&str>) -> Result<Option<u32>, Response<Body>> {
let Some(query) = query_string.filter(|value| !value.is_empty()) else {
return Ok(None);
};
let value = query
.split('&')
.filter_map(|pair| pair.split_once('='))
.find_map(|(key, value)| {
if key == "older_than_days" && !value.is_empty() {
Some(value)
} else {
None
fn parse_manual_cleanup_older_than_days(
value: Option<&serde_json::Value>,
) -> Result<(Option<u32>, bool), Response<Body>> {
match value {
None | Some(serde_json::Value::Null) => Ok((None, false)),
Some(value) => {
let Some(raw) = value.as_u64() else {
return Err(bad_manual_cleanup_request("older_than_days 必须为非负整数"));
};
if raw == 0 {
return Ok((None, true));
}
let days = u32::try_from(raw)
.ok()
.filter(|days| *days >= 1)
.ok_or_else(|| bad_manual_cleanup_request("older_than_days 必须为正整数"))?;
Ok((Some(days), false))
}
}
}
fn parse_manual_cleanup_targets(
value: Option<&serde_json::Value>,
mode: ManualUsageCleanupMode,
) -> Result<UsageCleanupTargets, Response<Body>> {
let Some(value) = value else {
return Ok(match mode {
ManualUsageCleanupMode::BeforeNow => UsageCleanupTargets::body_targets(),
ManualUsageCleanupMode::Policy | ManualUsageCleanupMode::OlderThanDays => {
UsageCleanupTargets::all_policy_targets()
}
});
let Some(raw) = value else {
};
if value.is_null() {
return Ok(match mode {
ManualUsageCleanupMode::BeforeNow => UsageCleanupTargets::body_targets(),
ManualUsageCleanupMode::Policy | ManualUsageCleanupMode::OlderThanDays => {
UsageCleanupTargets::all_policy_targets()
}
});
}
let raw_targets = match value {
serde_json::Value::Array(items) => items
.iter()
.map(|item| {
item.as_str()
.map(str::to_string)
.ok_or_else(|| bad_manual_cleanup_request("targets 必须为字符串数组"))
})
.collect::<Result<Vec<_>, _>>()?,
serde_json::Value::String(raw) => raw
.split(',')
.map(str::trim)
.filter(|item| !item.is_empty())
.map(str::to_string)
.collect(),
_ => return Err(bad_manual_cleanup_request("targets 必须为字符串数组")),
};
let mut targets = UsageCleanupTargets {
detail_body: false,
compressed_body: false,
headers: false,
records: false,
expired_keys: false,
};
for raw in raw_targets {
match raw.as_str() {
"detail_body" | "detail" | "raw_body" => targets.detail_body = true,
"compressed_body" | "compressed" => targets.compressed_body = true,
"headers" | "header" => targets.headers = true,
"records" | "log" | "logs" => targets.records = true,
"expired_keys" => targets.expired_keys = true,
"all" => targets = UsageCleanupTargets::all_policy_targets(),
_ => {
return Err(bad_manual_cleanup_request(
"targets 只能包含 detail_body、compressed_body、headers、records",
))
}
}
}
Ok(targets)
}
fn bad_manual_cleanup_request(detail: impl Into<String>) -> Response<Body> {
(
http::StatusCode::BAD_REQUEST,
Json(json!({ "detail": detail.into() })),
)
.into_response()
}
fn query_param(query_string: Option<&str>, name: &str) -> Option<String> {
let query = query_string.filter(|value| !value.is_empty())?;
form_urlencoded::parse(query.as_bytes())
.find_map(|(key, value)| (key == name && !value.is_empty()).then(|| value.into_owned()))
}
fn parse_older_than_days_query(query_string: Option<&str>) -> Result<Option<u32>, Response<Body>> {
let Some(value) = query_param(query_string, "older_than_days") else {
return Ok(None);
};
raw.parse::<u32>()
value
.parse::<u32>()
.ok()
.filter(|days| *days >= 1)
.map(Some)
@@ -790,3 +950,55 @@ fn parse_older_than_days_query(query_string: Option<&str>) -> Result<Option<u32>
.into_response()
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn manual_cleanup_request_defaults_to_policy_targets() {
let options = parse_manual_usage_cleanup_request(None).expect("default request is valid");
assert_eq!(options.mode, ManualUsageCleanupMode::Policy);
assert_eq!(options.requested_older_than_days, None);
assert_eq!(options.targets, UsageCleanupTargets::all_policy_targets());
}
#[test]
fn manual_cleanup_request_treats_zero_days_as_before_now_body_only() {
let body = Bytes::from_static(br#"{"older_than_days":0}"#);
let options =
parse_manual_usage_cleanup_request(Some(&body)).expect("before-now request is valid");
assert_eq!(options.mode, ManualUsageCleanupMode::BeforeNow);
assert_eq!(options.requested_older_than_days, None);
assert_eq!(options.targets, UsageCleanupTargets::body_targets());
}
#[test]
fn manual_cleanup_request_rejects_before_now_headers() {
let body = Bytes::from_static(br#"{"mode":"before_now","targets":["headers"]}"#);
assert!(parse_manual_usage_cleanup_request(Some(&body)).is_err());
}
#[test]
fn manual_cleanup_preview_query_decodes_comma_separated_targets() {
let options = parse_manual_usage_cleanup_query_options(
Some("mode=before_now&targets=detail_body%2Ccompressed_body"),
None,
)
.expect("encoded targets query is valid");
assert_eq!(options.mode, ManualUsageCleanupMode::BeforeNow);
assert_eq!(options.targets, UsageCleanupTargets::body_targets());
}
#[test]
fn manual_cleanup_request_rejects_non_object_body() {
let body = Bytes::from_static(br#"[]"#);
assert!(parse_manual_usage_cleanup_request(Some(&body)).is_err());
}
}

View File

@@ -7,7 +7,7 @@ use crate::constants::{
};
use crate::control::GatewayControlDecision;
use crate::control::GatewayPublicRequestContext;
use crate::middleware::{should_downgrade_access_log, RequestLogEmitted};
use crate::middleware::{sanitize_access_log_path, should_downgrade_access_log, RequestLogEmitted};
use crate::AppState;
use aether_runtime::{maybe_hold_axum_response_permit, AdmissionPermit};
use axum::body::{Body, Bytes};
@@ -95,11 +95,12 @@ pub(super) fn finalize_gateway_response(
.map(|auth_context| auth_context.api_key_id.as_str())
.unwrap_or("-");
let status_code = response.status().as_u16();
let sanitized_path_and_query = sanitize_access_log_path(path_and_query);
emit_admin_audit(
&mut response,
trace_id,
method,
path_and_query,
sanitized_path_and_query.as_str(),
control_decision,
);
if response.status().is_server_error() {
@@ -112,7 +113,7 @@ pub(super) fn finalize_gateway_response(
request_id,
remote_addr = %remote_addr,
method = %method,
path = %path_and_query,
path = %sanitized_path_and_query,
user_id,
api_key_id,
route_class,
@@ -122,7 +123,7 @@ pub(super) fn finalize_gateway_response(
elapsed_ms,
"gateway request failed"
);
} else if should_downgrade_access_log(method, path_and_query) {
} else if should_downgrade_access_log(method, sanitized_path_and_query.as_str()) {
trace!(
event_name = "http_request_completed",
log_type = "access",
@@ -132,7 +133,7 @@ pub(super) fn finalize_gateway_response(
request_id,
remote_addr = %remote_addr,
method = %method,
path = %path_and_query,
path = %sanitized_path_and_query,
user_id,
api_key_id,
route_class,
@@ -152,7 +153,7 @@ pub(super) fn finalize_gateway_response(
request_id,
remote_addr = %remote_addr,
method = %method,
path = %path_and_query,
path = %sanitized_path_and_query,
user_id,
api_key_id,
route_class,
@@ -254,3 +255,106 @@ pub(super) fn finalize_gateway_response_with_context(
request_permit,
)
}
#[cfg(test)]
mod tests {
use super::finalize_gateway_response;
use crate::control::GatewayControlDecision;
use crate::AppState;
use axum::body::Body;
use axum::http::{Method, Response, StatusCode};
use std::sync::{Arc, Mutex};
use std::time::Instant;
use tracing_subscriber::filter::LevelFilter;
use tracing_subscriber::prelude::*;
#[derive(Clone, Default)]
struct SharedBuffer(Arc<Mutex<Vec<u8>>>);
struct SharedBufferWriter(Arc<Mutex<Vec<u8>>>);
impl SharedBuffer {
fn lines(&self) -> Vec<serde_json::Value> {
String::from_utf8(self.0.lock().expect("buffer should lock").clone())
.expect("buffer should contain valid utf-8")
.lines()
.filter(|line| !line.trim().is_empty())
.map(|line| serde_json::from_str(line).expect("json log line should parse"))
.collect()
}
}
impl std::io::Write for SharedBufferWriter {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.0
.lock()
.expect("buffer should lock")
.extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
impl<'a> tracing_subscriber::fmt::writer::MakeWriter<'a> for SharedBuffer {
type Writer = SharedBufferWriter;
fn make_writer(&'a self) -> Self::Writer {
SharedBufferWriter(Arc::clone(&self.0))
}
}
#[test]
fn finalize_gateway_response_logs_sanitized_path_and_query() {
let state = AppState::new().expect("gateway state should build");
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.json()
.flatten_event(true)
.with_current_span(false)
.with_span_list(false)
.with_writer(writer.clone())
.with_filter(LevelFilter::INFO),
);
let dispatch = tracing::Dispatch::new(subscriber);
let _guard = tracing::dispatcher::set_default(&dispatch);
let remote_addr = "127.0.0.1:8080"
.parse()
.expect("remote address should parse");
let control_decision = GatewayControlDecision::synthetic(
"/v1beta/models/gemini-3-flash-preview:generateContent",
Some("ai_public".to_string()),
Some("gemini".to_string()),
Some("generate_content".to_string()),
Some("gemini:generate_content".to_string()),
);
let response = Response::builder()
.status(StatusCode::OK)
.body(Body::empty())
.expect("response should build");
let _response = finalize_gateway_response(
&state,
response,
"trace-finalize",
&remote_addr,
&Method::GET,
"/v1beta/models/gemini-3-flash-preview:generateContent?key=secret&alt=sse",
Some(&control_decision),
"execution_runtime_sync",
&Instant::now(),
None,
);
let logs = writer.lines();
assert_eq!(logs.len(), 1);
assert_eq!(
logs[0]["path"],
"/v1beta/models/gemini-3-flash-preview:generateContent?alt=sse"
);
}
}

View File

@@ -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"))

View File

@@ -59,6 +59,7 @@ mod rate_limit;
mod request_candidate_runtime;
mod roles;
mod router;
mod routing;
mod scheduler;
mod state;
mod system_features;

View File

@@ -21,9 +21,10 @@ pub(crate) use runtime::{
spawn_stats_aggregation_worker, spawn_stats_hourly_aggregation_worker,
spawn_usage_cleanup_worker, spawn_wallet_daily_usage_aggregation_worker,
start_admin_request_body_cleanup_task, start_admin_system_purge_task,
start_proxy_upgrade_rollout, AccountSelfCheckRunSummary, AdminCleanupRunRecord,
AdminCleanupTaskKind, AdminStatsRebuildSummary, AdminSystemCleanupSummary,
ManualUsageCleanupError, OAuthTokenRefreshRunSummary, PoolQuotaProbeRunSummary,
start_manual_usage_cleanup_task, start_proxy_upgrade_rollout, AccountSelfCheckRunSummary,
AdminCleanupRunRecord, AdminCleanupTaskKind, AdminStatsRebuildSummary,
AdminSystemCleanupSummary, ManualUsageCleanupError, ManualUsageCleanupMode,
ManualUsageCleanupOptions, OAuthTokenRefreshRunSummary, PoolQuotaProbeRunSummary,
PoolQuotaProbeWorkerConfig, ProviderCheckinRunSummary, ProxyUpgradeRolloutCancelSummary,
ProxyUpgradeRolloutConflictClearSummary, ProxyUpgradeRolloutNodeActionSummary,
ProxyUpgradeRolloutProbeConfig, ProxyUpgradeRolloutSkippedRestoreSummary,

View File

@@ -61,9 +61,9 @@ pub(crate) use aether_data_contracts::repository::usage::{
};
use audit_cleanup::*;
pub(crate) use cleanup_runs::{
list_admin_cleanup_run_records, record_completed_cleanup_run, record_failed_cleanup_run,
start_admin_request_body_cleanup_task, start_admin_system_purge_task, AdminCleanupRunRecord,
AdminCleanupTaskKind, USAGE_CLEANUP_KIND,
list_admin_cleanup_run_records, record_admin_cleanup_run, record_completed_cleanup_run,
record_failed_cleanup_run, start_admin_request_body_cleanup_task,
start_admin_system_purge_task, AdminCleanupRunRecord, AdminCleanupTaskKind, USAGE_CLEANUP_KIND,
};
use config::*;
use db_maintenance::*;
@@ -98,12 +98,16 @@ pub(crate) use proxy_upgrade_rollout::{
};
use request_candidate_cleanup::*;
use runners::*;
pub(crate) use runners::{run_manual_usage_cleanup_once, ManualUsageCleanupError};
pub(crate) use runners::{
run_manual_usage_cleanup_once, start_manual_usage_cleanup_task, ManualUsageCleanupError,
};
use schedule::*;
use stats_daily::*;
use stats_hourly::*;
pub(crate) use usage_cleanup::preview_manual_usage_cleanup;
use usage_cleanup::*;
pub(crate) use usage_cleanup::{
preview_manual_usage_cleanup, ManualUsageCleanupMode, ManualUsageCleanupOptions,
};
use wallet_daily_usage::*;
pub(crate) use workers::*;

View File

@@ -14,8 +14,8 @@ use serde_json::{json, Value};
use tracing::{debug, info, warn};
use crate::admin_api::{
admin_provider_pool_config, provider_account_self_check_endpoint_for_provider,
provider_type_supports_account_self_check, refresh_provider_pool_quota_locally, AdminAppState,
admin_provider_pool_config, provider_quota_refresh_endpoint_for_provider,
provider_type_supports_quota_refresh, refresh_provider_pool_quota_locally, AdminAppState,
};
use crate::{AppState, GatewayError};
@@ -547,7 +547,7 @@ fn endpoint_for_self_check(
provider_type: &str,
endpoints: &[StoredProviderCatalogEndpoint],
) -> Option<StoredProviderCatalogEndpoint> {
provider_account_self_check_endpoint_for_provider(provider_type, endpoints, true)
provider_quota_refresh_endpoint_for_provider(provider_type, endpoints, true)
}
fn gateway_error_message(err: GatewayError) -> String {
@@ -637,7 +637,7 @@ pub(crate) async fn perform_account_self_check_once_with_config(
summary.providers_skipped = summary.providers_skipped.saturating_add(1);
continue;
};
if !provider_type_supports_account_self_check(&provider_type) {
if !provider_type_supports_quota_refresh(&provider_type) {
summary.providers_skipped = summary.providers_skipped.saturating_add(1);
continue;
}

View File

@@ -489,7 +489,7 @@ fn request_body_cleanup_record(
}
}
async fn record_cleanup_run(
pub(crate) async fn record_admin_cleanup_run(
data: &GatewayDataState,
record: AdminCleanupRunRecord,
) -> Result<(), DataLayerError> {
@@ -509,6 +509,13 @@ async fn record_cleanup_run(
Ok(())
}
async fn record_cleanup_run(
data: &GatewayDataState,
record: AdminCleanupRunRecord,
) -> Result<(), DataLayerError> {
record_admin_cleanup_run(data, record).await
}
fn parse_cleanup_run_records(value: Value) -> Vec<AdminCleanupRunRecord> {
value
.as_array()

View File

@@ -1,5 +1,6 @@
use aether_data_contracts::DataLayerError;
use serde_json::json;
use std::sync::Arc;
use std::time::Instant;
use tracing::{info, warn};
@@ -11,11 +12,12 @@ use super::{
cleanup_expired_gemini_file_mappings_once, cleanup_proxy_node_metrics_once,
cleanup_request_candidates_once, cleanup_stale_pending_requests_once,
cleanup_stale_proxy_nodes_once, collect_proxy_upgrade_rollout_probes, now_unix_secs,
perform_db_maintenance_once, perform_provider_checkin_once, perform_stats_aggregation_once,
perform_stats_hourly_aggregation_once, perform_usage_cleanup_once,
perform_usage_cleanup_once_with_override, perform_wallet_daily_usage_aggregation_once,
record_completed_cleanup_run, record_failed_cleanup_run, record_proxy_upgrade_traffic_success,
summarize_database_pool,
perform_db_maintenance_once, perform_manual_usage_cleanup_once, perform_provider_checkin_once,
perform_stats_aggregation_once, perform_stats_hourly_aggregation_once,
perform_usage_cleanup_once, perform_wallet_daily_usage_aggregation_once,
record_admin_cleanup_run, record_completed_cleanup_run, record_failed_cleanup_run,
record_proxy_upgrade_traffic_success, summarize_database_pool, AdminCleanupRunRecord,
ManualUsageCleanupOptions,
};
pub(super) async fn run_audit_cleanup_once(data: &GatewayDataState) -> Result<(), DataLayerError> {
@@ -299,15 +301,14 @@ pub(super) async fn run_usage_cleanup_once(data: &GatewayDataState) -> Result<()
Ok(())
}
pub(crate) async fn run_manual_usage_cleanup_once(
data: &GatewayDataState,
override_older_than_days: Option<u32>,
pub(crate) async fn start_manual_usage_cleanup_task(
data: Arc<GatewayDataState>,
options: ManualUsageCleanupOptions,
actor_user_id: Option<String>,
) -> Result<aether_data_contracts::repository::usage::UsageCleanupSummary, ManualUsageCleanupError>
{
) -> Result<AdminCleanupRunRecord, ManualUsageCleanupError> {
use super::{list_admin_cleanup_run_records, USAGE_CLEANUP_KIND};
let existing = list_admin_cleanup_run_records(data)
let existing = list_admin_cleanup_run_records(&data)
.await
.map_err(ManualUsageCleanupError::DataLayer)?;
if existing
@@ -318,10 +319,42 @@ pub(crate) async fn run_manual_usage_cleanup_once(
}
let started_at_unix_secs = now_unix_secs();
let record = AdminCleanupRunRecord {
id: uuid::Uuid::new_v4().to_string(),
kind: USAGE_CLEANUP_KIND.to_string(),
trigger: "manual".to_string(),
status: "processing".to_string(),
message: manual_usage_cleanup_start_message(options),
started_at_unix_secs,
completed_at_unix_secs: None,
duration_ms: None,
summary: manual_usage_cleanup_progress_summary(options, 0, None, actor_user_id.as_deref()),
error: None,
};
record_admin_cleanup_run(&data, record.clone())
.await
.map_err(ManualUsageCleanupError::DataLayer)?;
tokio::spawn(run_manual_usage_cleanup_task(
data,
record.clone(),
options,
actor_user_id,
));
Ok(record)
}
pub(crate) async fn run_manual_usage_cleanup_once(
data: &GatewayDataState,
options: ManualUsageCleanupOptions,
actor_user_id: Option<String>,
) -> Result<aether_data_contracts::repository::usage::UsageCleanupSummary, ManualUsageCleanupError>
{
use super::USAGE_CLEANUP_KIND;
let started_at = Instant::now();
let override_duration =
override_older_than_days.map(|days| chrono::Duration::days(i64::from(days)));
let summary = match perform_usage_cleanup_once_with_override(data, override_duration).await {
let started_at_unix_secs = now_unix_secs();
let summary = match perform_manual_usage_cleanup_once(data, options).await {
Ok(summary) => summary,
Err(err) => {
record_failed_cleanup_run(
@@ -336,17 +369,8 @@ pub(crate) async fn run_manual_usage_cleanup_once(
return Err(ManualUsageCleanupError::DataLayer(err));
}
};
let total = summary
.body_externalized
.saturating_add(summary.legacy_body_refs_migrated)
.saturating_add(summary.body_cleaned)
.saturating_add(summary.header_cleaned)
.saturating_add(summary.keys_cleaned)
.saturating_add(summary.records_deleted);
let message = match override_older_than_days {
Some(days) => format!("请求记录手动清理完成,清理 {days} 天前的记录,影响 {total}"),
None => format!("请求记录手动清理完成(按当前策略),影响 {total}"),
};
let total = usage_cleanup_total(summary);
let message = manual_usage_cleanup_completed_message(options, total);
record_completed_cleanup_run(
data,
USAGE_CLEANUP_KIND,
@@ -360,7 +384,9 @@ pub(crate) async fn run_manual_usage_cleanup_once(
"header_cleaned": summary.header_cleaned,
"keys_cleaned": summary.keys_cleaned,
"records_deleted": summary.records_deleted,
"requested_older_than_days": override_older_than_days,
"mode": options.mode.as_str(),
"requested_older_than_days": options.requested_older_than_days,
"targets": options.targets,
"actor_user_id": actor_user_id,
}),
message,
@@ -371,7 +397,8 @@ pub(crate) async fn run_manual_usage_cleanup_once(
log_type = "ops",
worker = "usage_cleanup",
trigger = "manual",
requested_older_than_days = override_older_than_days,
mode = options.mode.as_str(),
requested_older_than_days = options.requested_older_than_days,
actor_user_id = actor_user_id.as_deref(),
total_affected = total,
"gateway finished manual usage cleanup"
@@ -379,6 +406,154 @@ pub(crate) async fn run_manual_usage_cleanup_once(
Ok(summary)
}
async fn run_manual_usage_cleanup_task(
data: Arc<GatewayDataState>,
initial_record: AdminCleanupRunRecord,
options: ManualUsageCleanupOptions,
actor_user_id: Option<String>,
) {
let started_at = Instant::now();
match perform_manual_usage_cleanup_once(&data, options).await {
Ok(summary) => {
let total = usage_cleanup_total(summary);
let record = AdminCleanupRunRecord {
id: initial_record.id,
kind: initial_record.kind,
trigger: initial_record.trigger,
status: "completed".to_string(),
message: manual_usage_cleanup_completed_message(options, total),
started_at_unix_secs: initial_record.started_at_unix_secs,
completed_at_unix_secs: Some(now_unix_secs()),
duration_ms: Some(
started_at
.elapsed()
.as_millis()
.try_into()
.unwrap_or(u64::MAX),
),
summary: manual_usage_cleanup_progress_summary(
options,
100,
Some(summary),
actor_user_id.as_deref(),
),
error: None,
};
if let Err(err) = record_admin_cleanup_run(&data, record).await {
warn!(error = %err, "failed to record manual usage cleanup completion");
}
info!(
event_name = "usage_cleanup_manual_completed",
log_type = "ops",
worker = "usage_cleanup",
trigger = "manual",
mode = options.mode.as_str(),
requested_older_than_days = options.requested_older_than_days,
actor_user_id = actor_user_id.as_deref(),
total_affected = total,
"gateway finished manual usage cleanup task"
);
}
Err(err) => {
let record = AdminCleanupRunRecord {
id: initial_record.id,
kind: initial_record.kind,
trigger: initial_record.trigger,
status: "failed".to_string(),
message: "请求记录手动清理失败".to_string(),
started_at_unix_secs: initial_record.started_at_unix_secs,
completed_at_unix_secs: Some(now_unix_secs()),
duration_ms: Some(
started_at
.elapsed()
.as_millis()
.try_into()
.unwrap_or(u64::MAX),
),
summary: manual_usage_cleanup_progress_summary(
options,
100,
None,
actor_user_id.as_deref(),
),
error: Some(err.to_string()),
};
if let Err(record_err) = record_admin_cleanup_run(&data, record).await {
warn!(error = %record_err, "failed to record manual usage cleanup failure");
}
warn!(error = %err, "manual usage cleanup task failed");
}
}
}
fn usage_cleanup_total(
summary: aether_data_contracts::repository::usage::UsageCleanupSummary,
) -> usize {
summary
.body_externalized
.saturating_add(summary.legacy_body_refs_migrated)
.saturating_add(summary.body_cleaned)
.saturating_add(summary.header_cleaned)
.saturating_add(summary.keys_cleaned)
.saturating_add(summary.records_deleted)
}
fn manual_usage_cleanup_start_message(options: ManualUsageCleanupOptions) -> String {
match options.mode {
super::ManualUsageCleanupMode::BeforeNow => {
"请求记录手动清理已开始,清理当前时刻之前的已选请求体".to_string()
}
super::ManualUsageCleanupMode::OlderThanDays => format!(
"请求记录手动清理已开始,清理 {} 天前的已选内容",
options.requested_older_than_days.unwrap_or_default()
),
super::ManualUsageCleanupMode::Policy => {
"请求记录手动清理已开始,按当前策略清理已选内容".to_string()
}
}
}
fn manual_usage_cleanup_completed_message(
options: ManualUsageCleanupOptions,
total: usize,
) -> String {
match options.mode {
super::ManualUsageCleanupMode::BeforeNow => {
format!("请求记录手动清理完成,已清理当前时刻之前的已选请求体,影响 {total}")
}
super::ManualUsageCleanupMode::OlderThanDays => format!(
"请求记录手动清理完成,清理 {} 天前的已选内容,影响 {total}",
options.requested_older_than_days.unwrap_or_default()
),
super::ManualUsageCleanupMode::Policy => {
format!("请求记录手动清理完成(按当前策略),影响 {total}")
}
}
}
fn manual_usage_cleanup_progress_summary(
options: ManualUsageCleanupOptions,
progress_percent: u8,
summary: Option<aether_data_contracts::repository::usage::UsageCleanupSummary>,
actor_user_id: Option<&str>,
) -> serde_json::Value {
let summary = summary.unwrap_or_default();
json!({
"mode": options.mode.as_str(),
"requested_older_than_days": options.requested_older_than_days,
"targets": options.targets,
"progress_percent": progress_percent,
"body_externalized": summary.body_externalized,
"legacy_body_refs_migrated": summary.legacy_body_refs_migrated,
"body_cleaned": summary.body_cleaned,
"header_cleaned": summary.header_cleaned,
"keys_cleaned": summary.keys_cleaned,
"records_deleted": summary.records_deleted,
"total": usage_cleanup_total(summary),
"actor_user_id": actor_user_id,
})
}
#[derive(Debug)]
pub(crate) enum ManualUsageCleanupError {
AlreadyRunning,

View File

@@ -28,12 +28,12 @@ use super::{
spawn_stats_hourly_aggregation_worker, spawn_usage_cleanup_worker,
spawn_wallet_daily_usage_aggregation_worker, start_proxy_upgrade_rollout,
stats_aggregation_target_day, stats_hourly_aggregation_target_hour, summarize_database_pool,
usage_cleanup_settings, usage_cleanup_window, usage_cleanup_window_with_override,
wallet_daily_usage_aggregation_target, AppState, DbMaintenanceRunSummary,
FailedPendingUsageRow, GatewayDataState, ProxyNodeMetricsCleanupSettings,
ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow, UsageCleanupSettings, USAGE_CLEANUP_HOUR,
USAGE_CLEANUP_MINUTE, WALLET_DAILY_USAGE_AGGREGATION_HOUR,
WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
usage_cleanup_settings, usage_cleanup_window, usage_cleanup_window_for_mode,
usage_cleanup_window_with_override, wallet_daily_usage_aggregation_target, AppState,
DbMaintenanceRunSummary, FailedPendingUsageRow, GatewayDataState, ManualUsageCleanupMode,
ProxyNodeMetricsCleanupSettings, ProxyUpgradeRolloutProbeConfig, StalePendingUsageRow,
UsageCleanupSettings, USAGE_CLEANUP_HOUR, USAGE_CLEANUP_MINUTE,
WALLET_DAILY_USAGE_AGGREGATION_HOUR, WALLET_DAILY_USAGE_AGGREGATION_MINUTE,
};
#[tokio::test]
@@ -986,6 +986,29 @@ fn usage_cleanup_window_with_override_is_always_non_aggressive() {
assert_eq!(passthrough, policy);
}
#[test]
fn usage_cleanup_before_now_window_uses_current_timestamp_only() {
let now_utc = "2026-03-18T03:00:00Z"
.parse::<DateTime<Utc>>()
.expect("timestamp should parse");
let settings = UsageCleanupSettings {
detail_retention_days: 7,
compressed_retention_days: 30,
header_retention_days: 90,
log_retention_days: 365,
batch_size: 123,
auto_delete_expired_keys: false,
};
let window =
usage_cleanup_window_for_mode(now_utc, settings, ManualUsageCleanupMode::BeforeNow, None);
assert_eq!(window.detail_cutoff, now_utc);
assert_eq!(window.compressed_cutoff, now_utc);
assert_eq!(window.header_cutoff, now_utc);
assert_eq!(window.log_cutoff, now_utc);
}
#[tokio::test]
async fn summarize_database_pool_uses_busy_connections_for_usage_rate() {
let data = GatewayDataState::from_config(crate::data::GatewayDataConfig::from_postgres_config(

View File

@@ -1,4 +1,6 @@
use aether_data_contracts::repository::usage::{UsageCleanupSummary, UsageCleanupWindow};
use aether_data_contracts::repository::usage::{
UsageCleanupExecutionMode, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow,
};
use aether_data_contracts::DataLayerError;
use chrono::Utc;
@@ -15,6 +17,8 @@ pub(crate) struct ManualUsageCleanupPreview {
pub compressed_cutoff: chrono::DateTime<Utc>,
pub header_cutoff: chrono::DateTime<Utc>,
pub log_cutoff: chrono::DateTime<Utc>,
pub mode: ManualUsageCleanupMode,
pub targets: UsageCleanupTargets,
pub requested_older_than_days: Option<u32>,
pub detail_count: u64,
pub compressed_count: u64,
@@ -22,49 +26,131 @@ pub(crate) struct ManualUsageCleanupPreview {
pub log_count: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum ManualUsageCleanupMode {
Policy,
OlderThanDays,
BeforeNow,
}
impl ManualUsageCleanupMode {
pub(crate) const fn as_str(self) -> &'static str {
match self {
Self::Policy => "policy",
Self::OlderThanDays => "older_than_days",
Self::BeforeNow => "before_now",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct ManualUsageCleanupOptions {
pub(crate) mode: ManualUsageCleanupMode,
pub(crate) requested_older_than_days: Option<u32>,
pub(crate) targets: UsageCleanupTargets,
}
impl ManualUsageCleanupOptions {
pub(crate) const fn policy() -> Self {
Self {
mode: ManualUsageCleanupMode::Policy,
requested_older_than_days: None,
targets: UsageCleanupTargets::all_policy_targets(),
}
}
}
pub(super) async fn perform_usage_cleanup_once(
data: &GatewayDataState,
) -> Result<UsageCleanupSummary, DataLayerError> {
perform_usage_cleanup_once_with_override(data, None).await
perform_usage_cleanup_once_with_override(data, None, true).await
}
pub(super) async fn perform_usage_cleanup_once_with_override(
data: &GatewayDataState,
override_older_than: Option<chrono::Duration>,
respect_auto_enabled: bool,
) -> Result<UsageCleanupSummary, DataLayerError> {
let options = ManualUsageCleanupOptions {
mode: if override_older_than.is_some() {
ManualUsageCleanupMode::OlderThanDays
} else {
ManualUsageCleanupMode::Policy
},
requested_older_than_days: None,
targets: UsageCleanupTargets::all_policy_targets(),
};
perform_usage_cleanup_once_with_options(
data,
options,
override_older_than,
respect_auto_enabled,
)
.await
}
pub(super) async fn perform_manual_usage_cleanup_once(
data: &GatewayDataState,
options: ManualUsageCleanupOptions,
) -> Result<UsageCleanupSummary, DataLayerError> {
let override_duration = options
.requested_older_than_days
.map(|days| chrono::Duration::days(i64::from(days)));
perform_usage_cleanup_once_with_options(data, options, override_duration, false).await
}
async fn perform_usage_cleanup_once_with_options(
data: &GatewayDataState,
options: ManualUsageCleanupOptions,
override_older_than: Option<chrono::Duration>,
respect_auto_enabled: bool,
) -> Result<UsageCleanupSummary, DataLayerError> {
if !data.has_usage_writer() {
return Ok(UsageCleanupSummary::default());
}
if override_older_than.is_none()
if respect_auto_enabled
&& override_older_than.is_none()
&& !system_config_bool(data, "enable_auto_cleanup", true).await?
{
return Ok(UsageCleanupSummary::default());
}
let window = compute_usage_cleanup_window(data, override_older_than).await?;
let window = compute_usage_cleanup_window(data, options.mode, override_older_than).await?;
let settings = usage_cleanup_settings(data).await?;
data.cleanup_usage(
&window,
settings.batch_size,
settings.auto_delete_expired_keys,
options.targets,
cleanup_execution_mode(options.mode),
)
.await
}
pub(crate) async fn preview_manual_usage_cleanup(
data: &GatewayDataState,
override_older_than_days: Option<u32>,
options: ManualUsageCleanupOptions,
) -> Result<ManualUsageCleanupPreview, DataLayerError> {
let override_duration =
override_older_than_days.map(|days| chrono::Duration::days(i64::from(days)));
let window = compute_usage_cleanup_window(data, override_duration).await?;
let counts = data.preview_usage_cleanup(&window).await?;
let override_duration = options
.requested_older_than_days
.map(|days| chrono::Duration::days(i64::from(days)));
let window = compute_usage_cleanup_window(data, options.mode, override_duration).await?;
let counts = data
.preview_usage_cleanup(
&window,
options.targets,
cleanup_execution_mode(options.mode),
)
.await?;
Ok(ManualUsageCleanupPreview {
detail_cutoff: window.detail_cutoff,
compressed_cutoff: window.compressed_cutoff,
header_cutoff: window.header_cutoff,
log_cutoff: window.log_cutoff,
requested_older_than_days: override_older_than_days,
mode: options.mode,
targets: options.targets,
requested_older_than_days: options.requested_older_than_days,
detail_count: counts.detail,
compressed_count: counts.compressed,
header_count: counts.header,
@@ -74,11 +160,43 @@ pub(crate) async fn preview_manual_usage_cleanup(
async fn compute_usage_cleanup_window(
data: &GatewayDataState,
mode: ManualUsageCleanupMode,
override_older_than: Option<chrono::Duration>,
) -> Result<UsageCleanupWindow, DataLayerError> {
let settings = usage_cleanup_settings(data).await?;
Ok(match override_older_than {
Some(duration) => usage_cleanup_window_with_override(Utc::now(), settings, Some(duration)),
None => usage_cleanup_window(Utc::now(), settings),
})
Ok(usage_cleanup_window_for_mode(
Utc::now(),
settings,
mode,
override_older_than,
))
}
pub(super) fn usage_cleanup_window_for_mode(
now_utc: chrono::DateTime<Utc>,
settings: super::UsageCleanupSettings,
mode: ManualUsageCleanupMode,
override_older_than: Option<chrono::Duration>,
) -> UsageCleanupWindow {
match mode {
ManualUsageCleanupMode::Policy => usage_cleanup_window(now_utc, settings),
ManualUsageCleanupMode::OlderThanDays => {
usage_cleanup_window_with_override(now_utc, settings, override_older_than)
}
ManualUsageCleanupMode::BeforeNow => UsageCleanupWindow {
detail_cutoff: now_utc,
compressed_cutoff: now_utc,
header_cutoff: now_utc,
log_cutoff: now_utc,
},
}
}
fn cleanup_execution_mode(mode: ManualUsageCleanupMode) -> UsageCleanupExecutionMode {
match mode {
ManualUsageCleanupMode::BeforeNow => UsageCleanupExecutionMode::BeforeNowBodyFields,
ManualUsageCleanupMode::Policy | ManualUsageCleanupMode::OlderThanDays => {
UsageCleanupExecutionMode::Policy
}
}
}

View File

@@ -8,6 +8,7 @@ use axum::middleware::Next;
use axum::response::Response;
use tracing::{info, trace, warn};
use crate::ai_serving::api::sanitize_request_path_and_query;
use crate::constants::{
CONTROL_REQUEST_ID_HEADER, CONTROL_ROUTE_CLASS_HEADER, EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
};
@@ -52,14 +53,19 @@ pub(crate) fn should_downgrade_access_log(method: &Method, path: &str) -> bool {
|| normalized_path.starts_with("/api/admin/monitoring/trace/")
}
pub(crate) fn sanitize_access_log_path(path: &str) -> String {
sanitize_request_path_and_query(path, None).unwrap_or_else(|| "/".to_string())
}
pub(crate) async fn access_log_middleware(mut request: Request<Body>, next: Next) -> Response {
let started_at = Instant::now();
let method = request.method().clone();
let path = request
let raw_path = request
.uri()
.path_and_query()
.map(|value| value.as_str().to_string())
.unwrap_or_else(|| "/".to_string());
let path = sanitize_access_log_path(&raw_path);
let trace_id = extract_or_generate_trace_id(request.headers());
if !request.headers().contains_key(TRACE_ID_HEADER) {
request.headers_mut().insert(
@@ -158,7 +164,7 @@ pub(crate) async fn access_log_middleware(mut request: Request<Body>, next: Next
#[cfg(test)]
mod tests {
use super::{access_log_middleware, should_downgrade_access_log};
use super::{access_log_middleware, sanitize_access_log_path, should_downgrade_access_log};
use crate::constants::{
CONTROL_REQUEST_ID_HEADER, CONTROL_ROUTE_CLASS_HEADER, EXECUTION_PATH_HEADER,
TRACE_ID_HEADER,
@@ -212,6 +218,56 @@ mod tests {
}
}
#[test]
fn access_log_path_redacts_credential_query_values() {
assert_eq!(
sanitize_access_log_path(
"/v1beta/models/gemini-3-flash-preview:generateContent?key=secret&alt=sse&pageSize=10&token=hidden"
),
"/v1beta/models/gemini-3-flash-preview:generateContent?alt=sse&pageSize=10"
);
}
#[tokio::test(flavor = "current_thread")]
async fn access_log_emits_sanitized_path() {
let writer = SharedBuffer::default();
let subscriber = tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.json()
.flatten_event(true)
.with_current_span(false)
.with_span_list(false)
.with_writer(writer.clone())
.with_filter(LevelFilter::INFO),
);
let dispatch = tracing::Dispatch::new(subscriber);
let _guard = tracing::dispatcher::set_default(&dispatch);
let app = Router::new()
.route(
"/v1beta/models/gemini-3-flash-preview:generateContent",
get(|| async { Response::new(Body::empty()) }),
)
.layer(axum::middleware::from_fn(access_log_middleware));
let _response = app
.oneshot(
Request::builder()
.uri("/v1beta/models/gemini-3-flash-preview:generateContent?key=secret&alt=sse")
.body(Body::empty())
.expect("request should build"),
)
.await
.expect("request should succeed");
let logs = writer.lines();
assert_eq!(logs.len(), 1);
assert_eq!(
logs[0]["path"],
"/v1beta/models/gemini-3-flash-preview:generateContent?alt=sse"
);
}
#[tokio::test(flavor = "current_thread")]
async fn access_log_emits_completed_events_by_default() {
let writer = SharedBuffer::default();

View File

@@ -3,7 +3,7 @@ mod frontdoor_cors;
mod strip_cf_headers;
pub(crate) use access_log::{
access_log_middleware, should_downgrade_access_log, RequestLogEmitted,
access_log_middleware, sanitize_access_log_path, should_downgrade_access_log, RequestLogEmitted,
};
pub(crate) use frontdoor_cors::frontdoor_cors_middleware;
pub use strip_cf_headers::strip_cf_headers_middleware;

View 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;

View 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(())
}

View 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(),
})
}

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

View 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,
}
}

View File

@@ -11,6 +11,7 @@ mod cors;
mod integrations;
mod oauth;
mod proxy;
mod routing_profiles;
mod runtime;
#[cfg(test)]
mod testing;

View 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()))
}
}

View File

@@ -1,5 +1,31 @@
use super::*;
#[test]
fn admin_system_build_version_contract_uses_explicit_local_build_arg() {
let build_rs = read_workspace_file("apps/aether-gateway/build.rs");
for pattern in [
"cargo:rerun-if-env-changed=AETHER_BUILD_VERSION",
"env::var(\"AETHER_BUILD_VERSION\")",
] {
assert!(
build_rs.contains(pattern),
"apps/aether-gateway/build.rs should consume explicit build version pattern {pattern}"
);
}
let dockerfile = read_workspace_file("Dockerfile.app.local");
for pattern in [
"ARG AETHER_BUILD_VERSION",
"ENV AETHER_BUILD_VERSION=${AETHER_BUILD_VERSION}",
"AETHER_VERSION=${AETHER_BUILD_VERSION}",
] {
assert!(
dockerfile.contains(pattern),
"Dockerfile.app.local should pass explicit build version pattern {pattern}"
);
}
}
#[test]
fn admin_system_and_endpoint_roots_stay_thin() {
let system_mod = read_workspace_file("apps/aether-gateway/src/handlers/admin/system/mod.rs");

View File

@@ -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);

View File

@@ -29,7 +29,6 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
struct SeenExecutionRuntimeRequest {
url: String,
authorization: String,
accept: String,
provider_api_format: String,
total_ms: Option<u64>,
}
@@ -69,7 +68,6 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
.get("authorization")
.cloned()
.unwrap_or_default(),
accept: plan.headers.get("accept").cloned().unwrap_or_default(),
provider_api_format: plan.provider_api_format.clone(),
total_ms: plan
.timeouts
@@ -80,22 +78,41 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
request_id: plan.request_id,
candidate_id: None,
status_code: 200,
headers: BTreeMap::new(),
headers: BTreeMap::from([
(
"x-codex-primary-reset-after-seconds".to_string(),
"18000".to_string(),
),
(
"x-codex-primary-reset-at".to_string(),
"1900000000".to_string(),
),
(
"x-codex-secondary-reset-after-seconds".to_string(),
"604800".to_string(),
),
(
"x-codex-secondary-reset-at".to_string(),
"1900500000".to_string(),
),
]),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"user": {
"id": "user-codex-123",
"email": "codex@example.com",
"name": "Codex User"
"plan_type": "plus",
"rate_limit": {
"primary_window": {
"used_percent": 12.5,
"window_minutes": 300
},
"secondary_window": {
"used_percent": 55.0,
"window_minutes": 10080
}
},
"account": {
"id": "acct-codex-123",
"name": "Personal",
"plan_type": "plus"
},
"plan": {
"type": "Plus",
"title": "ChatGPT Plus"
"credits": {
"has_credits": true,
"balance": 42.0,
"unlimited": false
}
})),
body_bytes_b64: None,
@@ -167,22 +184,18 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
);
assert_eq!(payload["results"][0]["quota_snapshot"]["plan_type"], "plus");
assert_eq!(
payload["results"][0]["quota_snapshot"]["exhausted"],
json!(false)
payload["results"][0]["quota_snapshot"]["reset_at"],
1_900_000_000u64
);
assert_eq!(
payload["results"][0]["quota_snapshot"]["credits"]["balance"],
json!(42.0)
);
assert_eq!(
payload["results"][0]["quota_snapshot"]["windows"]
.as_array()
.map(Vec::len),
Some(0usize)
);
assert_eq!(
payload["results"][0]["metadata"]["email"],
"codex@example.com"
);
assert_eq!(
payload["results"][0]["metadata"]["account_id"],
"acct-codex-123"
Some(2usize)
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
@@ -193,13 +206,12 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
.expect("execution runtime request should be captured");
assert_eq!(
seen_execution_runtime_request.url,
"https://chatgpt.com/backend-api/me"
"https://chatgpt.com/backend-api/wham/usage"
);
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer sk-codex-123"
);
assert_eq!(seen_execution_runtime_request.accept, "application/json");
assert_eq!(
seen_execution_runtime_request.provider_api_format,
"openai:responses"
@@ -225,23 +237,33 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("email")),
Some(&json!("codex@example.com"))
.and_then(|value| value.get("primary_used_percent")),
Some(&json!(55.0))
);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("account_id")),
Some(&json!("acct-codex-123"))
.and_then(|value| value.get("primary_reset_at")),
Some(&json!(1_900_500_000u64))
);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("secondary_used_percent")),
Some(&json!(12.5))
);
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("secondary_reset_at")),
Some(&json!(1_900_000_000u64))
);
assert!(reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("primary_used_percent"))
.is_none());
gateway_handle.abort();
execution_runtime_handle.abort();
@@ -249,7 +271,7 @@ async fn gateway_refreshes_admin_provider_quota_locally_for_codex_with_trusted_a
}
#[tokio::test]
async fn gateway_marks_codex_key_invalid_when_backend_me_returns_payment_required() {
async fn gateway_marks_codex_quota_exhausted_when_wham_usage_returns_payment_required() {
let upstream = Router::new().route(
"/api/admin/endpoints/providers/provider-codex/refresh-quota",
any(move |_request: Request| async move {
@@ -337,18 +359,31 @@ async fn gateway_marks_codex_key_invalid_when_backend_me_returns_payment_require
let payload: serde_json::Value = response.json().await.expect("json body should parse");
assert_eq!(payload["success"], 0);
assert_eq!(payload["failed"], 1);
assert_eq!(payload["results"][0]["status"], "payment_required");
assert_eq!(payload["results"][0]["status"], "quota_exhausted");
assert_eq!(payload["results"][0]["status_code"], 402);
assert_eq!(
payload["results"][0]["quota_snapshot"]["provider_type"],
"codex"
);
assert_eq!(
payload["results"][0]["quota_snapshot"]["exhausted"],
json!(true)
);
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-a".to_string()])
.await
.expect("keys should read");
assert_eq!(reloaded.len(), 1);
assert!(reloaded[0].oauth_invalid_at_unix_secs.is_some());
assert_eq!(reloaded[0].oauth_invalid_at_unix_secs, None);
assert_eq!(reloaded[0].oauth_invalid_reason, None);
assert_eq!(
reloaded[0].oauth_invalid_reason.as_deref(),
Some("[ACCOUNT_BLOCK] payment required")
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(|value| value.get("codex"))
.and_then(|value| value.get("primary_used_percent")),
Some(&json!(100.0))
);
gateway_handle.abort();
@@ -1345,7 +1380,7 @@ async fn gateway_reports_codex_quota_runtime_failures_locally_without_falling_ba
assert!(payload["results"][0]["message"]
.as_str()
.expect("message should be string")
.contains("backend-api/me 请求执行失败: execution runtime returned HTTP 500"));
.contains("wham/usage 请求执行失败: execution runtime returned HTTP 500"));
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
let reloaded = provider_catalog_repository

View File

@@ -2454,9 +2454,8 @@ async fn gateway_completes_admin_provider_oauth_key_locally_with_trusted_admin_p
.as_str()
.expect("account_state_recheck_error should be string when recheck is attempted");
assert!(
account_state_recheck_error == "backend-api/me API 返回状态码 401"
|| account_state_recheck_error == "backend-api/me API 返回状态码 403"
|| account_state_recheck_error.starts_with("backend-api/me 请求执行失败:"),
account_state_recheck_error == "wham/usage API 返回状态码 401"
|| account_state_recheck_error.starts_with("wham/usage 请求执行失败:"),
"unexpected account_state_recheck_error: {account_state_recheck_error}"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
@@ -4926,9 +4925,8 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
.as_str()
.expect("account_state_recheck_error should be string when attempted");
assert!(
account_state_recheck_error == "backend-api/me API 返回状态码 401"
|| account_state_recheck_error == "backend-api/me API 返回状态码 403"
|| account_state_recheck_error.starts_with("backend-api/me 请求执行失败:"),
account_state_recheck_error == "wham/usage API 返回状态码 401"
|| account_state_recheck_error.starts_with("wham/usage 请求执行失败:"),
"unexpected account_state_recheck_error: {account_state_recheck_error}"
);
} else {
@@ -4946,7 +4944,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
.expect("execution runtime request should be captured");
assert_eq!(
seen_execution_runtime_request.url,
"https://chatgpt.com/backend-api/me"
"https://chatgpt.com/backend-api/wham/usage"
);
assert_eq!(
seen_execution_runtime_request.authorization,
@@ -4970,7 +4968,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
.expect("refreshed api key should decrypt");
assert_eq!(decrypted_api_key, "refreshed-codex-access-token");
if account_state_recheck_attempted
&& payload["account_state_recheck_error"] == "backend-api/me API 返回状态码 401"
&& payload["account_state_recheck_error"] == "wham/usage API 返回状态码 401"
{
assert!(stored_key.oauth_invalid_at_unix_secs.is_some());
assert_eq!(

View File

@@ -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

View File

@@ -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;

Some files were not shown because too many files have changed in this diff Show More