i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Unify candidate ranking pipeline

Browse Source
This commit is contained in:
fawney19
2026-04-27 12:34:03 +08:00
parent 9b866a6d17
commit 3b542434a2
46 changed files with 3635 additions and 1914 deletions
Download Patch File Download Diff File Expand all files Collapse all files

974
apps/aether-gateway/src/ai_pipeline/planner/candidate_affinity.rs
View File

File diff suppressed because it is too large Load Diff

394
apps/aether-gateway/src/ai_pipeline/planner/candidate_eligibility.rs
View File

@@ -1,393 +1 @@
use std::sync::Arc;
use aether_provider_transport::provider_types::provider_type_is_fixed;
use tracing::warn;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
use crate::orchestration::LocalExecutionCandidateMetadata;
use super::candidate_affinity::rank_eligible_local_execution_candidates;
use super::pool_scheduler::apply_local_execution_pool_scheduler;
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct EligibleLocalExecutionCandidate {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) transport: Arc<GatewayProviderTransportSnapshot>,
pub(crate) provider_api_format: String,
pub(crate) orchestration: LocalExecutionCandidateMetadata,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct SkippedLocalExecutionCandidate {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) skip_reason: &'static str,
pub(crate) transport: Option<Arc<GatewayProviderTransportSnapshot>>,
pub(crate) extra_data: Option<serde_json::Value>,
}
impl SkippedLocalExecutionCandidate {
pub(crate) fn transport_ref(&self) -> Option<&GatewayProviderTransportSnapshot> {
self.transport.as_deref()
}
}
pub(crate) async fn filter_and_rank_local_execution_candidates(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
requested_model: &str,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let requested_model = requested_model.trim();
filter_and_rank_local_execution_candidates_with_gate(
state,
candidates,
client_api_format,
required_capabilities,
sticky_session_token,
|candidate, transport, normalized_client_api_format| {
current_local_execution_candidate_skip_reason_with_transport(
candidate,
transport,
normalized_client_api_format,
requested_model,
)
},
)
.await
}
pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport_pair_gate(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
requested_model: Option<&str>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let requested_model = requested_model.map(str::trim);
filter_and_rank_local_execution_candidates_with_gate(
state,
candidates,
client_api_format,
required_capabilities,
sticky_session_token,
|candidate, transport, _normalized_client_api_format| {
current_local_execution_candidate_common_skip_reason_with_transport(
candidate,
transport,
requested_model,
)
},
)
.await
}
async fn filter_and_rank_local_execution_candidates_with_gate<F>(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
runtime_skip_reason: F,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
)
where
F: Fn(
&SchedulerMinimalCandidateSelectionCandidate,
&GatewayProviderTransportSnapshot,
&str,
) -> Option<&'static str>,
{
let normalized_client_api_format = client_api_format.trim().to_ascii_lowercase();
let mut selectable = Vec::with_capacity(candidates.len());
let mut skipped = Vec::with_capacity(candidates.len());
for candidate in candidates {
let Some(transport) = read_candidate_transport_snapshot(state, &candidate).await else {
skipped.push(SkippedLocalExecutionCandidate {
candidate,
skip_reason: "transport_snapshot_missing",
transport: None,
extra_data: None,
});
continue;
};
let transport = Arc::new(transport);
if candidate_is_ineligible_due_to_disabled_format_conversion(
transport.as_ref(),
normalized_client_api_format.as_str(),
) {
continue;
}
match runtime_skip_reason(
&candidate,
transport.as_ref(),
normalized_client_api_format.as_str(),
) {
Some(skip_reason) => skipped.push(SkippedLocalExecutionCandidate {
candidate,
skip_reason,
transport: Some(transport),
extra_data: None,
}),
None => selectable.push(EligibleLocalExecutionCandidate {
provider_api_format: transport.endpoint.api_format.trim().to_ascii_lowercase(),
candidate,
transport,
orchestration: LocalExecutionCandidateMetadata::default(),
}),
}
}
let ranked = rank_eligible_local_execution_candidates(
state,
selectable,
normalized_client_api_format.as_str(),
required_capabilities,
)
.await;
let (ranked, pool_skipped) =
apply_local_execution_pool_scheduler(state, ranked, sticky_session_token).await;
skipped.extend(pool_skipped);
(ranked, skipped)
}
pub(crate) fn extract_pool_sticky_session_token(body_json: &serde_json::Value) -> Option<String> {
fn non_empty_str(value: Option<&serde_json::Value>) -> Option<&str> {
value
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
let object = body_json.as_object()?;
non_empty_str(object.get("prompt_cache_key"))
.or_else(|| non_empty_str(object.get("conversation_id")))
.or_else(|| non_empty_str(object.get("conversationId")))
.or_else(|| non_empty_str(object.get("session_id")))
.or_else(|| non_empty_str(object.get("sessionId")))
.or_else(|| {
object
.get("metadata")
.and_then(serde_json::Value::as_object)
.and_then(|metadata| {
non_empty_str(metadata.get("session_id"))
.or_else(|| non_empty_str(metadata.get("conversation_id")))
})
})
.or_else(|| {
object
.get("conversationState")
.and_then(serde_json::Value::as_object)
.and_then(|state| {
non_empty_str(state.get("conversationId"))
.or_else(|| non_empty_str(state.get("sessionId")))
})
})
.map(ToOwned::to_owned)
}
fn current_local_execution_candidate_common_skip_reason_with_transport(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
transport: &GatewayProviderTransportSnapshot,
requested_model: Option<&str>,
) -> Option<&'static str> {
let requested_model = requested_model.unwrap_or_default();
if !transport.provider.is_active {
return Some("provider_inactive");
}
if !transport.endpoint.is_active {
return Some("endpoint_inactive");
}
if !transport.key.is_active {
return Some("key_inactive");
}
let endpoint_api_format = transport.endpoint.api_format.trim();
if !candidate
.endpoint_api_format
.trim()
.eq_ignore_ascii_case(endpoint_api_format)
&& !api_format_matches(&candidate.endpoint_api_format, endpoint_api_format)
{
return Some("endpoint_api_format_changed");
}
if !transport_key_supports_api_format(transport, endpoint_api_format) {
return Some("key_api_format_disabled");
}
if !transport_key_allows_candidate_model(transport, requested_model, candidate) {
return Some("key_model_disabled");
}
None
}
fn candidate_is_ineligible_due_to_disabled_format_conversion(
transport: &GatewayProviderTransportSnapshot,
normalized_client_api_format: &str,
) -> bool {
let endpoint_api_format = transport.endpoint.api_format.trim();
if api_format_matches(endpoint_api_format, normalized_client_api_format) {
return false;
}
crate::ai_pipeline::conversion::request_conversion_kind(
normalized_client_api_format,
endpoint_api_format,
)
.is_some()
&& crate::ai_pipeline::conversion::request_conversion_requires_enable_flag(
normalized_client_api_format,
endpoint_api_format,
)
&& !crate::ai_pipeline::conversion::request_conversion_enabled_for_transport(
transport,
normalized_client_api_format,
endpoint_api_format,
)
}
fn current_local_execution_candidate_skip_reason_with_transport(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
transport: &GatewayProviderTransportSnapshot,
normalized_client_api_format: &str,
requested_model: &str,
) -> Option<&'static str> {
if let Some(skip_reason) = current_local_execution_candidate_common_skip_reason_with_transport(
candidate,
transport,
Some(requested_model),
) {
return Some(skip_reason);
}
let endpoint_api_format = transport.endpoint.api_format.trim();
if api_format_matches(endpoint_api_format, normalized_client_api_format) {
return None;
}
if !crate::ai_pipeline::conversion::request_pair_allowed_for_transport(
transport,
normalized_client_api_format,
endpoint_api_format,
) {
return Some("transport_unsupported");
}
None
}
fn transport_key_supports_api_format(
transport: &GatewayProviderTransportSnapshot,
endpoint_api_format: &str,
) -> bool {
let provider_type = transport.provider.provider_type.trim();
let auth_type = transport.key.auth_type.trim();
let inherits_provider_api_formats = provider_type_is_fixed(provider_type)
&& (auth_type.eq_ignore_ascii_case("oauth")
|| (provider_type.eq_ignore_ascii_case("kiro")
&& auth_type.eq_ignore_ascii_case("bearer")
&& transport
.key
.decrypted_auth_config
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty())));
if inherits_provider_api_formats {
return true;
}
match transport.key.api_formats.as_deref() {
None => true,
Some(formats) => formats
.iter()
.any(|value| api_format_matches(value, endpoint_api_format)),
}
}
fn normalize_api_format_alias(value: &str) -> String {
crate::ai_pipeline::normalize_legacy_openai_format_alias(value)
}
fn api_format_matches(left: &str, right: &str) -> bool {
normalize_api_format_alias(left) == normalize_api_format_alias(right)
}
fn transport_key_allows_candidate_model(
transport: &GatewayProviderTransportSnapshot,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
let Some(allowed_models) = transport.key.allowed_models.as_deref() else {
return true;
};
let requested_model = requested_model.trim();
let global_model_name = candidate.global_model_name.trim();
let selected_provider_model_name = candidate.selected_provider_model_name.trim();
let mapping_matched_model = candidate
.mapping_matched_model
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
for allowed_model in allowed_models.iter().map(String::as_str).map(str::trim) {
if allowed_model.is_empty() {
continue;
}
if allowed_model == requested_model
|| allowed_model == global_model_name
|| allowed_model == selected_provider_model_name
|| mapping_matched_model.is_some_and(|value| value == allowed_model)
{
return true;
}
}
false
}
pub(crate) async fn read_candidate_transport_snapshot(
state: PlannerAppState<'_>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> Option<GatewayProviderTransportSnapshot> {
match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(transport)) => Some(transport),
Ok(None) => None,
Err(error) => {
warn!(
event_name = "candidate_eligibility_transport_load_failed",
log_type = "event",
provider_id = %candidate.provider_id,
endpoint_id = %candidate.endpoint_id,
key_id = %candidate.key_id,
error = ?error,
"failed to load provider transport while evaluating local candidate eligibility"
);
None
}
}
}
pub(crate) use super::candidate_resolution::*;

