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

Improve gateway transport and usage runtime

Browse Source
This commit is contained in:
elky
2026-06-25 22:36:27 +08:00
parent d336d1a7fa
commit 6f00e9fc67
112 changed files with 12456 additions and 1387 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
+4
View File
@@ -740,6 +740,7 @@ where
request_auth_channel,
use_api_format_alias_match,
key_mode,
sticky_session_token.is_none(),
Some(trace_id),
)
.await;
@@ -2188,6 +2189,7 @@ mod tests {
None,
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel,
true,
Some("trace-no-session-affinity"),
)
.await;
@@ -2241,6 +2243,7 @@ mod tests {
None,
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel,
true,
Some("trace-session-affinity"),
)
.await;
@@ -2276,6 +2279,7 @@ mod tests {
None,
false,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModel,
true,
Some("trace-fixed-order"),
)
.await;
apps/aether-gateway/src/ai_serving/planner/candidate_source.rs
+213 -15
View File
@@ -12,18 +12,29 @@ use aether_scheduler_core::{
};
use async_trait::async_trait;
use std::collections::{BTreeMap, BTreeSet, VecDeque};
use std::sync::Arc;
use crate::ai_serving::planner::candidate_affinity_cache::has_explicit_session_affinity;
use crate::ai_serving::planner::candidate_resolution::SkippedLocalExecutionCandidate;
use crate::ai_serving::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::cache::{
candidate_page_cache_stale_ttl, candidate_page_cache_ttl_from_env,
record_candidate_page_cache_follower_wait, record_candidate_page_cache_hit,
record_candidate_page_cache_load, record_candidate_page_cache_miss,
record_candidate_page_cache_none, record_candidate_row_page_cache_follower_wait,
record_candidate_row_page_cache_hit, record_candidate_row_page_cache_load,
record_candidate_row_page_cache_miss, record_candidate_row_page_cache_none, CacheLoadObserver,
CandidatePageCacheKey, CandidatePageSnapshot, CandidateRowPageCacheKey,
};
use crate::clock::request_distribution_seed;
use crate::data::candidate_selection::{
read_requested_model_rows_fast_path_page, requested_model_candidate_names,
MinimalCandidateSelectionRowSource, REQUESTED_MODEL_CANDIDATE_PAGE_SIZE,
REQUESTED_MODEL_MAX_SCANNED_ROWS,
MinimalCandidateSelectionRowSource, RequestedModelCandidateRowsPage,
REQUESTED_MODEL_CANDIDATE_PAGE_SIZE, REQUESTED_MODEL_MAX_SCANNED_ROWS,
};
use crate::scheduler::candidate::SchedulerSkippedCandidate;
use crate::scheduler::config::{SchedulerOrderingConfig, SchedulerSchedulingMode};
use crate::stage_metrics::observe_gateway_stage_ms;
use crate::GatewayError;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -258,6 +269,7 @@ pub(crate) struct LocalCandidatePreselectionPageCursor<'a> {
request_auth_channel: Option<String>,
use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode,
allow_priority_page_cache: bool,
candidate_api_formats: Vec<String>,
model_directive_enabled_api_formats: BTreeSet<String>,
ordering_config: SchedulerOrderingConfig,
@@ -295,6 +307,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
request_auth_channel: Option<&str>,
use_api_format_alias_match: bool,
key_mode: LocalCandidatePreselectionKeyMode,
allow_priority_page_cache: bool,
trace_id: Option<&str>,
) -> Self {
let candidate_api_formats =
@@ -336,6 +349,7 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
request_auth_channel: request_auth_channel.map(str::to_string),
use_api_format_alias_match,
key_mode,
allow_priority_page_cache,
candidate_api_formats,
model_directive_enabled_api_formats,
ordering_config,
@@ -429,6 +443,10 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
}
}
fn should_cache_current_priority_page(&self) -> bool {
self.allow_priority_page_cache && self.should_cache_current_priority_resolved_page()
}
#[cfg(test)]
pub(crate) fn mark_priority_page_emitted_for_tests(&mut self) {
self.priority_page_emitted = true;
@@ -443,11 +461,75 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
>,
GatewayError,
> {
let page = self.next_priority_page_with_planning_gate().await?;
let page = if self.should_cache_current_priority_page() {
self.cached_next_priority_page_snapshot().await?
} else {
self.next_priority_page_with_planning_gate().await?
};
self.remember_seen_candidates_from_page(&page);
Ok(page)
}
async fn cached_next_priority_page_snapshot(
&mut self,
) -> Result<
AiCandidatePreselectionOutcome<
SchedulerMinimalCandidateSelectionCandidate,
SkippedLocalExecutionCandidate,
>,
GatewayError,
> {
let key = CandidatePageCacheKey::new(
&self.requested_model,
&self.client_api_format,
self.require_streaming,
&self.auth_snapshot,
self.required_capabilities.as_ref(),
self.routing_policy.as_ref(),
self.request_auth_channel.as_deref(),
self.state.app().scheduler_affinity_epoch(),
self.key_mode.cache_key_name(),
self.use_api_format_alias_match,
self.client_session_affinity.as_ref(),
);
let cache = self.state.app().candidate_page_cache.clone();
let ttl = candidate_page_cache_ttl_from_env();
let stale_ttl = candidate_page_cache_stale_ttl(ttl);
let cached = cache
.get_or_load_once_stale_while_refreshing(
key,
ttl,
stale_ttl,
|| async {
let page = self.next_priority_page_with_planning_gate().await?;
Ok::<_, GatewayError>(Some(Arc::new(page) as Arc<CandidatePageSnapshot>))
},
CacheLoadObserver::new()
.on_hit(record_candidate_page_cache_hit)
.on_miss(record_candidate_page_cache_miss)
.on_load(record_candidate_page_cache_load)
.on_follower_wait(record_candidate_page_cache_follower_wait),
)
.await?;
match cached {
Some(snapshot) => {
let page = snapshot.as_ref().clone();
if page.candidates.is_empty() && page.skipped_candidates.is_empty() {
record_candidate_page_cache_none();
}
Ok(page)
}
None => {
record_candidate_page_cache_none();
Ok(AiCandidatePreselectionOutcome {
candidates: Vec::new(),
skipped_candidates: Vec::new(),
})
}
}
}
fn remember_seen_candidates_from_page(
&mut self,
page: &AiCandidatePreselectionOutcome<
@@ -717,17 +799,15 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
return Ok(None);
}
let limit = REQUESTED_MODEL_CANDIDATE_PAGE_SIZE.min(remaining);
let page = read_requested_model_rows_fast_path_page(
self.state.app().data.as_ref(),
&normalized_api_format,
&self.requested_model,
requested_name,
offset,
limit,
enable_model_directives,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let page = self
.read_requested_model_rows_fast_path_page_cached(
&normalized_api_format,
requested_name,
offset,
limit,
enable_model_directives,
)
.await?;
self.scanned_rows_by_format.insert(
normalized_api_format.clone(),
scanned.saturating_add(page.scanned_rows),
@@ -765,6 +845,70 @@ impl<'a> LocalCandidatePreselectionPageCursor<'a> {
}
}
async fn read_requested_model_rows_fast_path_page_cached(
&self,
normalized_api_format: &str,
requested_name: &str,
offset: u32,
limit: u32,
enable_model_directives: bool,
) -> Result<RequestedModelCandidateRowsPage, GatewayError> {
let key = CandidateRowPageCacheKey::new(
normalized_api_format,
&self.requested_model,
requested_name,
offset,
limit,
enable_model_directives,
);
let cache = self.state.app().candidate_row_page_cache.clone();
let ttl = candidate_page_cache_ttl_from_env();
let stale_ttl = candidate_page_cache_stale_ttl(ttl);
let cached = cache
.get_or_load_once_stale_while_refreshing(
key,
ttl,
stale_ttl,
|| async {
let page = read_requested_model_rows_fast_path_page(
self.state.app().data.as_ref(),
normalized_api_format,
&self.requested_model,
requested_name,
offset,
limit,
enable_model_directives,
)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
Ok::<_, GatewayError>(Some(Arc::new(page)))
},
CacheLoadObserver::new()
.on_hit(record_candidate_row_page_cache_hit)
.on_miss(record_candidate_row_page_cache_miss)
.on_load(record_candidate_row_page_cache_load)
.on_follower_wait(record_candidate_row_page_cache_follower_wait),
)
.await?;
match cached {
Some(page) => {
if page.rows.is_empty() {
record_candidate_row_page_cache_none();
}
Ok(page.as_ref().clone())
}
None => {
record_candidate_row_page_cache_none();
Ok(RequestedModelCandidateRowsPage {
rows: Vec::new(),
scanned_rows: 0,
end_of_requested_name: true,
})
}
}
}
async fn next_fallback_page_for_api_format(
&mut self,
candidate_api_format: &str,
@@ -1022,8 +1166,15 @@ async fn acquire_candidate_planning_gate(
.app()
.frontdoor_runtime_guards
.internal_gate_queue_budget;
let gate_wait_started_at = std::time::Instant::now();
match tokio::time::timeout(budget, gate.acquire()).await {
Ok(Ok(permit)) => Ok(Some(permit)),
Ok(Ok(permit)) => {
observe_gateway_stage_ms(
"candidate_planning_gate_wait",
gate_wait_started_at.elapsed().as_millis() as u64,
);
Ok(Some(permit))
}
Ok(Err(err)) => Err(GatewayError::Internal(err.to_string())),
Err(_) => Err(GatewayError::AdmissionTimeout {
trace_id: trace_id.to_string(),
@@ -1140,6 +1291,50 @@ mod tests {
}
}
#[tokio::test]
async fn priority_page_cache_requires_fixed_order_or_explicit_affinity() {
let repository: Arc<dyn MinimalCandidateSelectionReadRepository> =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(
Vec::<StoredMinimalCandidateSelectionRow>::new(),
));
let data_state =
GatewayDataState::with_minimal_candidate_selection_reader_for_tests(repository);
let app = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state);
let auth_snapshot = unrestricted_auth_snapshot();
let mut cursor = LocalCandidatePreselectionPageCursor::new(
PlannerAppState::new(&app),
"openai:chat",
"gpt-5",
true,
None,
&auth_snapshot,
None,
None,
None,
true,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
true,
None,
)
.await;
cursor.mark_priority_page_emitted_for_tests();
cursor.ordering_config.scheduling_mode = SchedulerSchedulingMode::CacheAffinity;
assert!(!cursor.should_cache_current_priority_resolved_page());
cursor.client_session_affinity =
Some(aether_scheduler_core::ClientSessionAffinity::from_session_key("session-1"));
assert!(cursor.should_cache_current_priority_resolved_page());
cursor.ordering_config.scheduling_mode = SchedulerSchedulingMode::FixedOrder;
assert!(cursor.should_cache_current_priority_resolved_page());
cursor.ordering_config.scheduling_mode = SchedulerSchedulingMode::LoadBalance;
assert!(!cursor.should_cache_current_priority_resolved_page());
}
fn openai_responses_mapping_row() -> StoredMinimalCandidateSelectionRow {
StoredMinimalCandidateSelectionRow {
provider_id: "provider-openai-responses-mapped-1".to_string(),
@@ -1362,6 +1557,7 @@ mod tests {
None,
true,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
true,
None,
)
.await;
@@ -1421,6 +1617,7 @@ mod tests {
None,
true,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
true,
None,
)
.await;
@@ -1495,6 +1692,7 @@ mod tests {
None,
true,
LocalCandidatePreselectionKeyMode::ProviderEndpointKeyModelAndApiFormat,
true,
None,
)
.await;
apps/aether-gateway/src/ai_serving/planner/decision_input.rs
+260 -13
View File
@@ -1,4 +1,5 @@
use std::collections::BTreeMap;
use std::time::Duration;
use aether_ai_serving::{run_ai_authenticated_decision_input, AiAuthenticatedDecisionInputPort};
use aether_routing_core::{
@@ -18,11 +19,15 @@ use crate::client_session_affinity::client_session_affinity_from_request;
use crate::clock::current_unix_secs;
use crate::routing::{
apply_routing_mutation_plan, build_routing_trace_seed, resolve_gateway_routing_policy,
select_gateway_routing_group, GatewayRoutingPolicyInput, GatewayRoutingSelectionError,
GatewayRoutingSelectionInput, ROUTING_GROUP_HEADER,
resolve_gateway_static_default_routing_policy, select_gateway_routing_group,
GatewayRoutingPolicyInput, GatewayRoutingSelectionError, GatewayRoutingSelectionInput,
GatewayStaticRoutingPolicyInput, ROUTING_GROUP_HEADER,
};
use crate::stage_metrics::observe_gateway_stage_ms;
use crate::{AiExecutionDecision, AppState, GatewayError};
const ROUTING_GROUP_SELECTION_CACHE_TTL: Duration = Duration::from_secs(30);
#[derive(Debug, Clone)]
pub(crate) struct ResolvedLocalDecisionAuthInput {
pub(crate) auth_context: ExecutionRuntimeAuthContext,
@@ -221,6 +226,7 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
let explicit_group = routing_header_value_str(&parts.headers, ROUTING_GROUP_HEADER);
let selected_group = match state.routing_group_read_repository() {
Some(repository) => {
let user_groups_lookup_started_at = std::time::Instant::now();
let user_group_ids = match state
.list_user_groups_for_user(&input.auth_context.user_id)
.await
@@ -235,17 +241,50 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
Vec::new()
}
};
let selection = select_gateway_routing_group(
repository.as_ref(),
GatewayRoutingSelectionInput {
explicit_group: explicit_group.as_deref(),
user_id: Some(input.auth_context.user_id.as_str()),
api_key_id: Some(input.auth_context.api_key_id.as_str()),
user_group_ids: &user_group_ids,
},
)
.await
.map_err(routing_selection_error)?;
observe_gateway_stage_ms(
"routing_user_groups_lookup",
user_groups_lookup_started_at.elapsed().as_millis() as u64,
);
let selection_cache_key = routing_group_selection_cache_key(
explicit_group.as_deref(),
Some(input.auth_context.user_id.as_str()),
Some(input.auth_context.api_key_id.as_str()),
&user_group_ids,
);
let user_id = input.auth_context.user_id.clone();
let api_key_id = input.auth_context.api_key_id.clone();
let group_selection_started_at = std::time::Instant::now();
let selection = state
.routing_group_selection_cache
.get_or_load_once(
selection_cache_key,
ROUTING_GROUP_SELECTION_CACHE_TTL,
|| async move {
let selection_load_started_at = std::time::Instant::now();
let selection = select_gateway_routing_group(
repository.as_ref(),
GatewayRoutingSelectionInput {
explicit_group: explicit_group.as_deref(),
user_id: Some(user_id.as_str()),
api_key_id: Some(api_key_id.as_str()),
user_group_ids: &user_group_ids,
},
)
.await
.map_err(routing_selection_error)?;
observe_gateway_stage_ms(
"routing_group_selection_load",
selection_load_started_at.elapsed().as_millis() as u64,
);
Ok::<_, GatewayError>(Some(selection))
},
)
.await?
.unwrap_or_default();
observe_gateway_stage_ms(
"routing_group_selection",
group_selection_started_at.elapsed().as_millis() as u64,
);
selection.group.map(|group| {
(
Some(group.id),
@@ -279,7 +318,21 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
return Ok(());
};
if try_attach_static_default_routing_policy_to_input(
input,
parts,
body_json,
client_api_format,
group_id.as_deref(),
group_version,
&group_config_json,
selection_source.as_str(),
)? {
return Ok(());
}
let headers_json = headers_to_routing_value(&parts.headers);
let policy_resolve_started_at = std::time::Instant::now();
let policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput {
group_id: group_id.as_deref(),
group_version,
@@ -294,13 +347,22 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
body: body_json,
phase: RoutingRulePhase::ClientRequest,
})?;
observe_gateway_stage_ms(
"routing_policy_resolve",
policy_resolve_started_at.elapsed().as_millis() as u64,
);
let mut effective_body_json = body_json.clone();
let mut effective_headers = parts.headers.clone();
let mutation_apply_started_at = std::time::Instant::now();
apply_routing_mutation_plan(
&mut effective_body_json,
&mut effective_headers,
&policy.mutation_plan,
)?;
observe_gateway_stage_ms(
"routing_mutation_apply",
mutation_apply_started_at.elapsed().as_millis() as u64,
);
let mut requested_model_changed = false;
if let Some(mut mutated_model) = extract_standard_requested_model(&effective_body_json) {
@@ -324,6 +386,7 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
let effective_headers_json = headers_to_routing_value(&effective_headers);
input.client_session_affinity =
client_session_affinity_from_request(&effective_headers, Some(&effective_body_json));
let final_policy_resolve_started_at = std::time::Instant::now();
let mut final_policy = resolve_gateway_routing_policy(GatewayRoutingPolicyInput {
group_id: group_id.as_deref(),
group_version,
@@ -338,6 +401,10 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
body: &effective_body_json,
phase: RoutingRulePhase::ClientRequest,
})?;
observe_gateway_stage_ms(
"routing_policy_resolve",
final_policy_resolve_started_at.elapsed().as_millis() as u64,
);
final_policy.mutation_plan = policy.mutation_plan.clone();
input.routing_trace_seed = Some(build_routing_trace_seed(&final_policy, client_api_format));
input.routing_policy = Some(final_policy);
@@ -353,6 +420,46 @@ pub(crate) async fn attach_routing_policy_to_local_requested_model_input(
Ok(())
}
fn try_attach_static_default_routing_policy_to_input(
input: &mut LocalRequestedModelDecisionInput,
parts: &http::request::Parts,
body_json: &Value,
client_api_format: &str,
group_id: Option<&str>,
group_version: Option<i64>,
group_config_json: &Value,
selection_source: &str,
) -> Result<bool, GatewayError> {
let static_policy_resolve_started_at = std::time::Instant::now();
let Some(policy) =
resolve_gateway_static_default_routing_policy(GatewayStaticRoutingPolicyInput {
group_id,
group_version,
group_config_json,
selection_source,
requested_model: input.requested_model.as_str(),
resolved_model: input.requested_model.as_str(),
})?
else {
observe_gateway_stage_ms(
"routing_static_policy_resolve",
static_policy_resolve_started_at.elapsed().as_millis() as u64,
);
return Ok(false);
};
observe_gateway_stage_ms(
"routing_static_policy_resolve",
static_policy_resolve_started_at.elapsed().as_millis() as u64,
);
input.client_session_affinity =
client_session_affinity_from_request(&parts.headers, Some(body_json));
input.routing_trace_seed = Some(build_routing_trace_seed(&policy, client_api_format));
input.routing_policy = Some(policy);
input.routing_context = None;
Ok(true)
}
pub(crate) fn build_local_authenticated_decision_input(
resolved_input: ResolvedLocalDecisionAuthInput,
) -> LocalAuthenticatedDecisionInput {
@@ -410,6 +517,33 @@ fn routing_header_value_str(headers: &http::HeaderMap, key: &str) -> Option<Stri
.map(ToOwned::to_owned)
}
fn routing_group_selection_cache_key(
explicit_group: Option<&str>,
user_id: Option<&str>,
api_key_id: Option<&str>,
user_group_ids: &[String],
) -> String {
let groups = user_group_ids
.iter()
.map(|value| escape_cache_key_part(value))
.collect::<Vec<_>>()
.join(",");
format!(
"v1|explicit={}|user={}|api_key={}|groups={}",
escape_cache_key_part(explicit_group.unwrap_or_default()),
escape_cache_key_part(user_id.unwrap_or_default()),
escape_cache_key_part(api_key_id.unwrap_or_default()),
groups
)
}
fn escape_cache_key_part(value: &str) -> String {
value
.replace('%', "%25")
.replace('|', "%7C")
.replace(',', "%2C")
}
fn btree_headers_to_header_map(
headers: &BTreeMap<String, String>,
) -> Result<HeaderMap, GatewayError> {
@@ -754,6 +888,119 @@ mod tests {
.group_config_json = config;
}
#[test]
fn static_default_routing_policy_attaches_without_request_context() {
let request = http::Request::builder()
.header("content-type", "application/json")
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let mut input = LocalRequestedModelDecisionInput {
auth_context: sample_auth_context(),
requested_model: "mock-model".to_string(),
auth_snapshot: sample_auth_snapshot(),
required_capabilities: None,
request_auth_channel: None,
client_session_affinity: None,
routing_policy: None,
routing_trace_seed: None,
routing_context: Some(LocalRoutingRequestContext {
group_id: Some("stale".to_string()),
group_version: Some(1),
group_config_json: json!({}),
selection_source: "stale".to_string(),
client_api_format: "openai:chat".to_string(),
effective_body_json: json!({}),
effective_headers: HeaderMap::new(),
}),
};
let group_config_json = json!({
"default_policy": {
"priority_mode": "global_key",
"scheduling_mode": "load_balance",
"keep_priority_on_conversion": true
},
"allowed_models": [],
"model_policies": [],
"rules": []
});
let attached = try_attach_static_default_routing_policy_to_input(
&mut input,
&parts,
&json!({"model": "mock-model"}),
"openai:chat",
Some("group-1"),
Some(4),
&group_config_json,
"system_default",
)
.expect("static routing should attach");
assert!(attached);
assert!(input.routing_context.is_none());
let policy = input.routing_policy.as_ref().expect("policy should be set");
assert_eq!(policy.group_id.as_deref(), Some("group-1"));
assert_eq!(policy.group_version, Some(4));
assert_eq!(
policy.priority_mode,
aether_routing_core::RoutingSetPriorityMode::GlobalKey
);
assert_eq!(
policy.scheduling_mode,
aether_routing_core::RoutingSchedulingMode::LoadBalance
);
assert!(policy.keep_priority_on_conversion);
assert!(policy.mutation_plan.is_empty());
assert!(input.routing_trace_seed.is_some());
}
#[test]
fn dynamic_routing_policy_does_not_attach_static_fast_path() {
let request = http::Request::builder()
.body(())
.expect("request should build");
let (parts, _) = request.into_parts();
let mut input = LocalRequestedModelDecisionInput {
auth_context: sample_auth_context(),
requested_model: "mock-model".to_string(),
auth_snapshot: sample_auth_snapshot(),
required_capabilities: None,
request_auth_channel: None,
client_session_affinity: None,
routing_policy: None,
routing_trace_seed: None,
routing_context: None,
};
let group_config_json = json!({
"rules": [{
"id": "rule-1",
"conditions": {},
"actions": [{
"type": "restrict_providers",
"provider_ids": ["provider-1"]
}]
}]
});
let attached = try_attach_static_default_routing_policy_to_input(
&mut input,
&parts,
&json!({"model": "mock-model"}),
"openai:chat",
Some("group-1"),
Some(4),
&group_config_json,
"system_default",
)
.expect("dynamic config should not fail static detection");
assert!(!attached);
assert!(input.routing_policy.is_none());
assert!(input.routing_trace_seed.is_none());
assert!(input.routing_context.is_none());
}
#[test]
fn provider_request_routing_policy_mutates_decision_body_headers_and_report_context() {
let input = sample_decision_input();
apps/aether-gateway/src/ai_serving/planner/redaction.rs
+2
View File
@@ -77,6 +77,7 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>(
};
let request_cache_key = request_redaction_cache_key(format, body_json);
if let Some(cached) = slot.cached_request_redaction(&request_cache_key) {
crate::stage_metrics::record_chat_pii_redaction_request_cache_hit();
observe_gateway_stage_ms("chat_pii_redaction_request_cache_hit", 0);
return Ok(provider_redaction_from_cached(
slot,
@@ -85,6 +86,7 @@ pub(crate) async fn resolve_provider_chat_pii_redaction<'a>(
cached,
));
}
crate::stage_metrics::record_chat_pii_redaction_request_cache_miss();
let runtime_config_started_at = Instant::now();
let runtime_config = read_chat_pii_redaction_runtime_config(state)
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision.rs
+1
View File
@@ -6,6 +6,7 @@ mod request;
mod support;
pub(super) use self::payload::maybe_build_local_openai_chat_decision_payload_for_candidate;
pub(super) use self::request::LocalOpenAiChatRequestPreparation;
pub(super) use self::support::{
build_lazy_local_openai_chat_candidate_attempt_source,
build_local_openai_chat_candidate_attempt_source,
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/payload.rs
+5 -1
View File
@@ -18,7 +18,9 @@ use crate::{
AiExecutionDecision, AppState, GatewayError,
};
use super::request::resolve_local_openai_chat_candidate_payload_parts;
use super::request::{
resolve_local_openai_chat_candidate_payload_parts, LocalOpenAiChatRequestPreparation,
};
use super::support::{LocalOpenAiChatCandidateAttempt, LocalOpenAiChatDecisionInput};
#[allow(clippy::too_many_arguments)]
@@ -28,6 +30,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
preparation: Option<&mut LocalOpenAiChatRequestPreparation>,
attempt: LocalOpenAiChatCandidateAttempt,
decision_kind: &str,
report_kind: &str,
@@ -48,6 +51,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
trace_id,
body_json,
input,
preparation,
&eligible,
candidate_index,
&candidate_id,
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/decision/request.rs
+42 -2
View File
@@ -82,6 +82,39 @@ pub(crate) struct LocalOpenAiChatCandidatePayloadParts {
pub(super) image_request_summary: Option<Value>,
}
#[derive(Default)]
pub(crate) struct LocalOpenAiChatRequestPreparation {
model_directives_enabled: BTreeMap<(String, String), bool>,
}
impl LocalOpenAiChatRequestPreparation {
async fn model_directives_enabled(
&mut self,
state: &AppState,
provider_api_format: &str,
requested_model: &str,
) -> bool {
let key = (
provider_api_format.trim().to_ascii_lowercase(),
requested_model.trim().to_string(),
);
if let Some(enabled) = self.model_directives_enabled.get(&key) {
crate::stage_metrics::record_openai_chat_model_directive_cache_hit();
return *enabled;
}
crate::stage_metrics::record_openai_chat_model_directive_cache_miss();
let enabled =
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
state,
provider_api_format,
Some(requested_model),
)
.await;
self.model_directives_enabled.insert(key, enabled);
enabled
}
}
fn is_grok_text_provider_api_format(provider_api_format: &str) -> bool {
matches!(
crate::ai_serving::normalize_api_format_alias(provider_api_format).as_str(),
@@ -96,6 +129,7 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
trace_id: &str,
body_json: &serde_json::Value,
input: &LocalOpenAiChatDecisionInput,
mut preparation: Option<&mut LocalOpenAiChatRequestPreparation>,
eligible: &EligibleLocalExecutionCandidate,
candidate_index: u32,
candidate_id: &str,
@@ -112,13 +146,18 @@ pub(crate) async fn resolve_local_openai_chat_candidate_payload_parts(
let force_body_stream_field =
endpoint_config_forces_body_stream_field(transport.endpoint.config.as_ref());
let model_directives_started_at = std::time::Instant::now();
let enable_model_directives =
let enable_model_directives = if let Some(preparation) = preparation {
preparation
.model_directives_enabled(state, provider_api_format, &input.requested_model)
.await
} else {
crate::system_features::reasoning_model_directive_enabled_for_api_format_and_model(
state,
provider_api_format,
Some(&input.requested_model),
)
.await;
.await
};
observe_gateway_stage_ms(
"openai_chat_payload_model_directives",
model_directives_started_at.elapsed().as_millis() as u64,
@@ -1907,6 +1946,7 @@ mod tests {
"trace-openai-chat-gemini-cli",
&body_json,
&sample_input(),
None,
&sample_gemini_cli_eligible(),
0,
"candidate-0",
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/mod.rs
+3
View File
@@ -13,6 +13,7 @@ use self::decision::{
build_lazy_local_openai_chat_candidate_attempt_source,
maybe_build_local_openai_chat_decision_payload_for_candidate, LocalOpenAiChatCandidateAttempt,
LocalOpenAiChatCandidateAttemptSource, LocalOpenAiChatDecisionInput,
LocalOpenAiChatRequestPreparation,
};
use self::plans::{
build_local_openai_chat_stream_attempt_source, build_local_openai_chat_stream_plan_and_reports,
@@ -159,6 +160,7 @@ pub(crate) async fn maybe_build_sync_local_decision_payload(
trace_id,
body_json,
&input,
None,
attempt,
OPENAI_CHAT_SYNC_PLAN_KIND,
"openai_chat_sync_success",
@@ -211,6 +213,7 @@ pub(crate) async fn maybe_build_stream_local_decision_payload(
trace_id,
body_json,
&input,
None,
attempt,
OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success",
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/resolve.rs
+16
View File
@@ -9,6 +9,7 @@ use crate::ai_serving::planner::decision_input::{
};
use crate::ai_serving::resolve_local_decision_execution_runtime_auth_context;
use crate::client_session_affinity::client_session_affinity_from_parts;
use crate::stage_metrics::observe_gateway_stage_ms;
use crate::{AppState, GatewayError};
pub(crate) async fn resolve_local_openai_chat_decision_input(
@@ -59,6 +60,7 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
return Ok(None);
};
let auth_started_at = std::time::Instant::now();
let resolved_input = match resolve_local_authenticated_decision_input(
state,
auth_context.clone(),
@@ -106,10 +108,20 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
return Err(err);
}
};
observe_gateway_stage_ms(
"openai_chat_decision_input_auth",
auth_started_at.elapsed().as_millis() as u64,
);
let mut input = build_local_requested_model_decision_input(resolved_input, requested_model);
input.request_auth_channel = decision.request_auth_channel.clone();
let affinity_started_at = std::time::Instant::now();
input.client_session_affinity = client_session_affinity_from_parts(parts, Some(body_json));
observe_gateway_stage_ms(
"openai_chat_decision_input_affinity",
affinity_started_at.elapsed().as_millis() as u64,
);
let routing_started_at = std::time::Instant::now();
if let Err(err) = attach_routing_policy_to_local_requested_model_input(
state,
parts,
@@ -126,5 +138,9 @@ pub(crate) async fn resolve_local_openai_chat_decision_input(
);
return Err(err);
}
observe_gateway_stage_ms(
"openai_chat_decision_input_routing",
routing_started_at.elapsed().as_millis() as u64,
);
Ok(Some(input))
}
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/stream.rs
+339 -27
View File
@@ -1,11 +1,12 @@
use async_trait::async_trait;
use std::collections::VecDeque;
use tracing::warn;
use super::super::{
build_lazy_local_openai_chat_candidate_attempt_source,
maybe_build_local_openai_chat_decision_payload_for_candidate, AppState, GatewayControlDecision,
GatewayError, LocalOpenAiChatCandidateAttempt, LocalOpenAiChatCandidateAttemptSource,
LocalOpenAiChatDecisionInput,
LocalOpenAiChatDecisionInput, LocalOpenAiChatRequestPreparation,
};
use super::diagnostic::{
set_local_openai_chat_candidate_evaluation_diagnostic, set_local_openai_chat_miss_diagnostic,
@@ -18,7 +19,23 @@ use crate::ai_serving::planner::plan_builders::{
build_openai_chat_stream_plan_from_decision, AiStreamAttempt,
};
use crate::ai_serving::planner::runtime_miss::apply_local_runtime_candidate_terminal_reason;
use crate::stage_metrics::observe_gateway_stage_ms;
use crate::ai_serving::planner::standard::build_local_openai_chat_upstream_url;
use crate::ai_serving::transport::{
is_windsurf_provider_transport, local_openai_chat_transport_unsupported_reason,
};
use crate::clock::request_distribution_seed;
use crate::stage_metrics::{
observe_gateway_stage_ms, record_openai_chat_stream_payload_build_prefetch_avoided,
record_openai_chat_stream_payload_build_selected,
record_openai_chat_stream_raw_candidates_scanned,
record_openai_chat_stream_target_select_selected_rank,
};
use crate::upstream_admission::upstream_target_key_from_url;
const OPENAI_CHAT_STREAM_TARGET_SELECT_WINDOW_ENV: &str =
"AETHER_GATEWAY_OPENAI_CHAT_STREAM_TARGET_SELECT_WINDOW";
const DEFAULT_OPENAI_CHAT_STREAM_TARGET_SELECT_WINDOW: usize = 2;
const MAX_OPENAI_CHAT_STREAM_TARGET_SELECT_WINDOW: usize = 8;
pub(crate) struct LocalOpenAiChatStreamAttemptSource<'a> {
state: &'a AppState,
@@ -27,6 +44,8 @@ pub(crate) struct LocalOpenAiChatStreamAttemptSource<'a> {
body_json: serde_json::Value,
input: LocalOpenAiChatDecisionInput,
candidates: LocalOpenAiChatCandidateAttemptSource<'a>,
prefetched_attempts: VecDeque<LocalOpenAiChatCandidateAttempt>,
request_preparation: LocalOpenAiChatRequestPreparation,
}
pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
@@ -41,6 +60,7 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
return Ok(None);
}
let attempt_source_started_at = std::time::Instant::now();
let Some(input) = resolve_local_openai_chat_decision_input(
state, parts, trace_id, decision, body_json, plan_kind, true,
)
@@ -77,6 +97,10 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
Some(input.requested_model.as_str()),
candidate_count,
);
observe_gateway_stage_ms(
"openai_chat_attempt_source_build",
attempt_source_started_at.elapsed().as_millis() as u64,
);
Ok(Some((
LocalOpenAiChatStreamAttemptSource {
@@ -86,6 +110,8 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
body_json: effective_body_json,
input,
candidates,
prefetched_attempts: VecDeque::new(),
request_preparation: LocalOpenAiChatRequestPreparation::default(),
},
candidate_count,
)))
@@ -94,21 +120,42 @@ pub(crate) async fn build_local_openai_chat_stream_attempt_source<'a>(
#[async_trait]
impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiChatStreamAttemptSource<'_> {
async fn next_execution_attempt(&mut self) -> Result<Option<AiStreamAttempt>, GatewayError> {
let select_started_at = std::time::Instant::now();
let selected = self.next_execution_attempt_with_target_select().await?;
observe_gateway_stage_ms(
"openai_chat_stream_target_select",
select_started_at.elapsed().as_millis() as u64,
);
Ok(selected)
}
async fn drain_execution_attempts(&mut self) -> Result<Vec<AiStreamAttempt>, GatewayError> {
let mut drained = Vec::new();
for attempt in self.candidates.drain_static_attempts() {
if let Some(attempt) = self.build_stream_attempt(attempt).await? {
drained.push(attempt);
}
}
Ok(drained)
}
}
impl LocalOpenAiChatStreamAttemptSource<'_> {
async fn next_execution_attempt_with_target_select(
&mut self,
) -> Result<Option<AiStreamAttempt>, GatewayError> {
loop {
let source_started_at = std::time::Instant::now();
let Some(attempt) = self.candidates.next_attempt().await? else {
observe_gateway_stage_ms(
"stream_candidate_source_next",
source_started_at.elapsed().as_millis() as u64,
let Some(attempt) = self.next_raw_attempt_with_target_select().await? else {
apply_local_runtime_candidate_terminal_reason(
self.state,
self.trace_id,
"no_local_stream_plans",
);
break;
return Ok(None);
};
observe_gateway_stage_ms(
"stream_candidate_source_next",
source_started_at.elapsed().as_millis() as u64,
);
let plan_started_at = std::time::Instant::now();
record_openai_chat_stream_payload_build_selected();
match self.build_stream_attempt(attempt).await? {
Some(attempt) => {
observe_gateway_stage_ms(
@@ -126,28 +173,122 @@ impl LocalExecutionAttemptSource<AiStreamAttempt> for LocalOpenAiChatStreamAttem
}
}
}
apply_local_runtime_candidate_terminal_reason(
self.state,
self.trace_id,
"no_local_stream_plans",
);
Ok(None)
}
async fn drain_execution_attempts(&mut self) -> Result<Vec<AiStreamAttempt>, GatewayError> {
let mut drained = Vec::new();
for attempt in self.candidates.drain_static_attempts() {
if let Some(attempt) = self.build_stream_attempt(attempt).await? {
drained.push(attempt);
async fn next_raw_attempt_with_target_select(
&mut self,
) -> Result<Option<LocalOpenAiChatCandidateAttempt>, GatewayError> {
let select_window = openai_chat_stream_target_select_window();
if select_window <= 1 {
return self.next_raw_attempt_linear().await;
}
let mut attempts = Vec::with_capacity(select_window);
for _ in 0..select_window {
match self.next_raw_attempt_linear().await? {
Some(attempt) => attempts.push(attempt),
None => break,
}
}
Ok(drained)
if attempts.is_empty() {
return Ok(None);
}
record_openai_chat_stream_raw_candidates_scanned(attempts.len());
let seed = request_distribution_seed();
let target_keys = attempts
.iter()
.map(|attempt| self.lightweight_target_key_for_attempt(attempt))
.collect::<Vec<_>>();
for target_key in target_keys.iter().flatten() {
self.state
.upstream_target_admission
.record_raw_seen_for_target_key(target_key);
}
let selected_index = if target_keys.iter().all(Option::is_some) {
let choices = attempts
.iter()
.zip(target_keys.iter())
.map(|(attempt, target_key)| {
let target_key = target_key.as_deref().unwrap_or("-");
let snapshot = self
.state
.upstream_target_admission
.snapshot_for_target_key(target_key);
TargetSelectChoice {
target_key,
identity: target_select_candidate_identity(attempt),
in_flight: snapshot
.as_ref()
.map(|snapshot| snapshot.in_flight)
.unwrap_or(0),
selection_pressure_total: snapshot
.as_ref()
.map(|snapshot| snapshot.selection_pressure_total)
.unwrap_or(0),
}
})
.collect::<Vec<_>>();
select_target_index(seed, &choices)
} else {
0
};
record_openai_chat_stream_target_select_selected_rank(selected_index);
record_openai_chat_stream_payload_build_prefetch_avoided(attempts.len().saturating_sub(1));
if let Some(Some(target_key)) = target_keys.get(selected_index) {
self.state
.upstream_target_admission
.record_preselect_for_target_key(target_key);
}
let selected = attempts.remove(selected_index);
self.prefetched_attempts.extend(attempts);
Ok(Some(selected))
}
}
impl LocalOpenAiChatStreamAttemptSource<'_> {
async fn build_stream_attempt(
async fn next_raw_attempt_linear(
&mut self,
) -> Result<Option<LocalOpenAiChatCandidateAttempt>, GatewayError> {
if let Some(attempt) = self.prefetched_attempts.pop_front() {
return Ok(Some(attempt));
}
let source_started_at = std::time::Instant::now();
let attempt = self.candidates.next_attempt().await?;
observe_gateway_stage_ms(
"stream_candidate_source_next",
source_started_at.elapsed().as_millis() as u64,
);
Ok(attempt)
}
fn lightweight_target_key_for_attempt(
&self,
attempt: &LocalOpenAiChatCandidateAttempt,
) -> Option<String> {
let provider_api_format = attempt.eligible.provider_api_format.trim();
if !provider_api_format.eq_ignore_ascii_case("openai:chat") {
return None;
}
let transport = &attempt.eligible.transport;
if transport
.provider
.provider_type
.trim()
.eq_ignore_ascii_case("grok")
|| is_windsurf_provider_transport(transport)
|| local_openai_chat_transport_unsupported_reason(transport).is_some()
{
return None;
}
if transport.provider.proxy.is_some()
|| transport.endpoint.proxy.is_some()
|| transport.key.proxy.is_some()
{
return None;
}
let upstream_url = build_local_openai_chat_upstream_url(self.parts, transport)?;
upstream_target_key_from_url(upstream_url.as_str(), None)
}
async fn build_stream_attempt(
&mut self,
attempt: LocalOpenAiChatCandidateAttempt,
) -> Result<Option<AiStreamAttempt>, GatewayError> {
let upstream_is_stream = openai_chat_upstream_is_stream_for_candidate(
@@ -161,6 +302,7 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
self.trace_id,
&self.body_json,
&self.input,
Some(&mut self.request_preparation),
attempt,
OPENAI_CHAT_STREAM_PLAN_KIND,
"openai_chat_stream_success",
@@ -185,6 +327,97 @@ impl LocalOpenAiChatStreamAttemptSource<'_> {
}
}
fn openai_chat_stream_target_select_window() -> usize {
std::env::var(OPENAI_CHAT_STREAM_TARGET_SELECT_WINDOW_ENV)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.filter(|value| *value > 0)
.unwrap_or(DEFAULT_OPENAI_CHAT_STREAM_TARGET_SELECT_WINDOW)
.clamp(1, MAX_OPENAI_CHAT_STREAM_TARGET_SELECT_WINDOW)
}
#[derive(Clone, Copy)]
struct TargetSelectCandidateIdentity<'a> {
provider_id: &'a str,
endpoint_id: &'a str,
key_id: &'a str,
candidate_id: &'a str,
}
#[derive(Clone, Copy)]
struct TargetSelectChoice<'a> {
target_key: &'a str,
identity: TargetSelectCandidateIdentity<'a>,
in_flight: usize,
selection_pressure_total: u64,
}
fn select_target_index(seed: u64, choices: &[TargetSelectChoice<'_>]) -> usize {
choices
.iter()
.enumerate()
.min_by_key(|(index, choice)| {
target_select_score(
seed,
choice.target_key,
&choice.identity,
*index,
choice.in_flight,
choice.selection_pressure_total,
)
})
.map(|(index, _)| index)
.unwrap_or(0)
}
fn target_select_candidate_identity(
attempt: &LocalOpenAiChatCandidateAttempt,
) -> TargetSelectCandidateIdentity<'_> {
TargetSelectCandidateIdentity {
provider_id: &attempt.eligible.candidate.provider_id,
endpoint_id: &attempt.eligible.candidate.endpoint_id,
key_id: &attempt.eligible.candidate.key_id,
candidate_id: &attempt.candidate_id,
}
}
fn target_select_tie_break(
seed: u64,
target_key: &str,
identity: &TargetSelectCandidateIdentity<'_>,
index: usize,
) -> u64 {
let mut hash = seed ^ ((index as u64).wrapping_mul(0x9E37_79B9_7F4A_7C15));
hash = hash_string(hash, target_key);
hash = hash_string(hash, identity.provider_id);
hash = hash_string(hash, identity.endpoint_id);
hash = hash_string(hash, identity.key_id);
hash_string(hash, identity.candidate_id)
}
fn hash_string(mut hash: u64, value: &str) -> u64 {
for byte in value.as_bytes() {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(0x100_0000_01B3);
}
hash
}
fn target_select_score(
seed: u64,
target_key: &str,
identity: &TargetSelectCandidateIdentity<'_>,
index: usize,
in_flight: usize,
selected_total: u64,
) -> (usize, u64, u64) {
(
in_flight,
selected_total,
target_select_tie_break(seed, target_key, identity, index),
)
}
pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
state: &AppState,
parts: &http::request::Parts,
@@ -234,3 +467,82 @@ pub(crate) async fn build_local_openai_chat_stream_plan_and_reports(
Ok(plans)
}
#[cfg(test)]
mod tests {
use super::*;
fn identity<'a>(
endpoint_id: &'a str,
candidate_id: &'a str,
) -> TargetSelectCandidateIdentity<'a> {
TargetSelectCandidateIdentity {
provider_id: "provider",
endpoint_id,
key_id: "key",
candidate_id,
}
}
#[test]
fn target_select_score_prefers_lower_in_flight() {
let busy = identity("endpoint-a", "candidate-a");
let idle = identity("endpoint-b", "candidate-b");
assert!(
target_select_score(7, "http://127.0.0.1:18182|proxy=-", &idle, 1, 0, 10)
< target_select_score(7, "http://127.0.0.1:18181|proxy=-", &busy, 0, 5, 0)
);
}
#[test]
fn target_select_tie_break_distinguishes_equivalent_targets() {
let left = identity("endpoint-a", "candidate-a");
let right = identity("endpoint-b", "candidate-b");
assert_ne!(
target_select_tie_break(11, "http://127.0.0.1:18181|proxy=-", &left, 0),
target_select_tie_break(11, "http://127.0.0.1:18182|proxy=-", &right, 1)
);
}
#[test]
fn select_target_index_prefers_lower_in_flight_target() {
let choices = [
TargetSelectChoice {
target_key: "http://127.0.0.1:18181|proxy=-",
identity: identity("endpoint-a", "candidate-a"),
in_flight: 8,
selection_pressure_total: 0,
},
TargetSelectChoice {
target_key: "http://127.0.0.1:18182|proxy=-",
identity: identity("endpoint-b", "candidate-b"),
in_flight: 1,
selection_pressure_total: 100,
},
];
assert_eq!(select_target_index(17, &choices), 1);
}
#[test]
fn select_target_index_uses_selection_pressure_before_tie_break() {
let choices = [
TargetSelectChoice {
target_key: "http://127.0.0.1:18181|proxy=-",
identity: identity("endpoint-a", "candidate-a"),
in_flight: 0,
selection_pressure_total: 20,
},
TargetSelectChoice {
target_key: "http://127.0.0.1:18182|proxy=-",
identity: identity("endpoint-b", "candidate-b"),
in_flight: 0,
selection_pressure_total: 1,
},
];
assert_eq!(select_target_index(19, &choices), 1);
}
}
apps/aether-gateway/src/ai_serving/planner/standard/openai/chat/plans/sync.rs
+1
View File
@@ -134,6 +134,7 @@ impl LocalOpenAiChatSyncAttemptSource<'_> {
self.trace_id,
&self.body_json,
&self.input,
None,
attempt,
OPENAI_CHAT_SYNC_PLAN_KIND,
"openai_chat_sync_success",