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

refactor: isolate dispatch scheduling core

Browse Source
This commit is contained in:
fawney19
2026-05-12 13:15:41 +08:00
parent 81ee27cdea
commit fa73655134
50 changed files with 4467 additions and 3202 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
apps/aether-gateway/Cargo.toml
View File

@@ -16,6 +16,7 @@ aether-contracts.workspace = true
aether-crypto.workspace = true
aether-data.workspace = true
aether-data-contracts.workspace = true
aether-dispatch-core.workspace = true
aether-http.workspace = true
aether-model-fetch.workspace = true
aether-oauth.workspace = true

20
apps/aether-gateway/src/ai_serving/mod.rs
View File

@@ -20,8 +20,8 @@ pub(crate) use self::finalize::internal::{
SyncToStreamBridgeOutcome,
};
pub(crate) use self::planner::{
build_gemini_stream_plan_from_decision, build_gemini_sync_plan_from_decision,
build_local_gemini_files_stream_attempt_source_for_kind,
apply_local_runtime_candidate_terminal_reason, build_gemini_stream_plan_from_decision,
build_gemini_sync_plan_from_decision, build_local_gemini_files_stream_attempt_source_for_kind,
build_local_gemini_files_stream_plan_and_reports_for_kind,
build_local_gemini_files_sync_attempt_source_for_kind,
build_local_gemini_files_sync_plan_and_reports_for_kind,
@@ -46,13 +46,17 @@ pub(crate) use self::planner::{
build_provider_key_pool_score_upsert, build_standard_family_stream_attempt_source,
build_standard_family_stream_plan_and_reports, build_standard_family_sync_attempt_source,
build_standard_family_sync_plan_and_reports, build_standard_stream_plan_from_decision,
build_standard_sync_plan_from_decision, extract_pool_sticky_session_token,
maybe_build_stream_decision_payload, maybe_build_stream_plan_payload,
maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
planner_is_matching_stream_request, set_local_openai_chat_execution_exhausted_diagnostic,
build_standard_sync_plan_from_decision, candidate_auth_channel_skip_reason,
extract_pool_sticky_session_token, maybe_build_stream_decision_payload,
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
maybe_build_sync_plan_payload, planner_is_matching_stream_request,
provider_key_pool_score_scope, read_candidate_transport_snapshot,
record_local_runtime_candidate_skip_reason,
set_local_openai_chat_execution_exhausted_diagnostic,
set_local_openai_image_execution_exhausted_diagnostic, CandidateFailureDiagnostic,
CandidateFailureDiagnosticKind, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot,
LocalExecutionAttemptSource, LocalResolvedOAuthRequestAuth, PlannerAppState,
CandidateFailureDiagnosticKind, EligibleLocalExecutionCandidate, GatewayAuthApiKeySnapshot,
GatewayProviderTransportSnapshot, LocalExecutionAttemptSource, LocalExecutionCandidateKind,
LocalResolvedOAuthRequestAuth, PlannerAppState, SkippedLocalExecutionCandidate,
};
pub(crate) use self::pure::*;
pub(crate) use self::transport::{

286
apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
View File

@@ -4,14 +4,17 @@ use aether_ai_serving::{
run_ai_available_candidate_persistence, run_ai_candidate_materialization,
run_ai_skipped_candidate_persistence, AiAvailableCandidatePersistencePort,
AiCandidateMaterializationOutcome, AiCandidateMaterializationPort,
AiSkippedCandidatePersistencePort,
AiCandidatePreselectionOutcome, AiSkippedCandidatePersistencePort,
};
use aether_dispatch_core::{DispatchSequence, DispatchSequenceItem};
use aether_scheduler_core::{ClientSessionAffinity, SchedulerMinimalCandidateSelectionCandidate};
use async_trait::async_trait;
use serde_json::Value;
use std::collections::VecDeque;
use std::convert::Infallible;
use std::sync::Arc;
use std::time::Duration;
use tokio::time::Instant;
use tracing::warn;
use uuid::Uuid;
@@ -29,12 +32,16 @@ use crate::ai_serving::planner::runtime_miss::record_local_runtime_candidate_ski
use crate::ai_serving::planner::CandidateFailureDiagnostic;
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::current_unix_ms;
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::{AppState, GatewayError};
const POOL_KEY_RETRY_INDEX_STRIDE: u32 = 100;
const AUTH_API_KEY_CONCURRENCY_WAIT_BUDGET: Duration = Duration::from_millis(100);
const AUTH_API_KEY_CONCURRENCY_RETRY_DELAY: Duration = Duration::from_millis(10);
#[derive(Debug, Clone)]
pub(crate) struct LocalExecutionCandidateAttempt {
@@ -61,12 +68,12 @@ pub(crate) trait LocalExecutionAttemptSource<T>: Send {
enum LocalExecutionCandidateAttemptSourceItem<'a> {
Static {
attempts: VecDeque<LocalExecutionCandidateAttempt>,
attempts: DispatchSequence<LocalExecutionCandidateAttempt>,
},
Pool {
cursor: PoolKeyCursor<'a>,
candidate_index: u32,
pending_attempts: VecDeque<LocalExecutionCandidateAttempt>,
pending_attempts: DispatchSequence<LocalExecutionCandidateAttempt>,
},
RequestedModelPage {
cursor: Box<RequestedModelAttemptPageCursor<'a>>,
@@ -80,7 +87,7 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
let mut items = VecDeque::new();
if !attempts.is_empty() {
items.push_back(LocalExecutionCandidateAttemptSourceItem::Static {
attempts: VecDeque::from(attempts),
attempts: dispatch_sequence_from_attempts(attempts),
});
}
Self { items }
@@ -91,8 +98,8 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
let front = self.items.front_mut()?;
match front {
LocalExecutionCandidateAttemptSourceItem::Static { attempts } => {
if let Some(attempt) = attempts.pop_front() {
if attempts.is_empty() {
if let Some(attempt) = next_attempt_from_dispatch_sequence(attempts) {
if dispatch_sequence_exhausted(attempts) {
self.items.pop_front();
}
return Some(attempt);
@@ -104,7 +111,7 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
candidate_index,
pending_attempts,
} => {
if let Some(attempt) = pending_attempts.pop_front() {
if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) {
return Some(attempt);
}
let Some(candidate) = cursor.next_key().await else {
@@ -113,9 +120,12 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
self.items.pop_front();
continue;
};
*pending_attempts = build_unpersisted_local_execution_candidate_attempts(
candidate,
*candidate_index,
*pending_attempts = dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
candidate,
*candidate_index,
)
.into(),
);
}
LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { cursor } => {
@@ -311,10 +321,7 @@ where
}
fn build_extra_data(&self, candidate: &Self::Candidate) -> Option<Self::ExtraData> {
ai_candidate_extra_data_with_ranking(
(self.build_extra_data)(candidate),
candidate.ranking.as_ref(),
)
available_candidate_extra_data_with_dispatch_ref(candidate, &self.build_extra_data)
}
fn generate_candidate_id(&self) -> String {
@@ -506,15 +513,6 @@ where
.map(decorate_skipped_candidate)
.collect::<Vec<_>>();
let candidate_count = candidates.len() + skipped_candidate_count;
if persistence_policy.skipped.record_runtime_miss_diagnostic {
for skipped_candidate in &skipped_candidates {
record_local_runtime_candidate_skip_reason(
state.app(),
trace_id,
skipped_candidate.skip_reason,
);
}
}
if scheduler_cache_affinity_enabled {
remember_first_local_candidate_affinity(
@@ -526,6 +524,14 @@ where
&candidates,
);
}
persist_skipped_local_execution_candidates_with_context(
state.app(),
trace_id,
persistence_policy.skipped,
u32::try_from(candidates.len()).unwrap_or(u32::MAX),
skipped_candidates,
)
.await;
let (items, _) = build_logical_candidate_items(
state,
@@ -566,7 +572,9 @@ fn build_logical_candidate_items<'a>(
candidate_index,
);
if !attempts.is_empty() {
items.push_back(LocalExecutionCandidateAttemptSourceItem::Static { attempts });
items.push_back(LocalExecutionCandidateAttemptSourceItem::Static {
attempts: dispatch_sequence_from_attempts(attempts.into()),
});
}
}
LocalExecutionCandidateKind::PoolGroup => {
@@ -585,7 +593,7 @@ fn build_logical_candidate_items<'a>(
items.push_back(LocalExecutionCandidateAttemptSourceItem::Pool {
cursor,
candidate_index,
pending_attempts: VecDeque::new(),
pending_attempts: DispatchSequence::new(Vec::new()),
});
}
}
@@ -646,6 +654,10 @@ where
required_capabilities: required_capabilities.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(),
skipped_api_key_id: persistence_policy.skipped.api_key_id.to_string(),
skipped_required_capabilities: persistence_policy.skipped.required_capabilities.cloned(),
skipped_error_context: persistence_policy.skipped.error_context,
record_runtime_miss_diagnostic,
resolution_mode,
decorate_skipped_candidate,
@@ -655,6 +667,7 @@ where
next_candidate_index: 0,
remembered_affinity: false,
scheduler_cache_affinity_enabled,
auth_api_key_concurrency_wait_deadline: None,
};
cursor.load_next_page().await;
let candidate_count = cursor.candidate_count;
@@ -682,6 +695,10 @@ struct RequestedModelAttemptPageCursor<'a> {
required_capabilities: Option<Value>,
sticky_session_token: Option<String>,
request_auth_channel: Option<String>,
skipped_user_id: String,
skipped_api_key_id: String,
skipped_required_capabilities: Option<Value>,
skipped_error_context: &'static str,
record_runtime_miss_diagnostic: bool,
resolution_mode: LocalCandidateResolutionMode,
decorate_skipped_candidate: DecorateSkippedCandidateFn<'a>,
@@ -691,6 +708,7 @@ struct RequestedModelAttemptPageCursor<'a> {
next_candidate_index: u32,
remembered_affinity: bool,
scheduler_cache_affinity_enabled: bool,
auth_api_key_concurrency_wait_deadline: Option<Instant>,
}
impl<'a> RequestedModelAttemptPageCursor<'a> {
@@ -720,6 +738,15 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
}
};
if page_is_exact_auth_api_key_concurrency_limited(&page) {
if self.wait_for_auth_api_key_concurrency_retry().await {
continue;
}
self.persist_final_auth_api_key_concurrency_skips(page.skipped_candidates)
.await;
return false;
}
let (candidates, resolved_skipped) =
resolve_and_rank_logical_local_execution_candidates(
self.state,
@@ -740,18 +767,10 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
.chain(resolved_skipped)
.map(|skipped| (self.decorate_skipped_candidate)(skipped))
.collect::<Vec<_>>();
let skipped_candidate_count = skipped_candidates.len();
self.candidate_count = self
.candidate_count
.saturating_add(candidates.len() + skipped_candidates.len());
if self.record_runtime_miss_diagnostic {
for skipped_candidate in &skipped_candidates {
record_local_runtime_candidate_skip_reason(
self.state.app(),
&self.trace_id,
skipped_candidate.skip_reason,
);
}
}
.saturating_add(candidates.len() + skipped_candidate_count);
if self.scheduler_cache_affinity_enabled
&& !self.remembered_affinity
&& !candidates.is_empty()
@@ -776,13 +795,90 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
Some(&self.requested_model),
self.request_auth_channel.as_deref(),
);
self.next_candidate_index = next_candidate_index;
self.next_candidate_index = next_candidate_index
.saturating_add(u32::try_from(skipped_candidate_count).unwrap_or(u32::MAX));
if !items.is_empty() {
self.pending_items = items;
return true;
}
let skipped_starting_candidate_index = next_candidate_index;
let skipped_persistence = LocalSkippedCandidatePersistenceContext {
user_id: self.skipped_user_id.as_str(),
api_key_id: self.skipped_api_key_id.as_str(),
required_capabilities: self.skipped_required_capabilities.as_ref(),
error_context: self.skipped_error_context,
record_runtime_miss_diagnostic: self.record_runtime_miss_diagnostic,
};
persist_skipped_local_execution_candidates_with_context(
self.state.app(),
&self.trace_id,
skipped_persistence,
skipped_starting_candidate_index,
skipped_candidates,
)
.await;
}
}
async fn wait_for_auth_api_key_concurrency_retry(&mut self) -> bool {
let now = Instant::now();
let deadline = *self
.auth_api_key_concurrency_wait_deadline
.get_or_insert(now + AUTH_API_KEY_CONCURRENCY_WAIT_BUDGET);
if now >= deadline {
return false;
}
let sleep_duration =
AUTH_API_KEY_CONCURRENCY_RETRY_DELAY.min(deadline.saturating_duration_since(now));
tokio::time::sleep(sleep_duration).await;
self.page_cursor.restart_scan();
true
}
async fn persist_final_auth_api_key_concurrency_skips(
&mut self,
skipped_candidates: Vec<SkippedLocalExecutionCandidate>,
) {
let skipped_candidates = skipped_candidates
.into_iter()
.map(|skipped| (self.decorate_skipped_candidate)(skipped))
.collect::<Vec<_>>();
let skipped_candidate_count = skipped_candidates.len();
self.candidate_count = self.candidate_count.saturating_add(skipped_candidate_count);
let skipped_persistence = LocalSkippedCandidatePersistenceContext {
user_id: self.skipped_user_id.as_str(),
api_key_id: self.skipped_api_key_id.as_str(),
required_capabilities: self.skipped_required_capabilities.as_ref(),
error_context: self.skipped_error_context,
record_runtime_miss_diagnostic: self.record_runtime_miss_diagnostic,
};
persist_skipped_local_execution_candidates_with_context(
self.state.app(),
&self.trace_id,
skipped_persistence,
self.next_candidate_index,
skipped_candidates,
)
.await;
self.next_candidate_index = self
.next_candidate_index
.saturating_add(u32::try_from(skipped_candidate_count).unwrap_or(u32::MAX));
}
}
fn page_is_exact_auth_api_key_concurrency_limited(
page: &AiCandidatePreselectionOutcome<
SchedulerMinimalCandidateSelectionCandidate,
SkippedLocalExecutionCandidate,
>,
) -> bool {
page.candidates.is_empty()
&& !page.skipped_candidates.is_empty()
&& page
.skipped_candidates
.iter()
.all(|skipped| skipped.skip_reason == API_KEY_CONCURRENCY_LIMIT_SKIP_REASON)
}
async fn pop_attempt_from_items(
@@ -792,8 +888,8 @@ async fn pop_attempt_from_items(
let front = items.front_mut()?;
match front {
LocalExecutionCandidateAttemptSourceItem::Static { attempts } => {
if let Some(attempt) = attempts.pop_front() {
if attempts.is_empty() {
if let Some(attempt) = next_attempt_from_dispatch_sequence(attempts) {
if dispatch_sequence_exhausted(attempts) {
items.pop_front();
}
return Some(attempt);
@@ -805,7 +901,7 @@ async fn pop_attempt_from_items(
candidate_index,
pending_attempts,
} => {
if let Some(attempt) = pending_attempts.pop_front() {
if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) {
return Some(attempt);
}
let Some(candidate) = cursor.next_key().await else {
@@ -814,9 +910,12 @@ async fn pop_attempt_from_items(
items.pop_front();
continue;
};
*pending_attempts = build_unpersisted_local_execution_candidate_attempts(
candidate,
*candidate_index,
*pending_attempts = dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
candidate,
*candidate_index,
)
.into(),
);
}
LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { .. } => {
@@ -1005,7 +1104,7 @@ where
{
let attempt_slots = local_attempt_slot_count(&candidate.transport).max(1);
let extra_data = ai_candidate_extra_data_with_ranking(
build_extra_data(&candidate),
available_candidate_base_extra_data_with_dispatch_ref(&candidate, build_extra_data),
candidate.ranking.as_ref(),
);
let should_persist = should_persist_available_local_candidate(&candidate);
@@ -1057,6 +1156,70 @@ where
attempts
}
fn available_candidate_extra_data_with_dispatch_ref<F>(
candidate: &EligibleLocalExecutionCandidate,
build_extra_data: &F,
) -> Option<Value>
where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
{
ai_candidate_extra_data_with_ranking(
available_candidate_base_extra_data_with_dispatch_ref(candidate, build_extra_data),
candidate.ranking.as_ref(),
)
}
fn available_candidate_base_extra_data_with_dispatch_ref<F>(
candidate: &EligibleLocalExecutionCandidate,
build_extra_data: &F,
) -> Option<Value>
where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value> + Send + Sync,
{
let dispatch_ref = serde_json::to_value(dispatch_ref_for_local_candidate(candidate)).ok()?;
let mut object = match build_extra_data(candidate) {
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("dispatch_ref".to_string(), dispatch_ref);
Some(Value::Object(object))
}
fn dispatch_sequence_from_attempts(
attempts: Vec<LocalExecutionCandidateAttempt>,
) -> DispatchSequence<LocalExecutionCandidateAttempt> {
DispatchSequence::new(
attempts
.into_iter()
.map(|attempt| DispatchSequenceItem {
candidate_index: attempt.candidate_index,
retry_index: attempt.retry_index,
candidate: attempt,
mark: aether_dispatch_core::DispatchSequenceMark::Pending,
})
.collect(),
)
}
fn next_attempt_from_dispatch_sequence(
sequence: &mut DispatchSequence<LocalExecutionCandidateAttempt>,
) -> Option<LocalExecutionCandidateAttempt> {
let attempt = sequence.next()?.candidate.clone();
let _ = sequence.mark_succeeded();
Some(attempt)
}
fn dispatch_sequence_exhausted(
sequence: &mut DispatchSequence<LocalExecutionCandidateAttempt>,
) -> bool {
sequence.next().is_none()
}
fn build_unpersisted_local_execution_candidate_attempts(
candidate: EligibleLocalExecutionCandidate,
candidate_index: u32,
@@ -1457,6 +1620,16 @@ mod tests {
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[0].candidate_index, 2);
assert_eq!(
stored[0]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
.and_then(|value| value.get("SingleKey"))
.and_then(|value| value.get("key"))
.and_then(|value| value.get("key_id")),
Some(&json!("normal-key"))
);
}
#[test]
@@ -1562,6 +1735,16 @@ mod tests {
assert_eq!(stored.len(), 1);
assert_eq!(stored[0].key_id.as_deref(), Some("normal-key"));
assert_eq!(stored[0].candidate_index, 1);
assert_eq!(
stored[0]
.extra_data
.as_ref()
.and_then(|value| value.get("dispatch_ref"))
.and_then(|value| value.get("SingleKey"))
.and_then(|value| value.get("key"))
.and_then(|value| value.get("key_id")),
Some(&json!("normal-key"))
);
}
#[test]
@@ -1643,15 +1826,26 @@ mod tests {
Some(&json!("cached_affinity"))
);
assert_eq!(extra_data.get("demoted_by"), Some(&json!("cross_format")));
assert_eq!(
extra_data
.get("dispatch_ref")
.and_then(|value| value.get("SingleKey"))
.and_then(|value| value.get("key"))
.and_then(|value| value.get("key_id")),
Some(&json!("ranked-key"))
);
}
#[tokio::test]
async fn dynamic_attempt_source_does_not_drain_unexecuted_single_keys() {
let mut source = LocalExecutionCandidateAttemptSource {
items: VecDeque::from([LocalExecutionCandidateAttemptSourceItem::Static {
attempts: build_unpersisted_local_execution_candidate_attempts(
sample_eligible("normal-key", None),
0,
attempts: dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
sample_eligible("normal-key", None),
0,
)
.into(),
),
}]),
};

32
apps/aether-gateway/src/ai_serving/planner/candidate_resolution.rs
View File

@@ -21,7 +21,6 @@ use crate::ai_serving::{
use crate::orchestration::LocalExecutionCandidateMetadata;
use super::candidate_ranking::rank_eligible_local_execution_candidates;
use super::pool_scheduler::apply_local_execution_pool_scheduler;
#[derive(Debug, Clone, PartialEq)]
pub(crate) struct EligibleLocalExecutionCandidate {
@@ -61,7 +60,6 @@ struct GatewayLocalCandidateResolutionPort<'a> {
auth_snapshot: Option<&'a GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&'a ClientSessionAffinity>,
required_capabilities: Option<&'a serde_json::Value>,
sticky_session_token: Option<&'a str>,
request_auth_channel: Option<&'a str>,
}
@@ -182,14 +180,7 @@ impl AiCandidateResolutionPort for GatewayLocalCandidateResolutionPort<'_> {
&self,
candidates: Vec<Self::Eligible>,
) -> Result<(Vec<Self::Eligible>, Vec<Self::Skipped>), Self::Error> {
Ok(apply_local_execution_pool_scheduler(
self.state,
candidates,
self.sticky_session_token,
self.requested_model,
self.request_auth_channel,
)
.await)
Ok((candidates, Vec::new()))
}
}
@@ -201,7 +192,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>,
sticky_session_token: Option<&str>,
_sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
@@ -216,7 +207,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates(
auth_snapshot,
client_session_affinity,
required_capabilities,
sticky_session_token,
None,
request_auth_channel,
AiCandidateResolutionMode::Standard,
)
@@ -231,7 +222,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>,
sticky_session_token: Option<&str>,
_sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
) -> (
Vec<EligibleLocalExecutionCandidate>,
@@ -246,7 +237,7 @@ pub(crate) async fn resolve_and_rank_local_execution_candidates_without_transpor
auth_snapshot,
client_session_affinity,
required_capabilities,
sticky_session_token,
None,
request_auth_channel,
AiCandidateResolutionMode::WithoutTransportPairGate,
)
@@ -261,7 +252,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>,
sticky_session_token: Option<&str>,
_sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
mode: AiCandidateResolutionMode,
) -> (
@@ -276,7 +267,7 @@ pub(crate) async fn resolve_and_rank_logical_local_execution_candidates(
auth_snapshot,
client_session_affinity,
required_capabilities,
sticky_session_token,
None,
request_auth_channel,
mode,
false,
@@ -292,7 +283,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>,
sticky_session_token: Option<&str>,
_sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
mode: AiCandidateResolutionMode,
) -> (
@@ -307,10 +298,10 @@ async fn resolve_and_rank_local_execution_candidates_with_mode(
auth_snapshot,
client_session_affinity,
required_capabilities,
sticky_session_token,
None,
request_auth_channel,
mode,
true,
false,
)
.await
}
@@ -324,7 +315,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>,
sticky_session_token: Option<&str>,
_sticky_session_token: Option<&str>,
request_auth_channel: Option<&str>,
mode: AiCandidateResolutionMode,
expand_pool_groups: bool,
@@ -339,7 +330,6 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
auth_snapshot,
client_session_affinity,
required_capabilities,
sticky_session_token,
request_auth_channel,
};

10
apps/aether-gateway/src/ai_serving/planner/candidate_source.rs
View File

@@ -323,6 +323,16 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
Ok(None)
}
pub(crate) fn restart_scan(&mut self) {
self.format_index = 0;
self.requested_name_indexes.clear();
self.requested_name_offsets.clear();
self.scanned_rows_by_format.clear();
self.resolved_global_model_names.clear();
self.fallback_scanned_api_formats.clear();
self.seen_candidate_keys.clear();
}
async fn next_page_for_api_format(
&mut self,
candidate_api_format: &str,

8
apps/aether-gateway/src/ai_serving/planner/mod.rs
View File

@@ -26,6 +26,10 @@ mod standard;
mod state;
pub(crate) use self::candidate_materialization::LocalExecutionAttemptSource;
pub(crate) use self::candidate_resolution::{
candidate_auth_channel_skip_reason, read_candidate_transport_snapshot,
EligibleLocalExecutionCandidate, LocalExecutionCandidateKind, SkippedLocalExecutionCandidate,
};
pub(crate) use self::passthrough::{
build_local_same_format_stream_attempt_source, build_local_same_format_stream_plan_and_reports,
build_local_same_format_sync_attempt_source, build_local_same_format_sync_plan_and_reports,
@@ -38,7 +42,11 @@ pub(crate) use self::plan_builders::{
AiStreamAttempt, AiSyncAttempt,
};
pub(crate) use self::pool_scores::build_provider_key_pool_score_upsert;
pub(crate) use self::pool_scores::provider_key_pool_score_scope;
pub(crate) use self::route::is_matching_stream_request as planner_is_matching_stream_request;
pub(crate) use self::runtime_miss::{
apply_local_runtime_candidate_terminal_reason, record_local_runtime_candidate_skip_reason,
};
pub(crate) use self::specialized::{
build_local_gemini_files_stream_attempt_source_for_kind,
build_local_gemini_files_stream_plan_and_reports_for_kind,

10
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/family/build.rs
View File

@@ -9,7 +9,7 @@ use crate::{AiExecutionDecision, AppState, GatewayError};
use super::super::plans::{resolve_stream_spec, resolve_sync_spec};
use super::candidates::{
materialize_local_same_format_provider_candidate_attempts,
build_local_same_format_provider_candidate_attempt_source,
resolve_local_same_format_provider_decision_input,
};
use super::payload::maybe_build_local_same_format_provider_decision_payload_for_candidate;
@@ -55,7 +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 (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
@@ -65,7 +65,7 @@ pub(crate) async fn maybe_build_sync_local_same_format_provider_decision_payload
candidate_count,
);
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) =
maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
@@ -122,7 +122,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
@@ -132,7 +132,7 @@ pub(crate) async fn maybe_build_stream_local_same_format_provider_decision_paylo
candidate_count,
);
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) =
maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,

9
apps/aether-gateway/src/ai_serving/planner/passthrough/provider/plans.rs
View File

@@ -20,7 +20,6 @@ pub(crate) use crate::ai_serving::{
use super::{
build_local_same_format_provider_candidate_attempt_source,
materialize_local_same_format_provider_candidate_attempts,
maybe_build_local_same_format_provider_decision_payload_for_candidate,
resolve_local_same_format_provider_decision_input, AiStreamAttempt, AiSyncAttempt, AppState,
GatewayControlDecision, GatewayError, LocalSameFormatProviderCandidateAttempt,
@@ -348,7 +347,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
@@ -362,7 +361,7 @@ pub(crate) async fn build_local_sync_plan_and_reports(
}
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
@@ -433,7 +432,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) = materialize_local_same_format_provider_candidate_attempts(
let (mut source, candidate_count) = build_local_same_format_provider_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
@@ -447,7 +446,7 @@ pub(crate) async fn build_local_stream_plan_and_reports(
}
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_same_format_provider_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)

2728
apps/aether-gateway/src/ai_serving/planner/pool_scheduler.rs
View File

File diff suppressed because it is too large Load Diff

24
apps/aether-gateway/src/ai_serving/planner/specialized/files.rs
View File

@@ -318,10 +318,10 @@ pub(crate) async fn maybe_build_sync_local_gemini_files_decision_payload(
return Ok(None);
};
let attempts =
materialize_local_gemini_files_candidate_attempts(state, trace_id, &input).await?;
let (mut source, _) =
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
state,
parts,
@@ -359,11 +359,11 @@ pub(crate) async fn maybe_build_stream_local_gemini_files_decision_payload(
return Ok(None);
};
let attempts =
materialize_local_gemini_files_candidate_attempts(state, trace_id, &input).await?;
let (mut source, _) =
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
let empty_body_json = serde_json::Value::Null;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
state,
parts,
@@ -407,11 +407,11 @@ async fn build_local_sync_plan_and_reports(
return Ok(Vec::new());
};
let attempts =
materialize_local_gemini_files_candidate_attempts(state, trace_id, &input).await?;
let (mut source, _) =
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
state,
parts,
@@ -459,12 +459,12 @@ async fn build_local_stream_plan_and_reports(
return Ok(Vec::new());
};
let attempts =
materialize_local_gemini_files_candidate_attempts(state, trace_id, &input).await?;
let (mut source, _) =
build_local_gemini_files_candidate_attempt_source(state, trace_id, &input).await?;
let mut plans = Vec::new();
let empty_body_json = serde_json::Value::Null;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_gemini_files_decision_payload_for_candidate(
state,
parts,

24
apps/aether-gateway/src/ai_serving/planner/specialized/image.rs
View File

@@ -405,7 +405,7 @@ pub(crate) async fn maybe_build_sync_local_image_decision_payload(
return Ok(None);
};
let Some(attempts) = list_local_openai_image_candidate_attempts(
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
state,
trace_id,
&input,
@@ -413,12 +413,12 @@ pub(crate) async fn maybe_build_sync_local_image_decision_payload(
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
.await?
else {
return Ok(None);
};
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
state,
parts,
@@ -465,7 +465,7 @@ pub(crate) async fn maybe_build_stream_local_image_decision_payload(
return Ok(None);
};
let Some(attempts) = list_local_openai_image_candidate_attempts(
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
state,
trace_id,
&input,
@@ -473,12 +473,12 @@ pub(crate) async fn maybe_build_stream_local_image_decision_payload(
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
.await?
else {
return Ok(None);
};
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
state,
parts,
@@ -521,7 +521,7 @@ async fn build_local_sync_plan_and_reports(
return Ok(Vec::new());
};
let Some(attempts) = list_local_openai_image_candidate_attempts(
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
state,
trace_id,
&input,
@@ -529,13 +529,13 @@ async fn build_local_sync_plan_and_reports(
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
.await?
else {
return Ok(Vec::new());
};
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
state,
parts,
@@ -597,7 +597,7 @@ async fn build_local_stream_plan_and_reports(
return Ok(Vec::new());
};
let Some(attempts) = list_local_openai_image_candidate_attempts(
let Some((mut source, _)) = build_local_openai_image_candidate_attempt_source(
state,
trace_id,
&input,
@@ -605,13 +605,13 @@ async fn build_local_stream_plan_and_reports(
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
.await?
else {
return Ok(Vec::new());
};
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_openai_image_decision_payload_for_candidate(
state,
parts,

12
apps/aether-gateway/src/ai_serving/planner/specialized/video.rs
View File

@@ -180,7 +180,7 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
return Ok(None);
};
let Some(attempts) = list_local_video_create_candidate_attempts(
let Some((mut source, _)) = build_local_video_create_candidate_attempt_source(
state,
trace_id,
&input,
@@ -188,12 +188,12 @@ pub(crate) async fn maybe_build_sync_local_video_decision_payload(
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
.await?
else {
return Ok(None);
};
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
state, parts, body_json, trace_id, &input, attempt, spec,
)
@@ -223,7 +223,7 @@ async fn build_local_sync_plan_and_reports(
return Ok(Vec::new());
};
let Some(attempts) = list_local_video_create_candidate_attempts(
let Some((mut source, _)) = build_local_video_create_candidate_attempt_source(
state,
trace_id,
&input,
@@ -231,13 +231,13 @@ async fn build_local_sync_plan_and_reports(
spec_metadata.api_format,
spec_metadata.decision_kind,
)
.await
.await?
else {
return Ok(Vec::new());
};
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_video_create_decision_payload_for_candidate(
state, parts, body_json, trace_id, &input, attempt, spec,
)

27
apps/aether-gateway/src/ai_serving/planner/standard/family/build.rs
View File

@@ -20,8 +20,7 @@ use crate::ai_serving::GatewayControlDecision;
use crate::{AiExecutionDecision, AppState, GatewayError};
use super::candidates::{
build_local_standard_candidate_attempt_source, materialize_local_standard_candidate_attempts,
resolve_local_standard_decision_input,
build_local_standard_candidate_attempt_source, resolve_local_standard_decision_input,
};
use super::payload::maybe_build_local_standard_decision_payload_for_candidate;
use super::{LocalStandardDecisionInput, LocalStandardSpec};
@@ -323,12 +322,12 @@ pub(crate) async fn maybe_build_sync_via_standard_family_payload(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
let (mut source, candidate_count) =
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
.await?;
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
@@ -372,12 +371,12 @@ pub(crate) async fn maybe_build_stream_via_standard_family_payload(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
let (mut source, candidate_count) =
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
.await?;
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
@@ -427,15 +426,15 @@ pub(crate) async fn build_local_sync_plan_and_reports(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
let (mut source, candidate_count) =
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
.await?;
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
if candidate_count == 0 {
return Ok(Vec::new());
}
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
@@ -501,15 +500,15 @@ pub(crate) async fn build_local_stream_plan_and_reports(
Some(input.requested_model.as_str()),
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) =
materialize_local_standard_candidate_attempts(state, trace_id, &input, body_json, spec)
let (mut source, candidate_count) =
build_local_standard_candidate_attempt_source(state, trace_id, &input, body_json, spec)
.await?;
apply_local_runtime_candidate_evaluation_progress(state, trace_id, candidate_count);
if candidate_count == 0 {
return Ok(Vec::new());
}
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_standard_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)

75
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs
View File

@@ -1,28 +1,23 @@
use serde_json::Value;
use tracing::warn;
use crate::ai_serving::planner::common::{
OPENAI_CHAT_STREAM_PLAN_KIND, OPENAI_CHAT_SYNC_PLAN_KIND,
};
use crate::ai_serving::planner::runtime_miss::set_local_runtime_execution_exhausted_diagnostic;
use crate::ai_serving::GatewayControlDecision;
use crate::{AiExecutionDecision, AppState, GatewayError};
use tracing::warn;
mod decision;
mod plans;
use self::decision::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, LocalOpenAiChatCandidateAttempt,
LocalOpenAiChatCandidateAttemptSource, LocalOpenAiChatDecisionInput,
};
use self::plans::{
build_local_openai_chat_stream_attempt_source, build_local_openai_chat_stream_plan_and_reports,
build_local_openai_chat_sync_attempt_source, build_local_openai_chat_sync_plan_and_reports,
list_local_openai_chat_candidates, resolve_local_openai_chat_decision_input,
set_local_openai_chat_miss_diagnostic,
resolve_local_openai_chat_decision_input,
};
pub(crate) async fn build_local_openai_chat_sync_plan_and_reports_for_kind(
@@ -146,32 +141,17 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
return Ok(None);
};
let (candidates, skipped_candidates) =
match list_local_openai_chat_candidates(state, &input, false).await {
Ok(value) => value,
Err(err) => {
warn!(
event_name = "local_openai_chat_scheduler_selection_failed",
log_type = "event",
trace_id = %trace_id,
error = ?err,
"gateway local openai chat sync decision scheduler selection failed"
);
return Ok(None);
}
};
let attempts = materialize_local_openai_chat_candidate_attempts(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
let (mut source, _) = build_lazy_local_openai_chat_candidate_attempt_source(
state, trace_id, &input, body_json, false,
)
.await;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let upstream_is_stream = self::plans::openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
false,
);
if let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state,
parts,
@@ -181,7 +161,7 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
attempt,
OPENAI_CHAT_SYNC_PLAN_KIND,
"openai_chat_sync_success",
false,
upstream_is_stream,
)
.await
{
@@ -212,32 +192,17 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
return Ok(None);
};
let (candidates, skipped_candidates) =
match list_local_openai_chat_candidates(state, &input, true).await {
Ok(value) => value,
Err(err) => {
warn!(
event_name = "local_openai_chat_scheduler_selection_failed",
log_type = "event",
trace_id = %trace_id,
error = ?err,
"gateway local openai chat stream decision scheduler selection failed"
);
return Ok(None);
}
};
let attempts = materialize_local_openai_chat_candidate_attempts(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
let (mut source, _) = build_lazy_local_openai_chat_candidate_attempt_source(
state, trace_id, &input, body_json, true,
)
.await;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let upstream_is_stream = self::plans::openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
true,
);
if let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state,
parts,
@@ -247,7 +212,7 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
attempt,
OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success",
true,
upstream_is_stream,
)
.await
{

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

@@ -22,7 +22,7 @@ pub(super) use self::sync::{
build_local_openai_chat_sync_attempt_source, build_local_openai_chat_sync_plan_and_reports,
};
fn openai_chat_upstream_is_stream_for_candidate(
pub(super) fn openai_chat_upstream_is_stream_for_candidate(
transport: &GatewayProviderTransportSnapshot,
provider_api_format: &str,
client_is_stream: bool,

86
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs
View File

@@ -3,13 +3,10 @@ use tracing::warn;
use super::super::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
GatewayError, LocalOpenAiChatCandidateAttempt, LocalOpenAiChatCandidateAttemptSource,
LocalOpenAiChatDecisionInput,
};
use super::candidates::list_local_openai_chat_candidates;
use super::diagnostic::{
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
};
@@ -176,27 +173,12 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
return Ok(Vec::new());
};
let (candidates, skipped_candidates) =
match list_local_openai_chat_candidates(state, &input, true).await {
Ok(value) => value,
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai chat stream decision scheduler selection failed"
);
set_local_openai_chat_miss_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
"scheduler_selection_failed",
);
return Ok(Vec::new());
}
};
if candidates.is_empty() && skipped_candidates.is_empty() {
let Some((mut attempt_source, candidate_count)) =
build_local_openai_chat_stream_attempt_source(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await?
else {
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
@@ -206,59 +188,13 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
0,
);
return Ok(Vec::new());
}
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
candidates.len() + skipped_candidates.len(),
);
let attempts = materialize_local_openai_chat_candidate_attempts(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)
.await;
};
let mut plans = Vec::new();
for attempt in attempts {
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
true,
);
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state,
parts,
trace_id,
body_json,
&input,
attempt,
OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success",
upstream_is_stream,
)
.await
else {
continue;
};
match build_openai_chat_stream_plan_from_decision(parts, body_json, payload) {
Ok(Some(value)) => plans.push(value),
Ok(None) => {}
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai chat stream decision plan build failed"
);
}
while let Some(attempt) = attempt_source.next_execution_attempt().await? {
plans.push(attempt);
if plans.len() >= candidate_count {
break;
}
}

85
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs
View File

@@ -3,13 +3,10 @@ use tracing::warn;
use super::super::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
materialize_local_openai_chat_candidate_attempts,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
GatewayError, LocalOpenAiChatCandidateAttempt, LocalOpenAiChatCandidateAttemptSource,
LocalOpenAiChatDecisionInput,
};
use super::candidates::list_local_openai_chat_candidates;
use super::diagnostic::{
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
};
@@ -176,27 +173,11 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
return Ok(Vec::new());
};
let (candidates, skipped_candidates) =
match list_local_openai_chat_candidates(state, &input, false).await {
Ok(value) => value,
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai chat sync decision scheduler selection failed"
);
set_local_openai_chat_miss_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
"scheduler_selection_failed",
);
return Ok(Vec::new());
}
};
if candidates.is_empty() && skipped_candidates.is_empty() {
let Some((mut attempt_source, candidate_count)) = build_local_openai_chat_sync_attempt_source(
state, parts, trace_id, decision, body_json, plan_kind,
)
.await?
else {
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
@@ -206,59 +187,13 @@ pub(crate) async fn build_local_openai_chat_sync_plan_and_reports(
0,
);
return Ok(Vec::new());
}
set_local_openai_chat_candidate_evaluation_diagnostic(
state,
trace_id,
decision,
plan_kind,
Some(input.requested_model.as_str()),
candidates.len() + skipped_candidates.len(),
);
let attempts = materialize_local_openai_chat_candidate_attempts(
state,
trace_id,
&input,
body_json,
candidates,
skipped_candidates,
)
.await;
};
let mut plans = Vec::new();
for attempt in attempts {
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
&attempt.eligible.transport,
attempt.eligible.provider_api_format.as_str(),
false,
);
let Some(payload) = maybe_build_local_openai_chat_decision_payload_for_candidate(
state,
parts,
trace_id,
body_json,
&input,
attempt,
OPENAI_CHAT_SYNC_PLAN_KIND,
"openai_chat_sync_success",
upstream_is_stream,
)
.await
else {
continue;
};
match build_openai_chat_sync_plan_from_decision(parts, body_json, payload) {
Ok(Some(value)) => plans.push(value),
Ok(None) => {}
Err(err) => {
warn!(
trace_id = %trace_id,
error = ?err,
"gateway local openai chat sync decision plan build failed"
);
}
while let Some(attempt) = attempt_source.next_execution_attempt().await? {
plans.push(attempt);
if plans.len() >= candidate_count {
break;
}
}

10
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/mod.rs
View File

@@ -6,7 +6,7 @@ mod decision;
mod plans;
use self::decision::{
materialize_local_openai_responses_candidate_attempts,
build_local_openai_responses_candidate_attempt_source,
maybe_build_local_openai_responses_decision_payload_for_candidate,
resolve_local_openai_responses_decision_input,
};
@@ -108,12 +108,12 @@ pub(crate) async fn maybe_build_sync_local_openai_responses_decision_payload(
return Ok(None);
};
let (attempts, _) = materialize_local_openai_responses_candidate_attempts(
let (mut source, _) = build_local_openai_responses_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
@@ -146,12 +146,12 @@ pub(crate) async fn maybe_build_stream_local_openai_responses_decision_payload(
return Ok(None);
};
let (attempts, _) = materialize_local_openai_responses_candidate_attempts(
let (mut source, _) = build_local_openai_responses_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
if let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)

9
apps/aether-gateway/src/ai_serving/planner/standard/openai/responses/plans.rs
View File

@@ -3,7 +3,6 @@ use tracing::warn;
use super::decision::{
build_local_openai_responses_candidate_attempt_source,
materialize_local_openai_responses_candidate_attempts,
maybe_build_local_openai_responses_decision_payload_for_candidate,
resolve_local_openai_responses_decision_input, LocalOpenAiResponsesCandidateAttempt,
LocalOpenAiResponsesCandidateAttemptSource, LocalOpenAiResponsesDecisionInput,
@@ -312,7 +311,7 @@ pub(super) async fn build_local_sync_plan_and_reports(
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) = materialize_local_openai_responses_candidate_attempts(
let (mut source, candidate_count) = build_local_openai_responses_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
@@ -322,7 +321,7 @@ pub(super) async fn build_local_sync_plan_and_reports(
}
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)
@@ -384,7 +383,7 @@ pub(super) async fn build_local_stream_plan_and_reports(
"candidate_evaluation_incomplete",
);
let (attempts, candidate_count) = materialize_local_openai_responses_candidate_attempts(
let (mut source, candidate_count) = build_local_openai_responses_candidate_attempt_source(
state, trace_id, &input, body_json, spec,
)
.await?;
@@ -394,7 +393,7 @@ pub(super) async fn build_local_stream_plan_and_reports(
}
let mut plans = Vec::new();
for attempt in attempts {
while let Some(attempt) = source.next_attempt().await {
let Some(payload) = maybe_build_local_openai_responses_decision_payload_for_candidate(
state, parts, trace_id, body_json, &input, attempt, spec,
)

4
apps/aether-gateway/src/cache/auth_context.rs vendored
View File

@@ -28,4 +28,8 @@ impl AuthContextCache {
self.entries
.insert(cache_key, auth_context, ttl, max_entries);
}
pub(crate) fn clear(&self) {
self.entries.clear();
}
}

3
apps/aether-gateway/src/dispatch/mod.rs Normal file
View File

@@ -0,0 +1,3 @@
pub(crate) mod pool;
pub(crate) mod pool_scheduler;
pub(crate) mod refs;

5
apps/aether-gateway/src/dispatch/pool.rs Normal file
View File

@@ -0,0 +1,5 @@
use aether_dispatch_core::PoolWindowConfig;
pub(crate) fn default_pool_window_config() -> PoolWindowConfig {
PoolWindowConfig::default()
}

2716
apps/aether-gateway/src/dispatch/pool_scheduler.rs Normal file
View File

File diff suppressed because it is too large Load Diff

213
apps/aether-gateway/src/dispatch/refs.rs Normal file
View File

@@ -0,0 +1,213 @@
use aether_dispatch_core::{
DispatchCandidateRef, DispatchRankFacts, KeyRef, PoolRef, ProviderEndpointRef,
};
use crate::ai_serving::{EligibleLocalExecutionCandidate, LocalExecutionCandidateKind};
pub(crate) fn dispatch_ref_for_local_candidate(
eligible: &EligibleLocalExecutionCandidate,
) -> DispatchCandidateRef {
let rank = DispatchRankFacts {
provider_priority: eligible.candidate.provider_priority,
key_priority: Some(eligible.candidate.key_internal_priority),
ranking_reason: eligible.ranking.as_ref().and_then(|ranking| {
ranking
.promoted_by
.or(ranking.demoted_by)
.map(str::to_string)
}),
};
match eligible.kind {
LocalExecutionCandidateKind::SingleKey => DispatchCandidateRef::SingleKey {
key: key_ref_for_candidate(eligible),
rank,
},
LocalExecutionCandidateKind::PoolGroup => DispatchCandidateRef::PoolRef {
pool: pool_ref_for_candidate(eligible),
rank,
},
}
}
pub(crate) fn key_ref_for_candidate(eligible: &EligibleLocalExecutionCandidate) -> KeyRef {
KeyRef {
provider_id: eligible.candidate.provider_id.clone(),
endpoint_id: eligible.candidate.endpoint_id.clone(),
key_id: eligible.candidate.key_id.clone(),
model_id: eligible.candidate.model_id.clone(),
selected_provider_model_name: eligible.candidate.selected_provider_model_name.clone(),
api_format: eligible.candidate.endpoint_api_format.clone(),
}
}
pub(crate) fn pool_ref_for_candidate(eligible: &EligibleLocalExecutionCandidate) -> PoolRef {
PoolRef {
provider_id: eligible.candidate.provider_id.clone(),
endpoint_id: eligible.candidate.endpoint_id.clone(),
model_id: eligible.candidate.model_id.clone(),
selected_provider_model_name: eligible.candidate.selected_provider_model_name.clone(),
api_format: eligible.candidate.endpoint_api_format.clone(),
pool_group_id: eligible
.orchestration
.candidate_group_id
.clone()
.unwrap_or_else(|| pool_group_id_for_provider_endpoint(eligible)),
}
}
pub(crate) fn provider_endpoint_ref_for_candidate(
eligible: &EligibleLocalExecutionCandidate,
) -> ProviderEndpointRef {
ProviderEndpointRef {
provider_id: eligible.candidate.provider_id.clone(),
endpoint_id: eligible.candidate.endpoint_id.clone(),
model_id: eligible.candidate.model_id.clone(),
selected_provider_model_name: eligible.candidate.selected_provider_model_name.clone(),
api_format: eligible.candidate.endpoint_api_format.clone(),
}
}
fn pool_group_id_for_provider_endpoint(eligible: &EligibleLocalExecutionCandidate) -> String {
format!(
"provider={}|endpoint={}|model={}|selected_model={}|api_format={}",
eligible.candidate.provider_id,
eligible.candidate.endpoint_id,
eligible.candidate.model_id,
eligible.candidate.selected_provider_model_name,
eligible.candidate.endpoint_api_format
)
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aether_dispatch_core::DispatchCandidateRef;
use aether_provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider,
};
use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use super::dispatch_ref_for_local_candidate;
use crate::ai_serving::{EligibleLocalExecutionCandidate, LocalExecutionCandidateKind};
use crate::orchestration::LocalExecutionCandidateMetadata;
#[test]
fn pool_group_maps_to_pool_ref_without_exposing_internal_key() {
let eligible = sample_eligible(LocalExecutionCandidateKind::PoolGroup);
let dispatch_ref = dispatch_ref_for_local_candidate(&eligible);
match dispatch_ref {
DispatchCandidateRef::PoolRef { pool, rank } => {
assert_eq!(pool.provider_id, "provider-1");
assert_eq!(pool.endpoint_id, "endpoint-1");
assert_eq!(pool.pool_group_id, "group-1");
assert_eq!(rank.provider_priority, 10);
}
other => panic!("expected pool ref, got {other:?}"),
}
}
#[test]
fn single_key_maps_to_key_ref() {
let eligible = sample_eligible(LocalExecutionCandidateKind::SingleKey);
let dispatch_ref = dispatch_ref_for_local_candidate(&eligible);
match dispatch_ref {
DispatchCandidateRef::SingleKey { key, rank } => {
assert_eq!(key.key_id, "key-1");
assert_eq!(rank.key_priority, Some(7));
}
other => panic!("expected key ref, got {other:?}"),
}
}
fn sample_eligible(kind: LocalExecutionCandidateKind) -> EligibleLocalExecutionCandidate {
EligibleLocalExecutionCandidate {
kind,
candidate: SchedulerMinimalCandidateSelectionCandidate {
provider_id: "provider-1".to_string(),
provider_name: "Provider 1".to_string(),
provider_type: "openai".to_string(),
provider_priority: 10,
endpoint_id: "endpoint-1".to_string(),
endpoint_api_format: "openai:chat".to_string(),
key_id: "key-1".to_string(),
key_name: "Key 1".to_string(),
key_auth_type: "api_key".to_string(),
key_internal_priority: 7,
key_global_priority_for_format: None,
key_capabilities: None,
model_id: "model-1".to_string(),
global_model_id: "global-model-1".to_string(),
global_model_name: "gpt-5".to_string(),
selected_provider_model_name: "gpt-5".to_string(),
mapping_matched_model: None,
},
transport: Arc::new(crate::ai_serving::GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "Provider 1".to_string(),
provider_type: "openai".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: false,
concurrent_limit: None,
max_retries: None,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: None,
},
endpoint: GatewayProviderTransportEndpoint {
id: "endpoint-1".to_string(),
provider_id: "provider-1".to_string(),
api_format: "openai:chat".to_string(),
api_family: Some("openai".to_string()),
endpoint_kind: Some("chat".to_string()),
is_active: true,
base_url: "https://example.com".to_string(),
header_rules: None,
body_rules: None,
max_retries: None,
custom_path: None,
config: None,
format_acceptance_config: None,
proxy: None,
},
key: GatewayProviderTransportKey {
id: "key-1".to_string(),
provider_id: "provider-1".to_string(),
name: "Key 1".to_string(),
auth_type: "api_key".to_string(),
is_active: true,
api_formats: Some(vec!["openai:chat".to_string()]),
auth_type_by_format: None,
allow_auth_channel_mismatch_formats: None,
allowed_models: None,
capabilities: None,
rate_multipliers: None,
global_priority_by_format: None,
expires_at_unix_secs: None,
proxy: None,
fingerprint: None,
decrypted_api_key: "secret".to_string(),
decrypted_auth_config: None,
},
}),
provider_api_format: "openai:chat".to_string(),
orchestration: LocalExecutionCandidateMetadata {
candidate_group_id: Some("group-1".to_string()),
pool_key_index: None,
pool_key_lease: None,
scheduler_affinity_epoch: None,
},
ranking: None,
}
}
}

4
apps/aether-gateway/src/handlers/admin/observability/monitoring/tests/basics.rs
View File

@@ -265,8 +265,8 @@ async fn admin_monitoring_cache_stats_count_runtime_scheduler_affinities() {
"model-alpha",
)
.expect("scheduler affinity cache key should build");
state.scheduler_affinity_cache.insert(
affinity_cache_key,
state.remember_scheduler_affinity_target(
&affinity_cache_key,
crate::cache::SchedulerAffinityTarget {
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),

12
apps/aether-gateway/src/handlers/admin/observability/monitoring/tests/mod.rs
View File

@@ -239,8 +239,8 @@ async fn admin_monitoring_cache_affinities_and_delete_use_runtime_scheduler_affi
"model-alpha",
)
.expect("scheduler affinity cache key should build");
state.scheduler_affinity_cache.insert(
affinity_cache_key.clone(),
state.remember_scheduler_affinity_target(
&affinity_cache_key,
crate::cache::SchedulerAffinityTarget {
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
@@ -393,8 +393,8 @@ async fn admin_monitoring_cache_affinities_parse_session_scoped_scheduler_affini
.next()
.expect("session hash should exist")
.to_string();
state.scheduler_affinity_cache.insert(
affinity_cache_key.clone(),
state.remember_scheduler_affinity_target(
&affinity_cache_key,
crate::cache::SchedulerAffinityTarget {
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
@@ -403,8 +403,8 @@ async fn admin_monitoring_cache_affinities_parse_session_scoped_scheduler_affini
crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL,
128,
);
state.scheduler_affinity_cache.insert(
other_affinity_cache_key.clone(),
state.remember_scheduler_affinity_target(
&other_affinity_cache_key,
crate::cache::SchedulerAffinityTarget {
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),

1
apps/aether-gateway/src/lib.rs
View File

@@ -36,6 +36,7 @@ mod clock;
mod constants;
mod control;
mod data;
mod dispatch;
mod error;
mod execution_runtime;
mod executor;

20
apps/aether-gateway/src/orchestration/effects.rs
View File

@@ -1333,8 +1333,8 @@ mod tests {
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
state.scheduler_affinity_cache.insert(
cache_key.clone(),
state.remember_scheduler_affinity_target(
&cache_key,
SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
@@ -1376,8 +1376,8 @@ mod tests {
.expect("legacy scheduler affinity cache key should build");
for cache_key in [&session_cache_key, &legacy_cache_key] {
state.scheduler_affinity_cache.insert(
cache_key.to_string(),
state.remember_scheduler_affinity_target(
cache_key.as_str(),
SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
@@ -1422,8 +1422,8 @@ mod tests {
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
state.scheduler_affinity_cache.insert(
cache_key.clone(),
state.remember_scheduler_affinity_target(
&cache_key,
SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
@@ -1464,8 +1464,8 @@ mod tests {
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
state.scheduler_affinity_cache.insert(
cache_key.clone(),
state.remember_scheduler_affinity_target(
&cache_key,
SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),
@@ -1612,8 +1612,8 @@ mod tests {
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
.expect("scheduler affinity cache key should build");
state.scheduler_affinity_cache.insert(
cache_key.clone(),
state.remember_scheduler_affinity_target(
&cache_key,
SchedulerAffinityTarget {
provider_id: "prov-1".to_string(),
endpoint_id: "ep-1".to_string(),

4
apps/aether-gateway/src/rate_limit.rs
View File

@@ -120,6 +120,10 @@ impl FrontdoorUserRpmLimiter {
self.resolve_system_default_limit(state).await
}
pub(crate) fn clear_system_default_cache(&self) {
self.system_default_cache.clear();
}
pub(crate) fn current_bucket(&self, now_ts: u64) -> u64 {
self.config.current_bucket(now_ts)
}

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

@@ -31,7 +31,9 @@ use regex::Regex;
use sha2::{Digest, Sha256};
use std::collections::BTreeMap;
pub(crate) use self::selection::SchedulerSkippedCandidate;
pub(crate) use self::selection::{
SchedulerSkippedCandidate, API_KEY_CONCURRENCY_LIMIT_SKIP_REASON,
};
use crate::data::auth::GatewayAuthApiKeySnapshot;
use crate::data::candidate_selection::{

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

@@ -20,7 +20,7 @@ pub(crate) struct SchedulerSkippedCandidate {
pub(crate) skip_reason: &'static str,
}
pub(super) const API_KEY_CONCURRENCY_LIMIT_SKIP_REASON: &str = "api_key_concurrency_limit_reached";
pub(crate) const API_KEY_CONCURRENCY_LIMIT_SKIP_REASON: &str = "api_key_concurrency_limit_reached";
pub(super) fn is_exact_all_skipped_by_auth_limit(
selected: &[SchedulerMinimalCandidateSelectionCandidate],

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

@@ -210,8 +210,8 @@ async fn reuses_cached_scheduler_affinity_candidate_before_sorted_fallback() {
let cache_key =
build_scheduler_affinity_cache_key(Some(&auth_snapshot), "openai:chat", "gpt-4.1", None)
.expect("cache key should build");
state.scheduler_affinity_cache.insert(
cache_key,
state.remember_scheduler_affinity_target(
&cache_key,
SchedulerAffinityTarget {
provider_id: "provider-b".to_string(),
endpoint_id: "endpoint-b".to_string(),
@@ -315,8 +315,8 @@ async fn cached_affinity_candidate_cannot_use_reserved_provider_key_rpm_capacity
let cache_key =
build_scheduler_affinity_cache_key(Some(&auth_snapshot), "openai:chat", "gpt-4.1", None)
.expect("cache key should build");
state.scheduler_affinity_cache.insert(
cache_key,
state.remember_scheduler_affinity_target(
&cache_key,
SchedulerAffinityTarget {
provider_id: "provider-a".to_string(),
endpoint_id: "endpoint-a".to_string(),

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

@@ -176,8 +176,8 @@ async fn required_capability_without_model_uses_session_scoped_affinity() {
Some(&client_session_affinity),
)
.expect("session affinity cache key should build");
state.scheduler_affinity_cache.insert(
cache_key,
state.remember_scheduler_affinity_target(
&cache_key,
SchedulerAffinityTarget {
provider_id: "provider-b".to_string(),
endpoint_id: "endpoint-b".to_string(),

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

@@ -459,8 +459,8 @@ async fn fixed_order_ignores_cached_scheduler_affinity_promotion() {
);
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
state.scheduler_affinity_cache.insert(
"scheduler_affinity:affinity-key-1:openai:chat:gpt-4.1".to_string(),
state.remember_scheduler_affinity_target(
"scheduler_affinity:affinity-key-1:openai:chat:gpt-4.1",
SchedulerAffinityTarget {
provider_id: "provider-b".to_string(),
endpoint_id: "endpoint-b".to_string(),
@@ -578,8 +578,8 @@ async fn cache_affinity_promotes_cached_scheduler_affinity_candidate_when_enable
);
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
state.scheduler_affinity_cache.insert(
"scheduler_affinity:affinity-key-1:openai:chat:gpt-4.1".to_string(),
state.remember_scheduler_affinity_target(
"scheduler_affinity:affinity-key-1:openai:chat:gpt-4.1",
SchedulerAffinityTarget {
provider_id: "provider-b".to_string(),
endpoint_id: "endpoint-b".to_string(),
@@ -683,8 +683,8 @@ async fn load_balance_ignores_provider_priority_and_cached_affinity() {
);
let auth_snapshot = sample_auth_snapshot("affinity-key-1");
state.scheduler_affinity_cache.insert(
"scheduler_affinity:affinity-key-1:openai:chat:gpt-4.1".to_string(),
state.remember_scheduler_affinity_target(
"scheduler_affinity:affinity-key-1:openai:chat:gpt-4.1",
SchedulerAffinityTarget {
provider_id: "provider-b".to_string(),
endpoint_id: "endpoint-b".to_string(),

144
apps/aether-gateway/src/state/catalog.rs
View File

@@ -469,7 +469,7 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if created.is_some() {
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(created)
}
@@ -485,7 +485,7 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if created.is_some() {
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(created)
}
@@ -500,7 +500,7 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated.is_some() {
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(updated)
}
@@ -515,7 +515,7 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if deleted {
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(deleted)
}
@@ -550,7 +550,7 @@ impl AppState {
}
}
if !endpoint_ids.is_empty() || !key_ids.is_empty() {
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(())
}
@@ -565,7 +565,7 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if created.is_some() {
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(created)
}
@@ -580,7 +580,7 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated.is_some() {
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(updated)
}
@@ -595,7 +595,7 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if deleted {
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(deleted)
}
@@ -610,7 +610,7 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated.is_some() {
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(updated)
}
@@ -631,7 +631,7 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated {
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(updated)
}
@@ -672,7 +672,7 @@ impl AppState {
);
}
}
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(deleted)
}
@@ -806,8 +806,128 @@ impl AppState {
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated {
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(updated)
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::Duration;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use crate::cache::SchedulerAffinityTarget;
use crate::data::GatewayDataState;
use crate::AppState;
fn sample_provider() -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"Provider 1".to_string(),
Some("https://example.com".to_string()),
"openai".to_string(),
)
.expect("provider should build")
}
fn sample_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-1".to_string(),
"provider-1".to_string(),
"openai:chat".to_string(),
Some("openai".to_string()),
Some("chat".to_string()),
true,
)
.expect("endpoint should build")
.with_transport_fields(
"https://api.example.com/v1".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("endpoint transport should build")
}
fn sample_key() -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"Key 1".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build")
}
#[tokio::test]
async fn provider_catalog_update_invalidates_scheduler_affinity_and_transport_snapshot_cache() {
let provider = sample_provider();
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider.clone()],
vec![sample_endpoint()],
vec![sample_key()],
));
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository)
.with_encryption_key_for_tests("test-encryption-key"),
);
let snapshot = state
.read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1")
.await
.expect("provider transport should read")
.expect("provider transport should exist");
assert!(!snapshot.provider.keep_priority_on_conversion);
let cache_key = "scheduler_affinity:api-key-1:openai:chat:gpt-5";
let ttl = Duration::from_secs(300);
state.remember_scheduler_affinity_target(
cache_key,
SchedulerAffinityTarget {
provider_id: "provider-1".to_string(),
endpoint_id: "endpoint-1".to_string(),
key_id: "key-1".to_string(),
},
ttl,
128,
);
assert!(state
.read_scheduler_affinity_target(cache_key, ttl)
.is_some());
let initial_epoch = state.scheduler_affinity_epoch();
let mut updated_provider = provider;
updated_provider.keep_priority_on_conversion = true;
updated_provider.provider_priority = -10;
state
.update_provider_catalog_provider(&updated_provider)
.await
.expect("provider update should succeed")
.expect("provider should update");
assert!(state.scheduler_affinity_epoch() > initial_epoch);
assert!(state
.read_scheduler_affinity_target(cache_key, ttl)
.is_none());
let snapshot = state
.read_provider_transport_snapshot("provider-1", "endpoint-1", "key-1")
.await
.expect("provider transport should read after update")
.expect("provider transport should exist after update");
assert!(snapshot.provider.keep_priority_on_conversion);
}
}

160
apps/aether-gateway/src/state/core.rs
View File

@@ -59,6 +59,30 @@ use crate::maintenance::spawn_usage_cleanup_worker;
use crate::maintenance::spawn_wallet_daily_usage_aggregation_worker;
const SYSTEM_CONFIG_CACHE_TTL: Duration = Duration::from_secs(3);
const SCHEDULER_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &[
"enable_format_conversion",
"keep_priority_on_conversion",
"provider_priority_mode",
"scheduling_mode",
];
const AUTH_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] =
&[crate::constants::DEFAULT_USER_GROUP_CONFIG_KEY];
const FRONTDOOR_RPM_AFFECTING_SYSTEM_CONFIG_KEYS: &[&str] = &["rate_limit_per_minute"];
fn system_config_key_affects_scheduler(key: &str) -> bool {
let key = key.trim();
SCHEDULER_AFFECTING_SYSTEM_CONFIG_KEYS.contains(&key)
}
fn system_config_key_affects_auth(key: &str) -> bool {
let key = key.trim();
AUTH_AFFECTING_SYSTEM_CONFIG_KEYS.contains(&key)
}
fn system_config_key_affects_frontdoor_rpm(key: &str) -> bool {
let key = key.trim();
FRONTDOOR_RPM_AFFECTING_SYSTEM_CONFIG_KEYS.contains(&key)
}
impl AppState {
fn usage_worker_queue_for(
@@ -143,7 +167,10 @@ impl AppState {
pub(crate) fn replace_data_state(&mut self, data: Arc<GatewayDataState>) {
self.clear_provider_transport_snapshot_cache();
self.invalidate_scheduler_affinity_cache();
self.invalidate_auth_context_cache();
self.system_config_cache.clear();
self.frontdoor_user_rpm.clear_system_default_cache();
let data = Arc::new(
(*data)
.clone()
@@ -492,11 +519,7 @@ impl AppState {
.upsert_system_config_value(key, value, description)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
self.system_config_cache.insert(
key.to_string(),
Some(value.clone()),
SYSTEM_CONFIG_CACHE_TTL,
);
self.remember_system_config_write(key, Some(value.clone()));
Ok(value)
}
@@ -515,10 +538,13 @@ impl AppState {
value: &serde_json::Value,
description: Option<&str>,
) -> Result<crate::data::state::StoredSystemConfigEntry, GatewayError> {
self.data
let entry = self
.data
.upsert_system_config_entry(key, value, description)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
self.remember_system_config_write(entry.key.as_str(), Some(entry.value.clone()));
Ok(entry)
}
pub(crate) async fn delete_system_config_value(&self, key: &str) -> Result<bool, GatewayError> {
@@ -529,9 +555,41 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))?;
self.system_config_cache
.insert(key.to_string(), None, SYSTEM_CONFIG_CACHE_TTL);
if deleted && system_config_key_affects_scheduler(key) {
self.invalidate_scheduler_affinity_cache();
}
if deleted && system_config_key_affects_auth(key) {
self.invalidate_auth_context_cache();
}
if deleted && system_config_key_affects_frontdoor_rpm(key) {
self.frontdoor_user_rpm.clear_system_default_cache();
}
Ok(deleted)
}
pub(crate) fn invalidate_provider_routing_caches(&self) {
self.clear_provider_transport_snapshot_cache();
self.invalidate_scheduler_affinity_cache();
}
pub(crate) fn invalidate_auth_context_cache(&self) {
self.auth_context_cache.clear();
}
fn remember_system_config_write(&self, key: &str, value: Option<serde_json::Value>) {
self.system_config_cache
.insert(key.to_string(), value, SYSTEM_CONFIG_CACHE_TTL);
if system_config_key_affects_scheduler(key) {
self.invalidate_scheduler_affinity_cache();
}
if system_config_key_affects_auth(key) {
self.invalidate_auth_context_cache();
}
if system_config_key_affects_frontdoor_rpm(key) {
self.frontdoor_user_rpm.clear_system_default_cache();
}
}
pub(crate) async fn read_admin_system_stats(
&self,
) -> Result<aether_data::repository::system::AdminSystemStats, GatewayError> {
@@ -558,7 +616,7 @@ impl AppState {
| aether_data::repository::system::AdminSystemPurgeTarget::Stats
) {
self.system_config_cache.clear();
self.clear_provider_transport_snapshot_cache();
self.invalidate_provider_routing_caches();
}
Ok(summary)
}
@@ -1190,6 +1248,7 @@ mod tests {
use serde_json::json;
use super::AppState;
use crate::cache::SchedulerAffinityTarget;
use crate::data::GatewayDataState;
#[tokio::test]
@@ -1237,6 +1296,91 @@ mod tests {
);
}
#[tokio::test]
async fn system_config_entry_write_refreshes_cache_and_scheduler_affinity_for_routing_keys() {
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
GatewayDataState::disabled().with_system_config_values_for_tests([(
"keep_priority_on_conversion".to_string(),
json!(false),
)]),
);
let cache_key = "scheduler_affinity:api-key-1:openai:chat:gpt-5";
let ttl = std::time::Duration::from_secs(300);
assert_eq!(
state
.read_system_config_json_value("keep_priority_on_conversion")
.await
.expect("system config read should succeed"),
Some(json!(false))
);
state.remember_scheduler_affinity_target(
cache_key,
SchedulerAffinityTarget {
provider_id: "provider-old".to_string(),
endpoint_id: "endpoint-old".to_string(),
key_id: "key-old".to_string(),
},
ttl,
128,
);
assert!(state
.read_scheduler_affinity_target(cache_key, ttl)
.is_some());
let initial_epoch = state.scheduler_affinity_epoch();
state
.upsert_system_config_entry("keep_priority_on_conversion", &json!(true), None)
.await
.expect("admin config write should succeed");
assert_eq!(
state
.read_system_config_json_value("keep_priority_on_conversion")
.await
.expect("system config read should use refreshed cache"),
Some(json!(true))
);
assert!(state.scheduler_affinity_epoch() > initial_epoch);
assert_eq!(state.read_scheduler_affinity_target(cache_key, ttl), None);
}
#[tokio::test]
async fn system_config_write_refreshes_frontdoor_rpm_default_cache() {
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
GatewayDataState::disabled().with_system_config_values_for_tests([(
"rate_limit_per_minute".to_string(),
json!(1),
)]),
);
assert_eq!(
state
.frontdoor_user_rpm()
.current_system_default_limit(&state)
.await
.expect("default rpm limit should read"),
1
);
state
.upsert_system_config_entry("rate_limit_per_minute", &json!(0), None)
.await
.expect("rpm system config should update");
assert_eq!(
state
.frontdoor_user_rpm()
.current_system_default_limit(&state)
.await
.expect("default rpm limit should use refreshed value"),
0
);
}
#[tokio::test]
async fn replacing_data_state_clears_system_config_cache() {
let mut state = AppState::new()

99
apps/aether-gateway/src/state/runtime/api_key_exports.rs
View File

@@ -182,10 +182,15 @@ impl AppState {
record: aether_data::repository::auth::CreateUserApiKeyRecord,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.create_user_api_key(record)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn create_standalone_api_key(
@@ -193,10 +198,15 @@ impl AppState {
record: aether_data::repository::auth::CreateStandaloneApiKeyRecord,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.create_standalone_api_key(record)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn update_user_api_key_basic(
@@ -204,10 +214,15 @@ impl AppState {
record: aether_data::repository::auth::UpdateUserApiKeyBasicRecord,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.update_user_api_key_basic(record)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn update_standalone_api_key_basic(
@@ -215,10 +230,15 @@ impl AppState {
record: aether_data::repository::auth::UpdateStandaloneApiKeyBasicRecord,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.update_standalone_api_key_basic(record)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn set_user_api_key_active(
@@ -228,10 +248,15 @@ impl AppState {
is_active: bool,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.set_user_api_key_active(user_id, api_key_id, is_active)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn set_standalone_api_key_active(
@@ -240,10 +265,15 @@ impl AppState {
is_active: bool,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.set_standalone_api_key_active(api_key_id, is_active)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn set_user_api_key_locked(
@@ -252,10 +282,15 @@ impl AppState {
api_key_id: &str,
is_locked: bool,
) -> Result<bool, GatewayError> {
self.data
let updated = self
.data
.set_user_api_key_locked(user_id, api_key_id, is_locked)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if updated {
self.invalidate_auth_context_cache();
}
Ok(updated)
}
pub(crate) async fn set_user_api_key_allowed_providers(
@@ -265,10 +300,15 @@ impl AppState {
allowed_providers: Option<Vec<String>>,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.set_user_api_key_allowed_providers(user_id, api_key_id, allowed_providers)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn set_user_api_key_force_capabilities(
@@ -278,10 +318,15 @@ impl AppState {
force_capabilities: Option<serde_json::Value>,
) -> Result<Option<aether_data::repository::auth::StoredAuthApiKeyExportRecord>, GatewayError>
{
self.data
let api_key = self
.data
.set_user_api_key_force_capabilities(user_id, api_key_id, force_capabilities)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if api_key.is_some() {
self.invalidate_auth_context_cache();
}
Ok(api_key)
}
pub(crate) async fn delete_user_api_key(
@@ -289,19 +334,29 @@ impl AppState {
user_id: &str,
api_key_id: &str,
) -> Result<bool, GatewayError> {
self.data
let deleted = self
.data
.delete_user_api_key(user_id, api_key_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if deleted {
self.invalidate_auth_context_cache();
}
Ok(deleted)
}
pub(crate) async fn delete_standalone_api_key(
&self,
api_key_id: &str,
) -> Result<bool, GatewayError> {
self.data
let deleted = self
.data
.delete_standalone_api_key(api_key_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if deleted {
self.invalidate_auth_context_cache();
}
Ok(deleted)
}
}

175
apps/aether-gateway/src/state/runtime/auth/user_lifecycle.rs
View File

@@ -248,10 +248,15 @@ impl AppState {
&self,
record: aether_data::repository::users::UpsertUserGroupRecord,
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, GatewayError> {
self.data
let group = self
.data
.create_user_group(record)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if group.is_some() {
self.invalidate_auth_context_cache();
}
Ok(group)
}
pub(crate) async fn update_user_group(
@@ -259,17 +264,27 @@ impl AppState {
group_id: &str,
record: aether_data::repository::users::UpsertUserGroupRecord,
) -> Result<Option<aether_data::repository::users::StoredUserGroup>, GatewayError> {
self.data
let group = self
.data
.update_user_group(group_id, record)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if group.is_some() {
self.invalidate_auth_context_cache();
}
Ok(group)
}
pub(crate) async fn delete_user_group(&self, group_id: &str) -> Result<bool, GatewayError> {
self.data
let deleted = self
.data
.delete_user_group(group_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if deleted {
self.invalidate_auth_context_cache();
}
Ok(deleted)
}
pub(crate) async fn list_user_group_members(
@@ -287,10 +302,13 @@ impl AppState {
group_id: &str,
user_ids: &[String],
) -> Result<Vec<aether_data::repository::users::StoredUserGroupMember>, GatewayError> {
self.data
let members = self
.data
.replace_user_group_members(group_id, user_ids)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
self.invalidate_auth_context_cache();
Ok(members)
}
pub(crate) async fn list_user_groups_for_user(
@@ -318,10 +336,13 @@ impl AppState {
user_id: &str,
group_ids: &[String],
) -> Result<Vec<aether_data::repository::users::StoredUserGroup>, GatewayError> {
self.data
let groups = self
.data
.replace_user_groups_for_user(user_id, group_ids)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
self.invalidate_auth_context_cache();
Ok(groups)
}
pub(crate) async fn add_user_to_group(
@@ -329,10 +350,15 @@ impl AppState {
group_id: &str,
user_id: &str,
) -> Result<bool, GatewayError> {
self.data
let added = self
.data
.add_user_to_group(group_id, user_id)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if added {
self.invalidate_auth_context_cache();
}
Ok(added)
}
pub(crate) async fn is_other_user_auth_email_taken(
@@ -417,13 +443,19 @@ impl AppState {
.lock()
.expect("auth user store should lock")
.insert(user.id.clone(), user.clone());
self.invalidate_auth_context_cache();
return Ok(Some(user));
}
self.data
let user = self
.data
.update_local_auth_user_profile(user_id, email, username)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if user.is_some() {
self.invalidate_auth_context_cache();
}
Ok(user)
}
pub(crate) async fn update_local_auth_user_password_hash(
@@ -606,10 +638,14 @@ impl AppState {
user.is_active = is_active;
}
let _ = (rate_limit_present, rate_limit);
return Ok(Some(user.clone()));
let user = user.clone();
drop(guard);
self.invalidate_auth_context_cache();
return Ok(Some(user));
}
self.data
let user = self
.data
.update_local_auth_user_admin_fields(
user_id,
role,
@@ -624,7 +660,11 @@ impl AppState {
is_active,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if user.is_some() {
self.invalidate_auth_context_cache();
}
Ok(user)
}
pub(crate) async fn update_local_auth_user_policy_modes(
@@ -651,10 +691,14 @@ impl AppState {
user.allowed_models_mode = mode;
}
let _ = rate_limit_mode;
return Ok(Some(user.clone()));
let user = user.clone();
drop(guard);
self.invalidate_auth_context_cache();
return Ok(Some(user));
}
self.data
let user = self
.data
.update_local_auth_user_policy_modes(
user_id,
allowed_providers_mode,
@@ -663,7 +707,11 @@ impl AppState {
rate_limit_mode,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if user.is_some() {
self.invalidate_auth_context_cache();
}
Ok(user)
}
pub(crate) async fn touch_auth_user_last_login(
@@ -821,3 +869,90 @@ fn normalized_user_group_ids(group_ids: &[String]) -> BTreeSet<String> {
.map(ToOwned::to_owned)
.collect()
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::Duration;
use aether_data::repository::users::{InMemoryUserReadRepository, UpsertUserGroupRecord};
use crate::control::GatewayControlAuthContext;
use crate::data::GatewayDataState;
use crate::AppState;
fn user_group_record(
allowed_models: Option<Vec<&str>>,
allowed_models_mode: &str,
) -> UpsertUserGroupRecord {
UpsertUserGroupRecord {
name: "Team".to_string(),
description: None,
priority: 0,
allowed_providers: None,
allowed_providers_mode: "unrestricted".to_string(),
allowed_api_formats: None,
allowed_api_formats_mode: "unrestricted".to_string(),
allowed_models: allowed_models.map(|values| {
values
.into_iter()
.map(ToOwned::to_owned)
.collect::<Vec<_>>()
}),
allowed_models_mode: allowed_models_mode.to_string(),
rate_limit: None,
rate_limit_mode: "inherit".to_string(),
}
}
fn cached_auth_context() -> GatewayControlAuthContext {
GatewayControlAuthContext {
user_id: "user-1".to_string(),
api_key_id: "key-1".to_string(),
username: Some("alice".to_string()),
api_key_name: Some("default".to_string()),
balance_remaining: None,
access_allowed: true,
user_rate_limit: None,
api_key_rate_limit: None,
api_key_is_standalone: false,
admin_bypass_limits: false,
local_rejection: None,
allowed_models: Some(vec!["gpt-4.1".to_string()]),
}
}
#[tokio::test]
async fn updating_user_group_invalidates_cached_auth_context() {
let repository = Arc::new(InMemoryUserReadRepository::default());
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(GatewayDataState::with_user_reader_for_tests(repository));
let group = state
.create_user_group(user_group_record(Some(vec!["gpt-4.1"]), "specific"))
.await
.expect("group should create")
.expect("group should exist");
let cache_key = "auth-context-cache-key".to_string();
let ttl = Duration::from_secs(60);
state
.auth_context_cache
.insert(cache_key.clone(), cached_auth_context(), ttl, 10);
assert!(state
.auth_context_cache
.get_fresh(&cache_key, ttl)
.is_some());
state
.update_user_group(&group.id, user_group_record(None, "unrestricted"))
.await
.expect("group should update")
.expect("group should exist after update");
assert!(state
.auth_context_cache
.get_fresh(&cache_key, ttl)
.is_none());
}
}

76
apps/aether-gateway/src/state/runtime/auth/user_provisioning.rs
View File

@@ -234,13 +234,19 @@ impl AppState {
.lock()
.expect("auth wallet store should lock")
.insert(wallet.id.clone(), wallet.clone());
self.invalidate_auth_context_cache();
return Ok(Some(wallet));
}
self.data
let wallet = self
.data
.initialize_auth_user_wallet(user_id, initial_gift_usd, unlimited)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if wallet.is_some() {
self.invalidate_auth_context_cache();
}
Ok(wallet)
}
pub(crate) async fn initialize_auth_api_key_wallet(
@@ -284,13 +290,19 @@ impl AppState {
.lock()
.expect("auth wallet store should lock")
.insert(wallet.id.clone(), wallet.clone());
self.invalidate_auth_context_cache();
return Ok(Some(wallet));
}
self.data
let wallet = self
.data
.initialize_auth_api_key_wallet(api_key_id, initial_gift_usd, unlimited)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if wallet.is_some() {
self.invalidate_auth_context_cache();
}
Ok(wallet)
}
pub(crate) async fn update_auth_user_wallet_limit_mode(
@@ -310,13 +322,21 @@ impl AppState {
let _ = wallet_id;
wallet.limit_mode = limit_mode.to_string();
wallet.updated_at_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
return Ok(Some(wallet.clone()));
let wallet = wallet.clone();
drop(guard);
self.invalidate_auth_context_cache();
return Ok(Some(wallet));
}
self.data
let wallet = self
.data
.update_auth_user_wallet_limit_mode(user_id, limit_mode)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if wallet.is_some() {
self.invalidate_auth_context_cache();
}
Ok(wallet)
}
pub(crate) async fn update_auth_api_key_wallet_limit_mode(
@@ -336,13 +356,21 @@ impl AppState {
let _ = wallet_id;
wallet.limit_mode = limit_mode.to_string();
wallet.updated_at_unix_secs = chrono::Utc::now().timestamp().max(0) as u64;
return Ok(Some(wallet.clone()));
let wallet = wallet.clone();
drop(guard);
self.invalidate_auth_context_cache();
return Ok(Some(wallet));
}
self.data
let wallet = self
.data
.update_auth_api_key_wallet_limit_mode(api_key_id, limit_mode)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if wallet.is_some() {
self.invalidate_auth_context_cache();
}
Ok(wallet)
}
#[allow(clippy::too_many_arguments)]
@@ -381,10 +409,14 @@ impl AppState {
if let Some(updated_at_unix_secs) = updated_at_unix_secs {
wallet.updated_at_unix_secs = updated_at_unix_secs;
}
return Ok(Some(wallet.clone()));
let wallet = wallet.clone();
drop(guard);
self.invalidate_auth_context_cache();
return Ok(Some(wallet));
}
self.data
let wallet = self
.data
.update_auth_user_wallet_snapshot(
user_id,
balance,
@@ -399,7 +431,11 @@ impl AppState {
updated_at_unix_secs,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if wallet.is_some() {
self.invalidate_auth_context_cache();
}
Ok(wallet)
}
#[allow(clippy::too_many_arguments)]
@@ -438,10 +474,14 @@ impl AppState {
if let Some(updated_at_unix_secs) = updated_at_unix_secs {
wallet.updated_at_unix_secs = updated_at_unix_secs;
}
return Ok(Some(wallet.clone()));
let wallet = wallet.clone();
drop(guard);
self.invalidate_auth_context_cache();
return Ok(Some(wallet));
}
self.data
let wallet = self
.data
.update_auth_api_key_wallet_snapshot(
api_key_id,
balance,
@@ -456,6 +496,10 @@ impl AppState {
updated_at_unix_secs,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
.map_err(|err| GatewayError::Internal(err.to_string()))?;
if wallet.is_some() {
self.invalidate_auth_context_cache();
}
Ok(wallet)
}
}

67
apps/aether-gateway/src/tests/architecture/ai_serving.rs
View File

@@ -1177,12 +1177,9 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
let ranking_call = candidate_resolution
.find("rank_eligible_local_execution_candidates(")
.expect("candidate_resolution.rs should call core-backed local candidate ranking");
let pool_scheduler_call = candidate_resolution
.find("apply_local_execution_pool_scheduler(")
.expect("candidate_resolution.rs should call pool scheduler after ranking");
assert!(
ranking_call < pool_scheduler_call,
"candidate_resolution.rs should rank eligible candidates before applying pool-internal account scheduling"
!candidate_resolution.contains("apply_local_execution_pool_scheduler("),
"candidate_resolution.rs should leave pool-internal key scheduling to dispatch cursors"
);
for pattern in [
"pub(crate) async fn resolve_and_rank_local_execution_candidates(",
@@ -1270,18 +1267,45 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
);
}
let pool_scheduler =
let planner_pool_scheduler =
read_workspace_file("apps/aether-gateway/src/ai_serving/planner/pool_scheduler.rs");
for pattern in [
"pub(crate) use crate::dispatch::pool_scheduler::apply_local_execution_pool_scheduler;",
"pub(crate) use crate::dispatch::pool_scheduler::PoolKeyCursor;",
] {
assert!(
planner_pool_scheduler.contains(pattern),
"planner/pool_scheduler.rs should only re-export dispatch pool scheduler compatibility item {pattern}"
);
}
for forbidden in [
"pub(crate) async fn apply_local_execution_pool_scheduler(",
"run_ai_pool_scheduler(",
"fn ai_pool_candidate_facts(",
"fn ai_pool_scheduling_config(",
"fn ai_pool_runtime_state(",
] {
assert!(
!planner_pool_scheduler.contains(forbidden),
"planner/pool_scheduler.rs should not own dispatch implementation {forbidden}"
);
}
let pool_scheduler = read_workspace_file("apps/aether-gateway/src/dispatch/pool_scheduler.rs");
for pattern in [
"pub(crate) async fn apply_local_execution_pool_scheduler(",
"pub(crate) struct PoolKeyCursor",
"run_ai_pool_scheduler(",
"fn ai_pool_candidate_facts(",
"fn ai_pool_scheduling_config(",
"fn ai_pool_runtime_state(",
"DEFAULT_POOL_WINDOW_SIZE",
"DEFAULT_POOL_PAGE_SIZE",
"DEFAULT_POOL_MAX_SCAN",
] {
assert!(
pool_scheduler.contains(pattern),
"planner/pool_scheduler.rs should adapt gateway pool runtime data through serving pool scheduler helper {pattern}"
"dispatch/pool_scheduler.rs should adapt gateway pool runtime data through serving pool scheduler helper {pattern}"
);
}
for forbidden in [
@@ -1296,7 +1320,34 @@ fn ai_serving_planner_separates_local_candidate_resolution_from_ranking() {
] {
assert!(
!pool_scheduler.contains(forbidden),
"planner/pool_scheduler.rs should not own global candidate ranking or pool scheduling policy helper {forbidden}"
"dispatch/pool_scheduler.rs should not own global candidate ranking or pool scheduling policy helper {forbidden}"
);
}
let dispatch_refs = read_workspace_file("apps/aether-gateway/src/dispatch/refs.rs");
for pattern in [
"DispatchCandidateRef::SingleKey",
"DispatchCandidateRef::PoolRef",
"pub(crate) fn key_ref_for_candidate",
"pub(crate) fn pool_ref_for_candidate",
] {
assert!(
dispatch_refs.contains(pattern),
"dispatch/refs.rs should expose logical dispatch refs through {pattern}"
);
}
let dispatch_core = read_workspace_file("crates/aether-dispatch-core/src/lib.rs");
for pattern in [
"DispatchCandidateRef",
"DispatchSequence",
"PoolDispatchPort",
"PoolWindowConfig",
"DispatchEffect",
] {
assert!(
dispatch_core.contains(pattern),
"aether-dispatch-core should export pure dispatch primitive {pattern}"
);
}

12
apps/aether-gateway/src/tests/proxy.rs
View File

@@ -540,8 +540,8 @@ async fn gateway_forwards_public_request_to_remote_tunnel_owner_before_fallback_
state = state
.with_data_state_for_tests(data_state)
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080"));
state.scheduler_affinity_cache.insert(
"scheduler_affinity:api-key-affinity-1:openai:chat:gpt-4.1".to_string(),
state.remember_scheduler_affinity_target(
"scheduler_affinity:api-key-affinity-1:openai:chat:gpt-4.1",
crate::cache::SchedulerAffinityTarget {
provider_id: "provider-owner".to_string(),
endpoint_id: "endpoint-owner".to_string(),
@@ -761,8 +761,8 @@ async fn gateway_aggregates_sync_sse_from_remote_tunnel_owner_before_returning_t
state = state
.with_data_state_for_tests(data_state)
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080"));
state.scheduler_affinity_cache.insert(
"scheduler_affinity:api-key-affinity-cli-1:openai:responses:gpt-5.4".to_string(),
state.remember_scheduler_affinity_target(
"scheduler_affinity:api-key-affinity-cli-1:openai:responses:gpt-5.4",
crate::cache::SchedulerAffinityTarget {
provider_id: "provider-cli-owner".to_string(),
endpoint_id: "endpoint-cli-owner".to_string(),
@@ -1001,8 +1001,8 @@ async fn gateway_streamifies_sync_json_from_remote_tunnel_owner_before_returning
state = state
.with_data_state_for_tests(data_state)
.with_tunnel_identity_for_tests("gateway-a", Some("http://gateway-a:8080"));
state.scheduler_affinity_cache.insert(
"scheduler_affinity:api-key-affinity-cli-1:openai:responses:gpt-5.4".to_string(),
state.remember_scheduler_affinity_target(
"scheduler_affinity:api-key-affinity-cli-1:openai:responses:gpt-5.4",
crate::cache::SchedulerAffinityTarget {
provider_id: "provider-cli-owner".to_string(),
endpoint_id: "endpoint-cli-owner".to_string(),