135
apps/aether-gateway/src/ai_pipeline/planner/candidate_materialization.rs
View File

@@ -1,9 +1,10 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use aether_scheduler_core::SchedulerRankingOutcome;
use serde_json::Value;
use uuid::Uuid;
use crate::ai_pipeline::planner::candidate_affinity::remember_scheduler_affinity_for_candidate;
use crate::ai_pipeline::planner::candidate_eligibility::{
use crate::ai_pipeline::planner::candidate_resolution::{
EligibleLocalExecutionCandidate, SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::failure_diagnostic::CandidateFailureDiagnostic;
@@ -70,7 +71,10 @@ pub(crate) fn remember_first_local_candidate_affinity(
}
fn should_persist_available_local_candidate(eligible: &EligibleLocalExecutionCandidate) -> bool {
eligible.orchestration.pool_key_index.is_none()
eligible
.orchestration
.pool_key_index
.is_none_or(|index| index == 0)
}
fn should_persist_skipped_local_candidate(candidate: &SkippedLocalExecutionCandidate) -> bool {
@@ -104,7 +108,10 @@ where
let candidate_index = candidate_index as u32;
let attempt_slots = local_attempt_slot_count(&eligible.transport);
let pool_key_index = eligible.orchestration.pool_key_index;
let extra_data = build_extra_data(&eligible);
let extra_data = local_candidate_extra_data_with_ranking(
build_extra_data(&eligible),
eligible.ranking.as_ref(),
);
let mut owned_eligible = Some(eligible);
for retry_index in 0..attempt_slots {
@@ -153,6 +160,54 @@ where
materialized
}
fn local_candidate_extra_data_with_ranking(
extra_data: Option<Value>,
ranking: Option<&SchedulerRankingOutcome>,
) -> Option<Value> {
let Some(ranking) = ranking else {
return extra_data;
};
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(
"ranking_mode".to_string(),
Value::String(format!("{:?}", ranking.ranking_mode)),
);
object.insert(
"priority_mode".to_string(),
Value::String(format!("{:?}", ranking.priority_mode)),
);
object.insert(
"ranking_index".to_string(),
Value::Number(serde_json::Number::from(ranking.ranking_index as u64)),
);
object.insert(
"priority_slot".to_string(),
Value::Number(serde_json::Number::from(i64::from(ranking.priority_slot))),
);
if let Some(promoted_by) = ranking.promoted_by {
object.insert(
"promoted_by".to_string(),
Value::String(promoted_by.to_string()),
);
}
if let Some(demoted_by) = ranking.demoted_by {
object.insert(
"demoted_by".to_string(),
Value::String(demoted_by.to_string()),
);
}
Some(Value::Object(object))
}
pub(crate) async fn persist_available_local_execution_candidates_with_context<F>(
state: PlannerAppState<'_>,
trace_id: &str,
@@ -356,7 +411,10 @@ mod tests {
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider,
};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use aether_scheduler_core::{
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode, SchedulerRankingMode,
SchedulerRankingOutcome,
};
use serde_json::json;
use super::*;
@@ -456,11 +514,12 @@ mod tests {
candidate_group_id: pool_key_index.map(|_| "pool-group".to_string()),
pool_key_index,
},
ranking: None,
}
}
#[tokio::test]
async fn pool_candidates_are_not_persisted_as_available_before_attempt() {
async fn pool_group_representatives_are_persisted_as_available_before_attempt() {
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
let app = AppState::new()
.expect("state should build")
@@ -478,6 +537,7 @@ mod tests {
None,
vec![
sample_eligible("pool-key", Some(0)),
sample_eligible("pool-key-internal", Some(1)),
sample_eligible("normal-key", None),
],
"persist should not fail",
@@ -485,13 +545,74 @@ mod tests {
)
.await;
assert_eq!(attempts.len(), 2);
assert_eq!(attempts.len(), 3);
let stored = app
.read_request_candidates_by_request_id("trace-pool-lazy")
.await
.expect("request candidates should read");
assert_eq!(stored.len(), 2);
assert_eq!(stored[0].key_id.as_deref(), Some("pool-key"));
assert_eq!(stored[0].candidate_index, 0);
assert_eq!(stored[1].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[1].candidate_index, 2);
}
#[tokio::test]
async fn available_candidates_persist_ranking_metadata_in_extra_data() {
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
let app = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
GatewayDataState::with_request_candidate_repository_for_tests(Arc::clone(
&repository,
)),
);
let mut eligible = sample_eligible("ranked-key", None);
eligible.ranking = Some(SchedulerRankingOutcome {
original_index: 1,
ranking_index: 0,
priority_mode: SchedulerPriorityMode::Provider,
ranking_mode: SchedulerRankingMode::CacheAffinity,
priority_slot: 7,
promoted_by: Some("cached_affinity"),
demoted_by: Some("cross_format"),
});
persist_available_local_execution_candidates(
PlannerAppState::new(&app),
"trace-ranking-extra-data",
"user-1",
"api-key-1",
None,
vec![eligible],
"persist should not fail",
|_| Some(json!({ "existing": "value" })),
)
.await;
let stored = app
.read_request_candidates_by_request_id("trace-ranking-extra-data")
.await
.expect("request candidates should read");
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].key_id.as_deref(), Some("normal-key"));
let extra_data = stored[0]
.extra_data
.as_ref()
.and_then(serde_json::Value::as_object)
.expect("ranking metadata should persist as object extra data");
assert_eq!(extra_data.get("existing"), Some(&json!("value")));
assert_eq!(
extra_data.get("ranking_mode"),
Some(&json!("CacheAffinity"))
);
assert_eq!(extra_data.get("priority_mode"), Some(&json!("Provider")));
assert_eq!(extra_data.get("ranking_index"), Some(&json!(0)));
assert_eq!(extra_data.get("priority_slot"), Some(&json!(7)));
assert_eq!(
extra_data.get("promoted_by"),
Some(&json!("cached_affinity"))
);
assert_eq!(extra_data.get("demoted_by"), Some(&json!("cross_format")));
}
#[tokio::test]

2
apps/aether-gateway/src/ai_pipeline/planner/candidate_metadata.rs
View File

@@ -2,7 +2,7 @@ use aether_contracts::ProxySnapshot;
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use serde_json::{json, Map, Value};
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::passthrough::resolve_same_format_provider_transport_unsupported_reason_for_trace;
use crate::ai_pipeline::transport::{
body_rules_are_locally_supported, header_rules_are_locally_supported,

414
apps/aether-gateway/src/ai_pipeline/planner/candidate_resolution.rs Normal file
View File

@@ -0,0 +1,414 @@
use std::sync::Arc;
use aether_provider_transport::provider_types::provider_type_is_fixed;
use tracing::warn;
use aether_scheduler_core::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRankingOutcome};
use crate::ai_pipeline::{
GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot, PlannerAppState,
};
use crate::orchestration::LocalExecutionCandidateMetadata;
use super::candidate_affinity::rank_eligible_local_execution_candidates;
use super::pool_scheduler::apply_local_execution_pool_scheduler;
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct EligibleLocalExecutionCandidate {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) transport: Arc<GatewayProviderTransportSnapshot>,
pub(crate) provider_api_format: String,
pub(crate) orchestration: LocalExecutionCandidateMetadata,
pub(crate) ranking: Option<SchedulerRankingOutcome>,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct SkippedLocalExecutionCandidate {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) skip_reason: &'static str,
pub(crate) transport: Option<Arc<GatewayProviderTransportSnapshot>>,
pub(crate) extra_data: Option<serde_json::Value>,
}
impl SkippedLocalExecutionCandidate {
pub(crate) fn transport_ref(&self) -> Option<&GatewayProviderTransportSnapshot> {
self.transport.as_deref()
}
}
pub(crate) async fn filter_and_rank_local_execution_candidates(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
requested_model: &str,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let requested_model = requested_model.trim();
filter_and_rank_local_execution_candidates_with_gate(
state,
candidates,
client_api_format,
Some(requested_model),
auth_snapshot,
required_capabilities,
sticky_session_token,
|candidate, transport, normalized_client_api_format| {
current_local_execution_candidate_skip_reason_with_transport(
candidate,
transport,
normalized_client_api_format,
requested_model,
)
},
)
.await
}
pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport_pair_gate(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
requested_model: Option<&str>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let requested_model = requested_model.map(str::trim);
filter_and_rank_local_execution_candidates_with_gate(
state,
candidates,
client_api_format,
requested_model,
auth_snapshot,
required_capabilities,
sticky_session_token,
|candidate, transport, _normalized_client_api_format| {
current_local_execution_candidate_common_skip_reason_with_transport(
candidate,
transport,
requested_model,
)
},
)
.await
}
async fn filter_and_rank_local_execution_candidates_with_gate<F>(
state: PlannerAppState<'_>,
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
client_api_format: &str,
requested_model: Option<&str>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&serde_json::Value>,
sticky_session_token: Option<&str>,
runtime_skip_reason: F,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
)
where
F: Fn(
&SchedulerMinimalCandidateSelectionCandidate,
&GatewayProviderTransportSnapshot,
&str,
) -> Option<&'static str>,
{
let normalized_client_api_format = client_api_format.trim().to_ascii_lowercase();
let mut selectable = Vec::with_capacity(candidates.len());
let mut skipped = Vec::with_capacity(candidates.len());
for candidate in candidates {
let Some(transport) = read_candidate_transport_snapshot(state, &candidate).await else {
skipped.push(SkippedLocalExecutionCandidate {
candidate,
skip_reason: "transport_snapshot_missing",
transport: None,
extra_data: None,
});
continue;
};
let transport = Arc::new(transport);
match runtime_skip_reason(
&candidate,
transport.as_ref(),
normalized_client_api_format.as_str(),
) {
Some(skip_reason) => skipped.push(SkippedLocalExecutionCandidate {
candidate,
skip_reason,
transport: Some(transport),
extra_data: None,
}),
None => selectable.push(EligibleLocalExecutionCandidate {
provider_api_format: transport.endpoint.api_format.trim().to_ascii_lowercase(),
candidate,
transport,
orchestration: LocalExecutionCandidateMetadata::default(),
ranking: None,
}),
}
}
let ranked = rank_eligible_local_execution_candidates(
state,
selectable,
normalized_client_api_format.as_str(),
requested_model,
auth_snapshot,
required_capabilities,
)
.await;
let (ranked, pool_skipped) =
apply_local_execution_pool_scheduler(state, ranked, sticky_session_token).await;
skipped.extend(pool_skipped);
(ranked, skipped)
}
pub(crate) fn extract_pool_sticky_session_token(body_json: &serde_json::Value) -> Option<String> {
fn non_empty_str(value: Option<&serde_json::Value>) -> Option<&str> {
value
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
}
let object = body_json.as_object()?;
non_empty_str(object.get("prompt_cache_key"))
.or_else(|| non_empty_str(object.get("conversation_id")))
.or_else(|| non_empty_str(object.get("conversationId")))
.or_else(|| non_empty_str(object.get("session_id")))
.or_else(|| non_empty_str(object.get("sessionId")))
.or_else(|| {
object
.get("metadata")
.and_then(serde_json::Value::as_object)
.and_then(|metadata| {
non_empty_str(metadata.get("session_id"))
.or_else(|| non_empty_str(metadata.get("conversation_id")))
})
})
.or_else(|| {
object
.get("conversationState")
.and_then(serde_json::Value::as_object)
.and_then(|state| {
non_empty_str(state.get("conversationId"))
.or_else(|| non_empty_str(state.get("sessionId")))
})
})
.map(ToOwned::to_owned)
}
fn current_local_execution_candidate_common_skip_reason_with_transport(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
transport: &GatewayProviderTransportSnapshot,
requested_model: Option<&str>,
) -> Option<&'static str> {
let requested_model = requested_model.unwrap_or_default();
if !transport.provider.is_active {
return Some("provider_inactive");
}
if !transport.endpoint.is_active {
return Some("endpoint_inactive");
}
if !transport.key.is_active {
return Some("key_inactive");
}
let endpoint_api_format = transport.endpoint.api_format.trim();
if !candidate
.endpoint_api_format
.trim()
.eq_ignore_ascii_case(endpoint_api_format)
&& !api_format_matches(&candidate.endpoint_api_format, endpoint_api_format)
{
return Some("endpoint_api_format_changed");
}
if !transport_key_supports_api_format(transport, endpoint_api_format) {
return Some("key_api_format_disabled");
}
if !transport_key_allows_candidate_model(transport, requested_model, candidate) {
return Some("key_model_disabled");
}
None
}
fn disabled_format_conversion_skip_reason(
transport: &GatewayProviderTransportSnapshot,
normalized_client_api_format: &str,
) -> Option<&'static str> {
let endpoint_api_format = transport.endpoint.api_format.trim();
if api_format_matches(endpoint_api_format, normalized_client_api_format) {
return None;
}
if crate::ai_pipeline::conversion::request_conversion_kind(
normalized_client_api_format,
endpoint_api_format,
)
.is_none()
{
return None;
}
if crate::ai_pipeline::conversion::request_conversion_requires_enable_flag(
normalized_client_api_format,
endpoint_api_format,
) && !crate::ai_pipeline::conversion::request_conversion_enabled_for_transport(
transport,
normalized_client_api_format,
endpoint_api_format,
) {
return Some("format_conversion_disabled");
}
None
}
fn current_local_execution_candidate_skip_reason_with_transport(
candidate: &SchedulerMinimalCandidateSelectionCandidate,
transport: &GatewayProviderTransportSnapshot,
normalized_client_api_format: &str,
requested_model: &str,
) -> Option<&'static str> {
if let Some(skip_reason) = current_local_execution_candidate_common_skip_reason_with_transport(
candidate,
transport,
Some(requested_model),
) {
return Some(skip_reason);
}
let endpoint_api_format = transport.endpoint.api_format.trim();
if api_format_matches(endpoint_api_format, normalized_client_api_format) {
return None;
}
if let Some(skip_reason) =
disabled_format_conversion_skip_reason(transport, normalized_client_api_format)
{
return Some(skip_reason);
}
if !crate::ai_pipeline::conversion::request_pair_allowed_for_transport(
transport,
normalized_client_api_format,
endpoint_api_format,
) {
return Some("transport_unsupported");
}
None
}
fn transport_key_supports_api_format(
transport: &GatewayProviderTransportSnapshot,
endpoint_api_format: &str,
) -> bool {
let provider_type = transport.provider.provider_type.trim();
let auth_type = transport.key.auth_type.trim();
let inherits_provider_api_formats = provider_type_is_fixed(provider_type)
&& (auth_type.eq_ignore_ascii_case("oauth")
|| (provider_type.eq_ignore_ascii_case("kiro")
&& auth_type.eq_ignore_ascii_case("bearer")
&& transport
.key
.decrypted_auth_config
.as_deref()
.map(str::trim)
.is_some_and(|value| !value.is_empty())));
if inherits_provider_api_formats {
return true;
}
match transport.key.api_formats.as_deref() {
None => true,
Some(formats) => formats
.iter()
.any(|value| api_format_matches(value, endpoint_api_format)),
}
}
fn normalize_api_format_alias(value: &str) -> String {
crate::ai_pipeline::normalize_legacy_openai_format_alias(value)
}
fn api_format_matches(left: &str, right: &str) -> bool {
normalize_api_format_alias(left) == normalize_api_format_alias(right)
}
fn transport_key_allows_candidate_model(
transport: &GatewayProviderTransportSnapshot,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> bool {
let Some(allowed_models) = transport.key.allowed_models.as_deref() else {
return true;
};
let requested_model = requested_model.trim();
let global_model_name = candidate.global_model_name.trim();
let selected_provider_model_name = candidate.selected_provider_model_name.trim();
let mapping_matched_model = candidate
.mapping_matched_model
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty());
for allowed_model in allowed_models.iter().map(String::as_str).map(str::trim) {
if allowed_model.is_empty() {
continue;
}
if allowed_model == requested_model
|| allowed_model == global_model_name
|| allowed_model == selected_provider_model_name
|| mapping_matched_model.is_some_and(|value| value == allowed_model)
{
return true;
}
}
false
}
pub(crate) async fn read_candidate_transport_snapshot(
state: PlannerAppState<'_>,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) -> Option<GatewayProviderTransportSnapshot> {
match state
.read_provider_transport_snapshot(
&candidate.provider_id,
&candidate.endpoint_id,
&candidate.key_id,
)
.await
{
Ok(Some(transport)) => Some(transport),
Ok(None) => None,
Err(error) => {
warn!(
event_name = "candidate_resolution_transport_load_failed",
log_type = "event",
provider_id = %candidate.provider_id,
endpoint_id = %candidate.endpoint_id,
key_id = %candidate.key_id,
error = ?error,
"failed to load provider transport while evaluating local candidate eligibility"
);
None
}
}
}

3
apps/aether-gateway/src/ai_pipeline/planner/mod.rs
View File

@@ -9,6 +9,7 @@ mod candidate_eligibility;
mod candidate_materialization;
mod candidate_metadata;
mod candidate_preparation;
mod candidate_resolution;
mod candidate_source;
mod common;
mod decision;
@@ -27,7 +28,7 @@ mod specialized;
mod standard;
mod state;
pub(crate) use self::candidate_eligibility::extract_pool_sticky_session_token;
pub(crate) use self::candidate_resolution::extract_pool_sticky_session_token;
pub(crate) use self::failure_diagnostic::{
CandidateFailureDiagnostic, CandidateFailureDiagnosticKind,
};

9
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/candidates.rs
View File

@@ -1,9 +1,5 @@
use tracing::warn;
use crate::ai_pipeline::planner::candidate_eligibility::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{
persist_available_local_execution_candidates_with_context,
persist_skipped_local_execution_candidates_with_context,
@@ -14,6 +10,10 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata_for_candidate,
LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::candidate_resolution::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
use crate::ai_pipeline::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
@@ -112,6 +112,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
candidates,
spec_metadata.api_format,
&input.requested_model,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
)

2
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/request/prepare.rs
View File

@@ -1,9 +1,9 @@
use std::sync::Arc;
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_preparation::{
resolve_candidate_mapped_model, resolve_candidate_oauth_auth, OauthPreparationContext,
};
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::spec_metadata::local_same_format_provider_spec_metadata;
use crate::ai_pipeline::transport::kiro::KiroRequestAuth;
use crate::ai_pipeline::transport::vertex::resolve_local_vertex_api_key_query_auth;

7
apps/aether-gateway/src/ai_pipeline/planner/pool_scheduler.rs
View File

@@ -7,7 +7,7 @@ use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKe
use serde_json::{Map, Value};
use tracing::warn;
use crate::ai_pipeline::planner::candidate_eligibility::{
use crate::ai_pipeline::planner::candidate_resolution::{
EligibleLocalExecutionCandidate, SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::PlannerAppState;
@@ -434,6 +434,7 @@ fn schedule_pool_group(
transport,
provider_api_format,
orchestration,
ranking,
} = eligible;
let key_id = candidate.key_id.clone();
let mut key_context = key_context_by_id.get(&key_id).cloned().unwrap_or_default();
@@ -495,6 +496,7 @@ fn schedule_pool_group(
transport,
provider_api_format,
orchestration,
ranking,
},
key_context,
original_index,
@@ -1040,7 +1042,7 @@ mod tests {
apply_local_execution_pool_scheduler_with_runtime_map, build_pool_catalog_key_context,
normalize_enabled_pool_presets, PoolCatalogKeyContext,
};
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::PlannerAppState;
use crate::data::GatewayDataState;
use crate::handlers::shared::provider_pool::{
@@ -1836,6 +1838,7 @@ mod tests {
},
provider_api_format: "openai:chat".to_string(),
orchestration: LocalExecutionCandidateMetadata::default(),
ranking: None,
transport: Arc::new(crate::ai_pipeline::GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: provider_id.to_string(),

3
apps/aether-gateway/src/ai_pipeline/planner/specialized/files/support.rs
View File

@@ -3,7 +3,6 @@ use serde_json::json;
use tracing::warn;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_eligibility::filter_and_rank_local_execution_candidates_without_transport_pair_gate;
use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate,
mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -15,6 +14,7 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_metadata,
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::candidate_resolution::filter_and_rank_local_execution_candidates_without_transport_pair_gate;
use crate::ai_pipeline::planner::decision_input::{
build_local_authenticated_decision_input, resolve_local_authenticated_decision_input,
};
@@ -95,6 +95,7 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
candidates,
GEMINI_FILES_CLIENT_API_FORMAT,
None,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(),
None,
)

9
apps/aether-gateway/src/ai_pipeline/planner/specialized/image/support.rs
View File

@@ -1,10 +1,6 @@
use tracing::warn;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_eligibility::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate,
mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -16,6 +12,10 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_metadata,
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::candidate_resolution::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
};
@@ -157,6 +157,7 @@ async fn materialize_local_openai_image_candidate_attempts(
candidates,
api_format,
&input.requested_model,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
)

9
apps/aether-gateway/src/ai_pipeline/planner/specialized/video/support.rs
View File

@@ -3,10 +3,6 @@ use tracing::warn;
use super::{LocalVideoCreateFamily, LocalVideoCreateSpec};
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_eligibility::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate,
mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -18,6 +14,10 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_metadata,
build_local_execution_candidate_metadata_for_candidate, LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::candidate_resolution::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
use crate::ai_pipeline::planner::decision_input::{
build_local_requested_model_decision_input, resolve_local_authenticated_decision_input,
@@ -169,6 +169,7 @@ async fn materialize_local_video_create_candidate_attempts(
candidates,
api_format,
&input.requested_model,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
)

9
apps/aether-gateway/src/ai_pipeline/planner/standard/family/candidates.rs
View File

@@ -3,10 +3,6 @@ use std::collections::BTreeSet;
use tracing::warn;
use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
use crate::ai_pipeline::planner::candidate_eligibility::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{
persist_available_local_execution_candidates_with_context,
persist_skipped_local_execution_candidates_with_context,
@@ -17,6 +13,10 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata_for_candidate,
LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::candidate_resolution::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
use crate::ai_pipeline::planner::common::extract_requested_model_from_request;
use crate::ai_pipeline::planner::decision_input::{
@@ -178,6 +178,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
candidates,
spec_metadata.api_format,
&input.requested_model,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
)

2
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/decision/request.rs
View File

@@ -4,10 +4,10 @@ use std::sync::Arc;
use serde_json::Value;
use crate::ai_pipeline::conversion::{request_conversion_direct_auth, request_conversion_kind};
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_preparation::{
prepare_header_authenticated_candidate, OauthPreparationContext,
};
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::common::OPENAI_CHAT_STREAM_PLAN_KIND;
use crate::ai_pipeline::planner::standard::{
apply_codex_openai_responses_special_headers, build_cross_format_openai_chat_request_body,

9
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/decision/support.rs
View File

@@ -1,10 +1,6 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::planner::candidate_eligibility::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -17,6 +13,10 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata_for_candidate,
LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::candidate_resolution::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::materialization_policy::{
build_local_candidate_persistence_policy, LocalCandidatePersistencePolicyKind,
};
@@ -135,6 +135,7 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
candidates,
"openai:chat",
&input.requested_model,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
)

2
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/chat/plans/candidates.rs
View File

@@ -4,7 +4,7 @@ use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use super::super::{GatewayError, LocalOpenAiChatDecisionInput};
use crate::ai_pipeline::conversion::request_candidate_api_formats;
use crate::ai_pipeline::planner::candidate_eligibility::SkippedLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
use crate::ai_pipeline::PlannerAppState;
use crate::clock::current_unix_secs;

2
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/responses/decision/request.rs
View File

@@ -5,10 +5,10 @@ use serde_json::Value;
use tracing::debug;
use crate::ai_pipeline::conversion::{request_conversion_direct_auth, request_conversion_kind};
use crate::ai_pipeline::planner::candidate_eligibility::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::candidate_preparation::{
prepare_header_authenticated_candidate, OauthPreparationContext,
};
use crate::ai_pipeline::planner::candidate_resolution::EligibleLocalExecutionCandidate;
use crate::ai_pipeline::planner::common::force_upstream_streaming_for_provider;
use crate::ai_pipeline::planner::spec_metadata::local_openai_responses_spec_metadata;
use crate::ai_pipeline::planner::standard::{

9
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/responses/decision/support.rs
View File

@@ -5,10 +5,6 @@ use tracing::warn;
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::ai_pipeline::conversion::{request_candidate_api_formats, request_conversion_kind};
use crate::ai_pipeline::planner::candidate_eligibility::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_materialization::{
mark_skipped_local_execution_candidate, mark_skipped_local_execution_candidate_with_extra_data,
mark_skipped_local_execution_candidate_with_failure_diagnostic,
@@ -21,6 +17,10 @@ use crate::ai_pipeline::planner::candidate_metadata::{
build_local_execution_candidate_contract_metadata_for_candidate,
LocalExecutionCandidateMetadataParts,
};
use crate::ai_pipeline::planner::candidate_resolution::{
extract_pool_sticky_session_token, filter_and_rank_local_execution_candidates,
SkippedLocalExecutionCandidate,
};
use crate::ai_pipeline::planner::candidate_source::auth_snapshot_allows_cross_format_candidate;
use crate::ai_pipeline::planner::common::extract_standard_requested_model;
use crate::ai_pipeline::planner::decision_input::{
@@ -234,6 +234,7 @@ pub(crate) async fn materialize_local_openai_responses_candidate_attempts(
candidates,
spec_metadata.api_format,
&input.requested_model,
Some(&input.auth_snapshot),
input.required_capabilities.as_ref(),
sticky_session_token.as_deref(),
)

71
apps/aether-gateway/src/data/candidate_selection.rs
View File

@@ -2,9 +2,10 @@ use aether_data::DataLayerError;
use aether_data_contracts::repository::candidate_selection::StoredMinimalCandidateSelectionRow;
use aether_scheduler_core::{
auth_constraints_allow_api_format, build_minimal_candidate_selection,
collect_global_model_names_for_required_capability, normalize_api_format,
resolve_requested_global_model_name, BuildMinimalCandidateSelectionInput,
SchedulerAuthConstraints, SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
collect_global_model_names_for_required_capability, enumerate_minimal_candidate_selection,
normalize_api_format, resolve_requested_global_model_name, row_supports_requested_model,
BuildMinimalCandidateSelectionInput, SchedulerAuthConstraints,
SchedulerMinimalCandidateSelectionCandidate, SchedulerPriorityMode,
};
use async_trait::async_trait;
use std::collections::BTreeSet;
@@ -30,31 +31,24 @@ pub(crate) async fn read_requested_model_rows(
api_format: &str,
requested_model_name: &str,
) -> Result<Option<(String, Vec<StoredMinimalCandidateSelectionRow>)>, DataLayerError> {
let exact_rows = state
.read_minimal_candidate_selection_rows_for_api_format_and_global_model(
api_format,
requested_model_name,
)
.await?;
if !exact_rows.is_empty() {
return Ok(Some((requested_model_name.to_string(), exact_rows)));
}
let rows = state
.read_minimal_candidate_selection_rows_for_api_format(api_format)
.await?;
let rows = rows
.into_iter()
.filter(|row| row_supports_requested_model(row, requested_model_name, api_format))
.collect::<Vec<_>>();
if rows.is_empty() {
return Ok(None);
}
let Some(resolved_global_model_name) =
resolve_requested_global_model_name(&rows, requested_model_name, api_format)
else {
return Ok(None);
};
Ok(Some((
resolved_global_model_name.clone(),
rows.into_iter()
.filter(|row| row.global_model_name == resolved_global_model_name)
.collect(),
)))
Ok(Some((resolved_global_model_name, rows)))
}
pub(crate) async fn read_minimal_candidate_selection(
@@ -180,6 +174,45 @@ pub(crate) async fn read_minimal_candidate_selection_with_priority_mode_and_affi
})
}
pub(crate) async fn enumerate_minimal_candidate_selection_with_required_capabilities(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
requested_model_name: &str,
require_streaming: bool,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
required_capabilities: Option<&serde_json::Value>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, DataLayerError> {
let normalized_api_format = normalize_api_format(api_format);
if normalized_api_format.is_empty() {
return Ok(Vec::new());
}
if !auth_constraints_allow_api_format(
auth_snapshot.map(auth_snapshot_constraints).as_ref(),
&normalized_api_format,
) {
return Ok(Vec::new());
}
let Some((resolved_global_model_name, rows)) =
read_requested_model_rows(state, &normalized_api_format, requested_model_name).await?
else {
return Ok(Vec::new());
};
let auth_constraints = auth_snapshot.map(auth_snapshot_constraints);
enumerate_minimal_candidate_selection(BuildMinimalCandidateSelectionInput {
rows,
normalized_api_format: &normalized_api_format,
requested_model_name,
resolved_global_model_name: resolved_global_model_name.as_str(),
require_streaming,
required_capabilities,
auth_constraints: auth_constraints.as_ref(),
affinity_key: None,
priority_mode: SchedulerPriorityMode::Provider,
})
}
pub(crate) async fn read_global_model_names_for_required_capability(
state: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,

9
apps/aether-gateway/src/executor/candidate_loop.rs
View File

@@ -271,7 +271,10 @@ where
fn should_skip_unused_persistence(report_context: Option<&serde_json::Value>) -> bool {
let metadata = local_execution_candidate_metadata_from_report_context(report_context);
metadata.candidate_group_id.is_some() && metadata.pool_key_index.is_some()
metadata.candidate_group_id.is_some()
&& metadata
.pool_key_index
.is_some_and(|pool_key_index| pool_key_index > 0)
}
fn resolve_stream_candidate_watchdog_timeout(plan: &aether_contracts::ExecutionPlan) -> Duration {
@@ -484,6 +487,10 @@ mod tests {
"candidate_group_id": "pool-group",
"pool_key_index": 1,
}))));
assert!(!should_skip_unused_persistence(Some(&json!({
"candidate_group_id": "pool-group",
"pool_key_index": 0,
}))));
assert!(!should_skip_unused_persistence(Some(&json!({
"candidate_group_id": "pool-group",
}))));

28
apps/aether-gateway/src/scheduler/candidate/enumeration.rs Normal file
View File

@@ -0,0 +1,28 @@
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::{
enumerate_minimal_candidate_selection_with_required_capabilities,
MinimalCandidateSelectionRowSource,
};
use crate::GatewayError;
pub(super) async fn enumerate_scheduler_candidates(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
api_format: &str,
global_model_name: &str,
require_streaming: bool,
required_capabilities: Option<&serde_json::Value>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
) -> Result<Vec<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
enumerate_minimal_candidate_selection_with_required_capabilities(
selection_row_source,
api_format,
global_model_name,
require_streaming,
auth_snapshot,
required_capabilities,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}

4
apps/aether-gateway/src/scheduler/candidate/mod.rs
View File

@@ -1,10 +1,12 @@
use self::affinity::candidate_affinity_hash;
use self::selection::{
collect_selectable_candidates, collect_selectable_candidates_with_skip_reasons,
};
use super::state::SchedulerRuntimeState;
mod affinity;
mod enumeration;
mod ranking;
mod resolution;
mod runtime;
mod selection;

76
apps/aether-gateway/src/scheduler/candidate/ranking.rs Normal file
View File

@@ -0,0 +1,76 @@
use aether_scheduler_core::{
apply_scheduler_candidate_ranking, effective_provider_key_health_score,
matches_affinity_target, provider_key_health_bucket,
requested_capability_priority_for_candidate, SchedulerAffinityTarget,
SchedulerRankableCandidate, SchedulerRankingContext, SchedulerRankingMode,
};
use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
use super::affinity::candidate_affinity_hash;
use super::runtime::CandidateRuntimeSelectionSnapshot;
use super::SchedulerMinimalCandidateSelectionCandidate;
pub(super) fn rank_scheduler_candidates(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
runtime_snapshot: &CandidateRuntimeSelectionSnapshot,
ordering_config: SchedulerOrderingConfig,
required_capabilities: Option<&serde_json::Value>,
priority_affinity_key: Option<&str>,
cached_affinity_target: Option<&SchedulerAffinityTarget>,
now_unix_secs: u64,
) {
let rankables = candidates
.iter()
.enumerate()
.map(|(index, candidate)| {
let provider_key = runtime_snapshot
.provider_key_rpm_states
.get(&candidate.key_id);
SchedulerRankableCandidate::from_candidate(candidate, index)
.with_capability_priority(requested_capability_priority_for_candidate(
required_capabilities,
candidate,
))
.with_cached_affinity_match(
cached_affinity_target
.is_some_and(|target| matches_affinity_target(candidate, target)),
)
.with_affinity_hash(
priority_affinity_key.map(|key| candidate_affinity_hash(key, candidate)),
)
.with_health(
provider_key.and_then(|key| {
provider_key_health_bucket(key, candidate.endpoint_api_format.as_str())
}),
provider_key
.and_then(|key| {
effective_provider_key_health_score(
key,
candidate.endpoint_api_format.as_str(),
)
})
.unwrap_or(1.0),
)
})
.collect::<Vec<_>>();
apply_scheduler_candidate_ranking(
candidates,
&rankables,
SchedulerRankingContext {
priority_mode: ordering_config.priority_mode,
ranking_mode: scheduler_ranking_mode(ordering_config.scheduling_mode),
include_health: true,
load_balance_seed: now_unix_secs,
},
);
}
fn scheduler_ranking_mode(mode: SchedulerSchedulingMode) -> SchedulerRankingMode {
match mode {
SchedulerSchedulingMode::FixedOrder => SchedulerRankingMode::FixedOrder,
SchedulerSchedulingMode::CacheAffinity => SchedulerRankingMode::CacheAffinity,
SchedulerSchedulingMode::LoadBalance => SchedulerRankingMode::LoadBalance,
}
}

45
apps/aether-gateway/src/scheduler/candidate/resolution.rs Normal file
View File

@@ -0,0 +1,45 @@
use std::collections::BTreeSet;
use aether_scheduler_core::SchedulerAffinityTarget;
use super::affinity::candidate_key;
use super::runtime::{current_candidate_runtime_skip_reason, CandidateRuntimeSelectionSnapshot};
use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerSkippedCandidate};
pub(super) fn resolve_scheduler_candidate_selectability(
candidates: Vec<SchedulerMinimalCandidateSelectionCandidate>,
runtime_snapshot: &CandidateRuntimeSelectionSnapshot,
now_unix_secs: u64,
cached_affinity_target: Option<&SchedulerAffinityTarget>,
) -> (
Vec<SchedulerMinimalCandidateSelectionCandidate>,
Vec<SchedulerSkippedCandidate>,
) {
let mut selected = Vec::with_capacity(candidates.len());
let mut skipped = Vec::new();
let mut emitted_selected_keys = BTreeSet::new();
let mut emitted_skipped_keys = BTreeSet::new();
for candidate in candidates {
let key = candidate_key(&candidate);
if let Some(skip_reason) = current_candidate_runtime_skip_reason(
&candidate,
runtime_snapshot,
now_unix_secs,
cached_affinity_target,
) {
if emitted_skipped_keys.insert(key) {
skipped.push(SchedulerSkippedCandidate {
candidate,
skip_reason,
});
}
continue;
}
if emitted_selected_keys.insert(key) {
selected.push(candidate);
}
}
(selected, skipped)
}

166
apps/aether-gateway/src/scheduler/candidate/selection.rs
View File

@@ -1,27 +1,15 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use aether_scheduler_core::{
collect_selectable_candidates_from_keys,
reorder_candidates_by_scheduler_health as reorder_candidates_by_scheduler_health_in_core,
SchedulerPriorityMode,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::{
read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities,
MinimalCandidateSelectionRowSource,
};
use crate::data::candidate_selection::MinimalCandidateSelectionRowSource;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::scheduler::config::SchedulerSchedulingMode;
use crate::GatewayError;
use super::affinity::{
build_scheduler_affinity_cache_key, candidate_key, remember_scheduler_affinity,
};
use super::affinity::{build_scheduler_affinity_cache_key, remember_scheduler_affinity};
use super::enumeration::enumerate_scheduler_candidates;
use super::ranking::rank_scheduler_candidates;
use super::resolution::resolve_scheduler_candidate_selectability;
use super::runtime::{
auth_snapshot_concurrency_limit_reached, current_candidate_runtime_skip_reason,
read_candidate_runtime_selection_snapshot,
auth_snapshot_concurrency_limit_reached, read_candidate_runtime_selection_snapshot,
};
use super::{SchedulerMinimalCandidateSelectionCandidate, SchedulerRuntimeState};
@@ -44,69 +32,6 @@ pub(super) fn is_exact_all_skipped_by_auth_limit(
.all(|candidate| candidate.skip_reason == API_KEY_CONCURRENCY_LIMIT_SKIP_REASON)
}
pub(super) fn reorder_candidates_by_scheduler_health(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
provider_key_rpm_states: &BTreeMap<String, StoredProviderCatalogKey>,
required_capabilities: Option<&serde_json::Value>,
affinity_key: Option<&str>,
priority_mode: SchedulerPriorityMode,
) {
reorder_candidates_by_scheduler_health_in_core(
candidates,
provider_key_rpm_states,
required_capabilities,
affinity_key,
priority_mode,
);
}
fn apply_load_balance_rotation(
candidates: &mut [SchedulerMinimalCandidateSelectionCandidate],
priority_mode: SchedulerPriorityMode,
now_unix_secs: u64,
) {
if candidates.len() < 2 {
return;
}
let mut start = 0usize;
while start < candidates.len() {
let mut end = start + 1;
while end < candidates.len()
&& candidates_share_load_balance_group(
&candidates[start],
&candidates[end],
priority_mode,
)
{
end += 1;
}
let group_len = end - start;
if group_len > 1 {
let offset = usize::try_from(now_unix_secs).unwrap_or(0) % group_len;
candidates[start..end].rotate_left(offset);
}
start = end;
}
}
fn candidates_share_load_balance_group(
left: &SchedulerMinimalCandidateSelectionCandidate,
right: &SchedulerMinimalCandidateSelectionCandidate,
priority_mode: SchedulerPriorityMode,
) -> bool {
match priority_mode {
SchedulerPriorityMode::Provider => {
left.provider_priority == right.provider_priority
&& left.key_internal_priority == right.key_internal_priority
}
SchedulerPriorityMode::GlobalKey => {
left.key_global_priority_for_format == right.key_global_priority_for_format
}
}
}
#[cfg_attr(not(test), allow(dead_code))]
pub(super) async fn select_minimal_candidate(
selection_row_source: &(impl MinimalCandidateSelectionRowSource + Sync),
@@ -182,36 +107,18 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
let ordering_config = runtime_state.read_scheduler_ordering_config().await?;
let priority_affinity_key =
scheduling_priority_affinity_key(auth_snapshot, ordering_config.scheduling_mode);
let mut candidates =
read_minimal_candidate_selection_with_priority_mode_and_affinity_key_and_required_capabilities(
let mut candidates = enumerate_scheduler_candidates(
selection_row_source,
api_format,
global_model_name,
require_streaming,
auth_snapshot,
ordering_config.priority_mode,
priority_affinity_key,
required_capabilities,
auth_snapshot,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
.await?;
let runtime_snapshot =
read_candidate_runtime_selection_snapshot(runtime_state, &candidates, now_unix_secs)
.await?;
reorder_candidates_by_scheduler_health(
&mut candidates,
&runtime_snapshot.provider_key_rpm_states,
required_capabilities,
priority_affinity_key,
ordering_config.priority_mode,
);
if ordering_config.scheduling_mode == SchedulerSchedulingMode::LoadBalance {
apply_load_balance_rotation(
&mut candidates,
ordering_config.priority_mode,
now_unix_secs,
);
}
let affinity_cache_key =
build_scheduler_affinity_cache_key(auth_snapshot, api_format, global_model_name);
let cached_affinity_target = if ordering_config.scheduling_mode
@@ -225,6 +132,15 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
};
if auth_snapshot_concurrency_limit_reached(auth_snapshot, &runtime_snapshot, now_unix_secs) {
rank_scheduler_candidates(
&mut candidates,
&runtime_snapshot,
ordering_config,
required_capabilities,
priority_affinity_key,
cached_affinity_target.as_ref(),
now_unix_secs,
);
return Ok((
Vec::new(),
candidates
@@ -237,37 +153,23 @@ pub(super) async fn collect_selectable_candidates_with_skip_reasons(
));
}
let mut selected_keys = BTreeSet::new();
let mut skipped = Vec::new();
let mut emitted_skipped_keys = BTreeSet::new();
let (mut selected, skipped) = resolve_scheduler_candidate_selectability(
candidates,
&runtime_snapshot,
now_unix_secs,
cached_affinity_target.as_ref(),
);
rank_scheduler_candidates(
&mut selected,
&runtime_snapshot,
ordering_config,
required_capabilities,
priority_affinity_key,
cached_affinity_target.as_ref(),
now_unix_secs,
);
for candidate in &candidates {
let key = candidate_key(candidate);
if let Some(skip_reason) = current_candidate_runtime_skip_reason(
candidate,
&runtime_snapshot,
now_unix_secs,
cached_affinity_target.as_ref(),
) {
if emitted_skipped_keys.insert(key) {
skipped.push(SchedulerSkippedCandidate {
candidate: candidate.clone(),
skip_reason,
});
}
continue;
}
selected_keys.insert(key);
}
Ok((
collect_selectable_candidates_from_keys(
candidates,
&selected_keys,
cached_affinity_target.as_ref(),
),
skipped,
))
Ok((selected, skipped))
}
fn scheduling_priority_affinity_key<'a>(

48
apps/aether-gateway/src/scheduler/candidate/tests/model.rs
View File

@@ -144,6 +144,54 @@ async fn read_minimal_candidate_selection_resolves_provider_model_alias() {
assert_eq!(selection[0].selected_provider_model_name, "gpt-5.2");
}
#[tokio::test]
async fn read_minimal_candidate_selection_keeps_all_rows_supporting_requested_model() {
let mut exact = sample_row();
exact.provider_id = "provider-exact".to_string();
exact.endpoint_id = "endpoint-exact".to_string();
exact.key_id = "key-exact".to_string();
exact.model_id = "model-exact".to_string();
exact.global_model_id = "global-exact".to_string();
exact.global_model_name = "gpt-5".to_string();
exact.model_provider_model_name = "gpt-5".to_string();
exact.model_provider_model_mappings = None;
let mut mapped = sample_row();
mapped.provider_id = "provider-mapped".to_string();
mapped.endpoint_id = "endpoint-mapped".to_string();
mapped.key_id = "key-mapped".to_string();
mapped.model_id = "model-mapped".to_string();
mapped.global_model_id = "global-mapped".to_string();
mapped.global_model_name = "claude-sonnet".to_string();
mapped.global_model_mappings = Some(vec!["gpt-5".to_string()]);
mapped.model_provider_model_name = "claude-sonnet-upstream".to_string();
mapped.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "claude-sonnet-upstream".to_string(),
priority: 1,
api_formats: Some(vec!["openai:chat".to_string()]),
}]);
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
exact, mapped,
]));
let quotas = Arc::new(InMemoryProviderQuotaRepository::seed(vec![]));
let state = GatewayDataState::with_candidate_selection_and_quota_for_tests(candidates, quotas);
let selection = read_minimal_candidate_selection(&state, "openai:chat", "gpt-5", false, None)
.await
.expect("selection should succeed");
let provider_ids = selection
.iter()
.map(|candidate| candidate.provider_id.as_str())
.collect::<Vec<_>>();
assert_eq!(provider_ids, vec!["provider-exact", "provider-mapped"]);
assert_eq!(
selection[1].selected_provider_model_name,
"claude-sonnet-upstream"
);
}
#[tokio::test]
async fn read_minimal_candidate_selection_allows_resolved_global_model_in_auth_snapshot() {
let mut row = sample_row();

8
apps/aether-gateway/src/scheduler/candidate/tests/selection.rs
View File

@@ -384,9 +384,9 @@ async fn fixed_order_disables_same_priority_affinity_hash_tiebreaker() {
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.provider_priority = 0;
second.provider_priority = 1;
second.key_internal_priority = 0;
second.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
second.key_global_priority_by_format = Some(json!({"openai:chat": 1}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,
@@ -438,9 +438,9 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable
second.endpoint_id = "endpoint-b".to_string();
second.key_id = "key-b".to_string();
second.key_name = "beta".to_string();
second.provider_priority = 1;
second.provider_priority = 0;
second.key_internal_priority = 0;
second.key_global_priority_by_format = Some(json!({"openai:chat": 1}));
second.key_global_priority_by_format = Some(json!({"openai:chat": 0}));
let candidates = Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
first, second,

132
apps/aether-gateway/src/tests/ai_execute/sync/chat/local_decision.rs
View File

@@ -476,6 +476,23 @@ async fn gateway_executes_openai_chat_sync_via_local_cross_format_gemini_candida
}
}
fn sample_disabled_conversion_candidate_row() -> StoredMinimalCandidateSelectionRow {
let mut row = sample_candidate_row();
row.provider_id = "provider-openai-chat-gemini-local-disabled".to_string();
row.provider_priority = 20;
row.endpoint_id = "endpoint-openai-chat-gemini-local-disabled".to_string();
row.key_id = "key-openai-chat-gemini-local-disabled".to_string();
row.model_id = "model-openai-chat-gemini-local-disabled".to_string();
row.global_model_id = "global-model-openai-chat-gemini-local-disabled".to_string();
row.model_provider_model_name = "gemini-2.5-flash-upstream".to_string();
row.model_provider_model_mappings = Some(vec![StoredProviderModelMapping {
name: "gemini-2.5-flash-upstream".to_string(),
priority: 2,
api_formats: Some(vec!["gemini:chat".to_string()]),
}]);
row
}
fn sample_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-chat-gemini-local-1".to_string(),
@@ -497,6 +514,27 @@ async fn gateway_executes_openai_chat_sync_via_local_cross_format_gemini_candida
)
}
fn sample_disabled_conversion_provider_catalog_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-openai-chat-gemini-local-disabled".to_string(),
"gemini".to_string(),
Some("https://example.com".to_string()),
"custom".to_string(),
)
.expect("provider should build")
.with_transport_fields(
true,
false,
false,
None,
Some(2),
None,
Some(20.0),
None,
None,
)
}
fn sample_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat-gemini-local-1".to_string(),
@@ -522,6 +560,29 @@ async fn gateway_executes_openai_chat_sync_via_local_cross_format_gemini_candida
.expect("endpoint transport should build")
}
fn sample_disabled_conversion_provider_catalog_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-openai-chat-gemini-local-disabled".to_string(),
"provider-openai-chat-gemini-local-disabled".to_string(),
"gemini:chat".to_string(),
Some("gemini".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://generativelanguage.googleapis.com".to_string(),
None,
None,
Some(2),
Some("/custom/v1beta/models/gemini-2.5-flash-upstream:generateContent".to_string()),
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-openai-chat-gemini-local-1".to_string(),
@@ -550,6 +611,34 @@ async fn gateway_executes_openai_chat_sync_via_local_cross_format_gemini_candida
.expect("key transport should build")
}
fn sample_disabled_conversion_provider_catalog_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-openai-chat-gemini-local-disabled".to_string(),
"provider-openai-chat-gemini-local-disabled".to_string(),
"prod".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(serde_json::json!(["gemini:chat"])),
encrypt_python_fernet_plaintext(
DEVELOPMENT_ENCRYPTION_KEY,
"sk-upstream-openai-chat-gemini-disabled",
)
.expect("api key should encrypt"),
None,
None,
Some(serde_json::json!({"gemini:chat": 2})),
None,
None,
None,
None,
)
.expect("key transport should build")
}
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeSyncRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let seen_report = Arc::new(Mutex::new(false));
@@ -730,11 +819,21 @@ async fn gateway_executes_openai_chat_sync_via_local_cross_format_gemini_candida
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_candidate_row(),
sample_disabled_conversion_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider()],
vec![sample_provider_catalog_endpoint()],
vec![sample_provider_catalog_key()],
vec![
sample_provider_catalog_provider(),
sample_disabled_conversion_provider_catalog_provider(),
],
vec![
sample_provider_catalog_endpoint(),
sample_disabled_conversion_provider_catalog_endpoint(),
],
vec![
sample_provider_catalog_key(),
sample_disabled_conversion_provider_catalog_key(),
],
));
let (upstream_url, upstream_handle) = start_server(upstream).await;
@@ -833,8 +932,31 @@ async fn gateway_executes_openai_chat_sync_via_local_cross_format_gemini_candida
.list_by_request_id("trace-openai-chat-gemini-local-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Success);
assert_eq!(stored_candidates.len(), 2);
assert_eq!(
stored_candidates
.iter()
.filter(|candidate| candidate.status == RequestCandidateStatus::Success)
.count(),
1
);
let skipped_candidate = stored_candidates
.iter()
.find(|candidate| candidate.status == RequestCandidateStatus::Skipped)
.expect("disabled conversion candidate should be persisted as skipped");
assert_eq!(
skipped_candidate.skip_reason.as_deref(),
Some("format_conversion_disabled")
);
let extra_data = skipped_candidate
.extra_data
.as_ref()
.expect("skipped cross-format candidate extra_data should exist");
assert_eq!(extra_data["execution_strategy"], "local_cross_format");
assert_eq!(
extra_data["transport_diagnostics"]["request_pair"]["conversion_enabled"],
false
);
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
assert!(

24
apps/aether-gateway/src/tests/architecture/ai_pipeline.rs
View File

@@ -423,6 +423,7 @@ fn ai_pipeline_planner_separates_local_candidate_eligibility_from_affinity_ranki
for pattern in [
"mod candidate_affinity;",
"mod candidate_eligibility;",
"mod candidate_resolution;",
"mod candidate_preparation;",
] {
assert!(
@@ -431,19 +432,30 @@ fn ai_pipeline_planner_separates_local_candidate_eligibility_from_affinity_ranki
);
}
let candidate_eligibility =
read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_eligibility.rs");
let candidate_resolution =
read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_resolution.rs");
for pattern in [
"pub(crate) async fn filter_and_rank_local_execution_candidates(",
"pub(crate) async fn filter_and_rank_local_execution_candidates_without_transport_pair_gate(",
"pub(crate) async fn read_candidate_transport_snapshot(",
] {
assert!(
candidate_eligibility.contains(pattern),
"planner/candidate_eligibility.rs should own {pattern}"
candidate_resolution.contains(pattern),
"planner/candidate_resolution.rs should own {pattern}"
);
}
let candidate_eligibility =
read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_eligibility.rs");
assert!(
candidate_eligibility.contains("pub(crate) use super::candidate_resolution::*;"),
"planner/candidate_eligibility.rs should remain a compatibility shim"
);
assert!(
!candidate_eligibility.contains("async fn filter_and_rank_local_execution_candidates("),
"planner/candidate_eligibility.rs should not keep resolution implementation"
);
let candidate_affinity =
read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_affinity.rs");
assert!(
@@ -1233,11 +1245,11 @@ fn ai_pipeline_specialized_files_attempts_consume_eligible_local_candidates_with
assert!(
specialized_files_support
.contains("filter_and_rank_local_execution_candidates_without_transport_pair_gate("),
"specialized files support should source runtime gating from candidate_eligibility"
"specialized files support should source runtime gating from candidate_resolution"
);
assert!(
!specialized_files_support.contains("rank_local_execution_candidates("),
"specialized files support should not bypass candidate_eligibility with raw affinity ranking"
"specialized files support should not bypass candidate_resolution with raw affinity ranking"
);
let specialized_files_decision = read_workspace_file(

22
apps/aether-gateway/src/tests/architecture/runtime_and_security.rs
View File

@@ -254,12 +254,20 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
"selection.rs should not depend on gateway-local SchedulerAffinityTarget"
);
assert!(
selection.contains("reorder_candidates_by_scheduler_health_in_core"),
"selection.rs should depend on core candidate reorder helper"
selection.contains("enumerate_scheduler_candidates("),
"selection.rs should delegate candidate enumeration"
);
assert!(
selection.contains("collect_selectable_candidates_from_keys"),
"selection.rs should depend on core selectable-candidate collector"
selection.contains("read_candidate_runtime_selection_snapshot("),
"selection.rs should delegate runtime snapshot loading"
);
assert!(
selection.contains("resolve_scheduler_candidate_selectability("),
"selection.rs should delegate selectability resolution"
);
assert!(
selection.contains("rank_scheduler_candidates("),
"selection.rs should delegate final ranking"
);
for pattern in [
"async fn collect_selectable_candidates(",
@@ -282,8 +290,10 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
"read_provider_concurrent_limits(",
"read_provider_key_rpm_states(",
"candidate_is_selectable_with_runtime_state",
"collect_selectable_candidates_from_keys",
"auth_api_key_concurrency_limit_reached",
"build_provider_concurrent_limit_map(",
"reorder_candidates_by_scheduler_health",
] {
assert!(
!selection.contains(pattern),
@@ -558,8 +568,8 @@ fn scheduler_candidate_runtime_paths_depend_on_scheduler_core_and_state_trait()
let planner_candidate_affinity =
read_workspace_file("apps/aether-gateway/src/ai_pipeline/planner/candidate_affinity.rs");
assert!(
planner_candidate_affinity
.contains("aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate"),
planner_candidate_affinity.contains("use aether_scheduler_core::{")
&& planner_candidate_affinity.contains("SchedulerMinimalCandidateSelectionCandidate"),
"planner/candidate_affinity.rs should depend directly on core minimal candidate DTO"
);
assert!(