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 gateway orchestration and failover effects

Browse Source
This commit is contained in:
fawney19
2026-04-18 11:35:11 +08:00
parent 3321bb3ccc
commit 569242d72f
42 changed files with 4365 additions and 1486 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -4,6 +4,7 @@ use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
use std::collections::BTreeSet;
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
use crate::orchestration::LocalExecutionCandidateMetadata;
use super::candidate_affinity::rank_eligible_local_execution_candidates;
use super::pool_scheduler::apply_local_execution_pool_scheduler;
@@ -13,6 +14,7 @@ pub(crate) struct EligibleLocalExecutionCandidate {
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
pub(crate) transport: GatewayProviderTransportSnapshot,
pub(crate) provider_api_format: String,
pub(crate) orchestration: LocalExecutionCandidateMetadata,
}
#[derive(Debug, Clone, PartialEq)]
@@ -122,6 +124,7 @@ where
provider_api_format: transport.endpoint.api_format.trim().to_ascii_lowercase(),
candidate,
transport,
orchestration: LocalExecutionCandidateMetadata::default(),
}),
}
}

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

@@ -9,15 +9,26 @@ use crate::ai_pipeline::planner::candidate_eligibility::{
use crate::ai_pipeline::planner::runtime_miss::record_local_runtime_candidate_skip_reason;
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
use crate::clock::current_unix_ms;
use crate::orchestration::{build_local_attempt_identities, ExecutionAttemptIdentity};
use crate::AppState;
#[derive(Debug, Clone)]
pub(crate) struct LocalExecutionCandidateAttempt {
pub(crate) eligible: EligibleLocalExecutionCandidate,
pub(crate) candidate_index: u32,
pub(crate) retry_index: u32,
pub(crate) pool_key_index: Option<u32>,
pub(crate) candidate_group_id: Option<String>,
pub(crate) candidate_id: String,
}
impl LocalExecutionCandidateAttempt {
pub(crate) fn attempt_identity(&self) -> ExecutionAttemptIdentity {
ExecutionAttemptIdentity::new(self.candidate_index, self.retry_index)
.with_pool_key_index(self.pool_key_index)
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct LocalAvailableCandidatePersistenceContext<'a> {
pub(crate) user_id: &'a str,
@@ -73,30 +84,43 @@ where
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value>,
{
let created_at_unix_ms = current_unix_ms();
let mut materialized = Vec::with_capacity(candidates.len());
let mut materialized = Vec::new();
for (candidate_index, eligible) in candidates.into_iter().enumerate() {
let generated_candidate_id = Uuid::new_v4().to_string();
let candidate_id = state
.persist_available_local_candidate(
trace_id,
user_id,
api_key_id,
&eligible.candidate,
candidate_index as u32,
&generated_candidate_id,
required_capabilities,
build_extra_data(&eligible),
created_at_unix_ms,
error_context,
)
.await;
let candidate_index = candidate_index as u32;
let attempt_identities =
build_local_attempt_identities(candidate_index, &eligible.transport)
.into_iter()
.map(|identity| identity.with_pool_key_index(eligible.orchestration.pool_key_index))
.collect::<Vec<_>>();
materialized.push(LocalExecutionCandidateAttempt {
eligible,
candidate_index: candidate_index as u32,
candidate_id,
});
for attempt_identity in attempt_identities {
let generated_candidate_id = Uuid::new_v4().to_string();
let candidate_id = state
.persist_available_local_candidate(
trace_id,
user_id,
api_key_id,
&eligible.candidate,
attempt_identity.candidate_index,
attempt_identity.retry_index,
&generated_candidate_id,
required_capabilities,
build_extra_data(&eligible),
created_at_unix_ms,
error_context,
)
.await;
materialized.push(LocalExecutionCandidateAttempt {
eligible: eligible.clone(),
candidate_index: attempt_identity.candidate_index,
retry_index: attempt_identity.retry_index,
pool_key_index: attempt_identity.pool_key_index,
candidate_group_id: eligible.orchestration.candidate_group_id.clone(),
candidate_id,
});
}
}
materialized
@@ -151,6 +175,7 @@ pub(crate) async fn persist_skipped_local_execution_candidate(
api_key_id,
candidate,
candidate_index,
0,
candidate_id,
required_capabilities,
skip_reason,

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

@@ -155,6 +155,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
Some(&input.requested_model),
&candidates,
);
let available_candidate_count = candidates.len() as u32;
let attempts = persist_available_local_execution_candidates_with_context(
planner_state,
trace_id,
@@ -180,7 +181,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
state,
trace_id,
persistence_policy.skipped,
attempts.len() as u32,
available_candidate_count,
skipped_candidates,
)
.await;

6
apps/aether-gateway/src/ai_pipeline/planner/passthrough/provider/family/payload.rs
View File

@@ -41,7 +41,9 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
let LocalSameFormatProviderCandidateAttempt {
eligible,
candidate_index,
candidate_group_id,
candidate_id,
..
} = &attempt;
let candidate = &eligible.candidate;
let resolved = resolve_local_same_format_provider_candidate_payload_parts(
@@ -83,8 +85,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id,
candidate_index: *candidate_index,
retry_index: 0,
attempt_identity: attempt.attempt_identity(),
model: &input.requested_model,
provider_name: &resolved.transport.provider.name,
provider_id: &candidate.provider_id,
@@ -94,6 +95,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
provider_api_format: spec_metadata.api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: candidate_group_id.as_deref(),
upstream_url: Some(&resolved.upstream_url),
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),

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

@@ -20,6 +20,7 @@ use crate::handlers::shared::{
parse_catalog_auth_config_json, provider_key_health_summary,
provider_key_status_snapshot_payload,
};
use crate::orchestration::LocalExecutionCandidateMetadata;
use crate::provider_key_auth::provider_key_auth_semantics;
const POOL_ACCOUNT_BLOCKED_SKIP_REASON: &str = "pool_account_blocked";
@@ -337,17 +338,27 @@ fn apply_local_execution_pool_scheduler_with_runtime_map(
let Some(group) = groups.remove(&group_key) else {
continue;
};
let candidate_group_id = local_execution_candidate_group_id(&group_key);
let Some(pool_config) =
pool_config_for_candidate(group.first().expect("group should exist"))
else {
reordered.extend(group);
reordered.extend(annotate_local_execution_group_candidates(
group,
candidate_group_id.as_str(),
false,
));
continue;
};
let runtime = runtime_by_provider
.get(&group_key.provider_id)
.unwrap_or(&default_runtime);
let (group_candidates, group_skipped) =
schedule_pool_group(group, pool_config, runtime, key_context_by_id);
let (group_candidates, group_skipped) = schedule_pool_group(
group,
pool_config,
runtime,
key_context_by_id,
candidate_group_id.as_str(),
);
reordered.extend(group_candidates);
skipped.extend(group_skipped);
}
@@ -366,6 +377,18 @@ fn pool_group_key(candidate: &EligibleLocalExecutionCandidate, pool_enabled: boo
}
}
fn local_execution_candidate_group_id(group_key: &PoolGroupKey) -> String {
format!(
"provider={}|endpoint={}|model={}|selected_model={}|api_format={}|singleton_key={}",
group_key.provider_id,
group_key.endpoint_id,
group_key.model_id,
group_key.selected_provider_model_name,
group_key.provider_api_format,
group_key.singleton_key_id.as_deref().unwrap_or("*"),
)
}
fn pool_config_for_candidate(
candidate: &EligibleLocalExecutionCandidate,
) -> Option<AdminProviderPoolConfig> {
@@ -377,6 +400,7 @@ fn schedule_pool_group(
pool_config: AdminProviderPoolConfig,
runtime: &AdminProviderPoolRuntimeState,
key_context_by_id: &BTreeMap<String, PoolCatalogKeyContext>,
candidate_group_id: &str,
) -> (
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
@@ -403,6 +427,7 @@ fn schedule_pool_group(
candidate,
transport,
provider_api_format,
orchestration,
} = eligible;
let key_id = candidate.key_id.clone();
let mut key_context = key_context_by_id.get(&key_id).cloned().unwrap_or_default();
@@ -460,6 +485,7 @@ fn schedule_pool_group(
candidate,
transport,
provider_api_format,
orchestration,
},
key_context,
original_index,
@@ -516,7 +542,28 @@ fn schedule_pool_group(
}
ordered.extend(available.into_iter().map(|item| item.eligible));
(ordered, skipped)
(
annotate_local_execution_group_candidates(ordered, candidate_group_id, true),
skipped,
)
}
fn annotate_local_execution_group_candidates(
candidates: Vec<EligibleLocalExecutionCandidate>,
candidate_group_id: &str,
pool_enabled: bool,
) -> Vec<EligibleLocalExecutionCandidate> {
candidates
.into_iter()
.enumerate()
.map(|(index, mut candidate)| {
candidate.orchestration = LocalExecutionCandidateMetadata {
candidate_group_id: Some(candidate_group_id.to_string()),
pool_key_index: pool_enabled.then_some(index as u32),
};
candidate
})
.collect()
}
#[derive(Debug)]
@@ -980,6 +1027,7 @@ mod tests {
use crate::handlers::shared::provider_pool::{
AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,
};
use crate::orchestration::LocalExecutionCandidateMetadata;
use crate::AppState;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
@@ -1039,6 +1087,72 @@ mod tests {
);
}
#[test]
fn pool_scheduler_attaches_group_and_pool_metadata_to_ranked_candidates() {
let pool_first = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-pool-a",
10,
Some(json!({ "pool_advanced": { "lru_enabled": true } })),
);
let other =
sample_eligible_candidate("provider-other", "endpoint-2", "key-other", 10, None);
let pool_second = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"key-pool-b",
10,
Some(json!({ "pool_advanced": { "lru_enabled": true } })),
);
let mut runtime_by_provider = BTreeMap::new();
runtime_by_provider.insert(
"provider-pool".to_string(),
AdminProviderPoolRuntimeState {
lru_score_by_key: BTreeMap::from([
("key-pool-a".to_string(), 20.0),
("key-pool-b".to_string(), 10.0),
]),
..AdminProviderPoolRuntimeState::default()
},
);
let (reordered, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
vec![pool_first, other, pool_second],
&runtime_by_provider,
&BTreeMap::new(),
);
assert!(skipped.is_empty());
assert_eq!(reordered.len(), 3);
assert_eq!(
reordered[0].orchestration,
LocalExecutionCandidateMetadata {
candidate_group_id: Some(
"provider=provider-pool|endpoint=endpoint-1|model=model-1|selected_model=gpt-5|api_format=openai:chat|singleton_key=*"
.to_string(),
),
pool_key_index: Some(0),
}
);
assert_eq!(reordered[1].orchestration.pool_key_index, Some(1));
assert_eq!(
reordered[1].orchestration.candidate_group_id,
reordered[0].orchestration.candidate_group_id
);
assert_eq!(
reordered[2].orchestration,
LocalExecutionCandidateMetadata {
candidate_group_id: Some(
"provider=provider-other|endpoint=endpoint-2|model=model-1|selected_model=gpt-5|api_format=openai:chat|singleton_key=key-other"
.to_string(),
),
pool_key_index: None,
}
);
}
#[test]
fn pool_scheduler_promotes_sticky_hit_before_other_sorted_keys() {
let key_a = sample_eligible_candidate(
@@ -1394,6 +1508,7 @@ mod tests {
mapping_matched_model: None,
},
provider_api_format: "openai:chat".to_string(),
orchestration: LocalExecutionCandidateMetadata::default(),
transport: crate::ai_pipeline::GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: provider_id.to_string(),

21
apps/aether-gateway/src/ai_pipeline/planner/report_context.rs
View File

@@ -3,13 +3,13 @@ use std::collections::BTreeMap;
use serde_json::{Map, Value};
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
use crate::orchestration::ExecutionAttemptIdentity;
pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) auth_context: &'a ExecutionRuntimeAuthContext,
pub(crate) request_id: &'a str,
pub(crate) candidate_id: &'a str,
pub(crate) candidate_index: u32,
pub(crate) retry_index: u32,
pub(crate) attempt_identity: ExecutionAttemptIdentity,
pub(crate) model: &'a str,
pub(crate) provider_name: &'a str,
pub(crate) provider_id: &'a str,
@@ -19,6 +19,7 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) provider_api_format: &'a str,
pub(crate) client_api_format: &'a str,
pub(crate) mapped_model: Option<&'a str>,
pub(crate) candidate_group_id: Option<&'a str>,
pub(crate) upstream_url: Option<&'a str>,
pub(crate) provider_request_method: Option<Value>,
pub(crate) provider_request_headers: Option<&'a BTreeMap<String, String>>,
@@ -69,11 +70,11 @@ pub(crate) fn build_local_execution_report_context(
);
object.insert(
"candidate_index".to_string(),
Value::Number(parts.candidate_index.into()),
Value::Number(parts.attempt_identity.candidate_index.into()),
);
object.insert(
"retry_index".to_string(),
Value::Number(parts.retry_index.into()),
Value::Number(parts.attempt_identity.retry_index.into()),
);
object.insert("model".to_string(), Value::String(parts.model.to_string()));
object.insert(
@@ -127,6 +128,12 @@ pub(crate) fn build_local_execution_report_context(
Value::String(mapped_model.to_string()),
);
}
if let Some(candidate_group_id) = parts.candidate_group_id {
object.insert(
"candidate_group_id".to_string(),
Value::String(candidate_group_id.to_string()),
);
}
if let Some(upstream_url) = parts.upstream_url {
object.insert(
"upstream_url".to_string(),
@@ -146,6 +153,12 @@ pub(crate) fn build_local_execution_report_context(
.expect("provider request headers should serialize"),
);
}
if let Some(pool_key_index) = parts.attempt_identity.pool_key_index {
object.insert(
"pool_key_index".to_string(),
Value::Number(pool_key_index.into()),
);
}
object.extend(parts.extra_fields);
Value::Object(object)

8
apps/aether-gateway/src/ai_pipeline/planner/specialized/files/decision.rs
View File

@@ -34,6 +34,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
) -> Option<GatewayControlSyncDecisionResponse> {
let spec_metadata = local_gemini_files_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let attempt_identity = attempt.attempt_identity();
let resolved = resolve_local_gemini_files_candidate_payload_parts(
state,
parts,
@@ -48,8 +49,9 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
.await?;
let LocalGeminiFilesCandidateAttempt {
eligible,
candidate_index,
candidate_group_id,
candidate_id,
..
} = attempt;
let candidate = eligible.candidate;
let transport = resolved.transport;
@@ -107,8 +109,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
candidate_index,
retry_index: 0,
attempt_identity,
model: "gemini-files",
provider_name: &transport.provider.name,
provider_id: &candidate.provider_id,
@@ -118,6 +119,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
mapped_model: None,
candidate_group_id: candidate_group_id.as_deref(),
upstream_url: None,
provider_request_method: None,
provider_request_headers: None,

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

@@ -105,6 +105,7 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
None,
&candidates,
);
let available_candidate_count = candidates.len() as u32;
let attempts = persist_available_local_execution_candidates_with_context(
planner_state,
trace_id,
@@ -132,7 +133,7 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
state,
trace_id,
persistence_policy.skipped,
attempts.len() as u32,
available_candidate_count,
skipped_candidates
.into_iter()
.map(|mut skipped_candidate| {

8
apps/aether-gateway/src/ai_pipeline/planner/specialized/video/decision.rs
View File

@@ -27,14 +27,16 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
) -> Option<GatewayControlSyncDecisionResponse> {
let spec_metadata = local_video_create_spec_metadata(spec);
let planner_state = PlannerAppState::new(state);
let attempt_identity = attempt.attempt_identity();
let resolved = resolve_local_video_create_candidate_payload_parts(
state, parts, body_json, trace_id, input, &attempt, spec,
)
.await?;
let LocalVideoCreateCandidateAttempt {
eligible,
candidate_index,
candidate_group_id,
candidate_id,
..
} = attempt;
let candidate = eligible.candidate;
let transport = resolved.transport;
@@ -90,8 +92,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
candidate_index,
retry_index: 0,
attempt_identity,
model: &input.requested_model,
provider_name: &transport.provider.name,
provider_id: &candidate.provider_id,
@@ -101,6 +102,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
provider_api_format: spec_metadata.api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: candidate_group_id.as_deref(),
upstream_url: None,
provider_request_method: None,
provider_request_headers: None,

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

@@ -182,6 +182,7 @@ async fn materialize_local_video_create_candidate_attempts(
Some(&input.requested_model),
&candidates,
);
let available_candidate_count = candidates.len() as u32;
let attempts = persist_available_local_execution_candidates_with_context(
state,
trace_id,
@@ -204,7 +205,7 @@ async fn materialize_local_video_create_candidate_attempts(
state.app(),
trace_id,
persistence_policy.skipped,
attempts.len() as u32,
available_candidate_count,
skipped_candidates
.into_iter()
.map(|mut skipped_candidate| {

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

@@ -234,6 +234,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
Some(&input.requested_model),
&candidates,
);
let available_candidate_count = candidates.len() as u32;
let attempts = persist_available_local_execution_candidates_with_context(
planner_state,
trace_id,
@@ -273,7 +274,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
state,
trace_id,
persistence_policy.skipped,
attempts.len() as u32,
available_candidate_count,
skipped_candidates,
)
.await;

6
apps/aether-gateway/src/ai_pipeline/planner/standard/family/payload.rs
View File

@@ -35,7 +35,9 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
let LocalStandardCandidateAttempt {
eligible,
candidate_index,
candidate_group_id,
candidate_id,
..
} = &attempt;
let candidate = &eligible.candidate;
let resolved = resolve_local_standard_candidate_payload_parts(
@@ -89,8 +91,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id,
candidate_index: *candidate_index,
retry_index: 0,
attempt_identity: attempt.attempt_identity(),
model: &input.requested_model,
provider_name: &candidate.provider_name,
provider_id: &candidate.provider_id,
@@ -100,6 +101,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
provider_api_format: &resolved.provider_api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: candidate_group_id.as_deref(),
upstream_url: Some(&resolved.upstream_url),
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),

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

@@ -28,10 +28,13 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
report_kind: &str,
upstream_is_stream: bool,
) -> Option<GatewayControlSyncDecisionResponse> {
let attempt_identity = attempt.attempt_identity();
let LocalOpenAiChatCandidateAttempt {
eligible,
candidate_index,
candidate_group_id,
candidate_id,
..
} = attempt;
let resolved = resolve_local_openai_chat_candidate_payload_parts(
state,
@@ -105,8 +108,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
candidate_index,
retry_index: 0,
attempt_identity,
model: &input.requested_model,
provider_name: &resolved.transport.provider.name,
provider_id: &candidate.provider_id,
@@ -116,6 +118,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
provider_api_format: &resolved.provider_api_format,
client_api_format: "openai:chat",
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: candidate_group_id.as_deref(),
upstream_url: Some(&resolved.upstream_url),
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),

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

@@ -122,6 +122,7 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
Some(&input.requested_model),
&candidates,
);
let available_candidate_count = candidates.len() as u32;
let attempts = persist_available_local_execution_candidates_with_context(
planner_state,
trace_id,
@@ -156,7 +157,7 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
state,
trace_id,
persistence_policy.skipped,
attempts.len() as u32,
available_candidate_count,
skipped_candidates,
)
.await;

7
apps/aether-gateway/src/ai_pipeline/planner/standard/openai/cli/decision/payload.rs
View File

@@ -31,10 +31,13 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
spec: LocalOpenAiCliSpec,
) -> Option<GatewayControlSyncDecisionResponse> {
let spec_metadata = local_openai_cli_spec_metadata(spec);
let attempt_identity = attempt.attempt_identity();
let LocalOpenAiCliCandidateAttempt {
eligible,
candidate_index,
candidate_group_id,
candidate_id,
..
} = attempt;
let resolved = resolve_local_openai_cli_candidate_payload_parts(
state,
@@ -133,8 +136,7 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
auth_context: &input.auth_context,
request_id: trace_id,
candidate_id: &candidate_id,
candidate_index,
retry_index: 0,
attempt_identity,
model: &input.requested_model,
provider_name: &resolved.transport.provider.name,
provider_id: &candidate.provider_id,
@@ -144,6 +146,7 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
provider_api_format: &resolved.provider_api_format,
client_api_format: spec_metadata.api_format,
mapped_model: Some(&resolved.mapped_model),
candidate_group_id: candidate_group_id.as_deref(),
upstream_url: Some(&resolved.upstream_url),
provider_request_method: Some(serde_json::Value::Null),
provider_request_headers: Some(&resolved.provider_request_headers),

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

@@ -230,6 +230,7 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
Some(&input.requested_model),
&candidates,
);
let available_candidate_count = candidates.len() as u32;
let attempts = persist_available_local_execution_candidates_with_context(
planner_state,
trace_id,
@@ -269,7 +270,7 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
state,
trace_id,
persistence_policy.skipped,
attempts.len() as u32,
available_candidate_count,
skipped_candidates,
)
.await;

4
apps/aether-gateway/src/ai_pipeline/planner/state/candidate_runtime.rs
View File

@@ -29,6 +29,7 @@ impl<'a> PlannerAppState<'a> {
api_key_id: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
retry_index: u32,
candidate_id: &str,
required_capabilities: Option<&Value>,
extra_data: Option<Value>,
@@ -42,6 +43,7 @@ impl<'a> PlannerAppState<'a> {
api_key_id,
candidate,
candidate_index,
retry_index,
candidate_id,
required_capabilities,
extra_data,
@@ -59,6 +61,7 @@ impl<'a> PlannerAppState<'a> {
api_key_id: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
retry_index: u32,
candidate_id: &str,
required_capabilities: Option<&Value>,
skip_reason: &str,
@@ -73,6 +76,7 @@ impl<'a> PlannerAppState<'a> {
api_key_id,
candidate,
candidate_index,
retry_index,
candidate_id,
required_capabilities,
skip_reason,

452
apps/aether-gateway/src/execution_runtime/fallback.rs
View File

@@ -1,24 +1,11 @@
use std::collections::BTreeSet;
use aether_contracts::{ExecutionPlan, ExecutionResult};
use regex::Regex;
use serde_json::{json, Value};
use tracing::debug;
use serde_json::Value;
use crate::provider_transport::GatewayProviderTransportSnapshot;
use crate::orchestration::{
resolve_local_failover_analysis_for_attempt, LocalFailoverAnalysis, LocalFailoverDecision,
};
use crate::AppState;
fn local_candidate_index(report_context: Option<&serde_json::Value>) -> Option<u64> {
report_context
.and_then(serde_json::Value::as_object)
.and_then(|context| context.get("candidate_index"))
.and_then(serde_json::Value::as_u64)
}
fn should_failover_local_upstream_status(status_code: u16) -> bool {
status_code >= 400
}
fn sync_plan_kind_disables_local_candidate_failover(plan_kind: &str) -> bool {
matches!(
plan_kind,
@@ -26,38 +13,6 @@ fn sync_plan_kind_disables_local_candidate_failover(plan_kind: &str) -> bool {
)
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct LocalFailoverPolicy {
max_retries: Option<u64>,
stop_status_codes: BTreeSet<u16>,
continue_status_codes: BTreeSet<u16>,
success_failover_patterns: Vec<LocalFailoverRegexRule>,
error_stop_patterns: Vec<LocalFailoverRegexRule>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct LocalFailoverRegexRule {
pattern: String,
status_codes: BTreeSet<u16>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum LocalFailoverDecision {
UseDefault,
RetryNextCandidate,
StopLocalFailover,
}
impl LocalFailoverDecision {
pub(crate) const fn as_str(self) -> &'static str {
match self {
Self::UseDefault => "use_default",
Self::RetryNextCandidate => "retry_next_candidate",
Self::StopLocalFailover => "stop_local_failover",
}
}
}
pub(crate) async fn should_retry_next_local_candidate_sync(
state: &AppState,
plan: &ExecutionPlan,
@@ -66,22 +21,43 @@ pub(crate) async fn should_retry_next_local_candidate_sync(
result: &ExecutionResult,
response_text: Option<&str>,
) -> bool {
if sync_plan_kind_disables_local_candidate_failover(plan_kind) {
return false;
}
matches!(
resolve_local_failover_decision(
analyze_local_candidate_failover_sync(
state,
plan,
plan_kind,
report_context,
result.status_code,
result,
response_text,
)
.await,
.await
.decision,
LocalFailoverDecision::RetryNextCandidate
)
}
pub(crate) async fn analyze_local_candidate_failover_sync(
state: &AppState,
plan: &ExecutionPlan,
plan_kind: &str,
report_context: Option<&serde_json::Value>,
result: &ExecutionResult,
response_text: Option<&str>,
) -> LocalFailoverAnalysis {
if sync_plan_kind_disables_local_candidate_failover(plan_kind) {
return LocalFailoverAnalysis::use_default();
}
resolve_local_failover_analysis_for_attempt(
state,
plan,
report_context,
result.status_code,
response_text,
)
.await
}
pub(crate) async fn should_stop_local_candidate_failover_sync(
state: &AppState,
plan: &ExecutionPlan,
@@ -90,19 +66,20 @@ pub(crate) async fn should_stop_local_candidate_failover_sync(
result: &ExecutionResult,
response_text: Option<&str>,
) -> bool {
if sync_plan_kind_disables_local_candidate_failover(plan_kind) {
return false;
}
matches!(
resolve_local_failover_decision(
analyze_local_candidate_failover_sync(
state,
plan,
plan_kind,
report_context,
result.status_code,
result,
response_text,
)
.await,
LocalFailoverDecision::StopLocalFailover
LocalFailoverAnalysis {
decision: LocalFailoverDecision::StopLocalFailover,
..
}
)
}
@@ -195,7 +172,7 @@ pub(crate) async fn should_retry_next_local_candidate_stream(
response_text: Option<&str>,
) -> bool {
matches!(
resolve_local_candidate_failover_decision_stream(
resolve_local_candidate_failover_analysis_stream(
state,
plan,
report_context,
@@ -203,7 +180,10 @@ pub(crate) async fn should_retry_next_local_candidate_stream(
response_text,
)
.await,
LocalFailoverDecision::RetryNextCandidate
LocalFailoverAnalysis {
decision: LocalFailoverDecision::RetryNextCandidate,
..
}
)
}
@@ -216,7 +196,7 @@ pub(crate) async fn should_stop_local_candidate_failover_stream(
response_text: Option<&str>,
) -> bool {
matches!(
resolve_local_candidate_failover_decision_stream(
resolve_local_candidate_failover_analysis_stream(
state,
plan,
report_context,
@@ -224,10 +204,30 @@ pub(crate) async fn should_stop_local_candidate_failover_stream(
response_text,
)
.await,
LocalFailoverDecision::StopLocalFailover
LocalFailoverAnalysis {
decision: LocalFailoverDecision::StopLocalFailover,
..
}
)
}
pub(crate) async fn resolve_local_candidate_failover_analysis_stream(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
status_code: u16,
response_text: Option<&str>,
) -> LocalFailoverAnalysis {
resolve_local_failover_analysis_for_attempt(
state,
plan,
report_context,
status_code,
response_text,
)
.await
}
pub(crate) async fn resolve_local_candidate_failover_decision_stream(
state: &AppState,
plan: &ExecutionPlan,
@@ -235,7 +235,15 @@ pub(crate) async fn resolve_local_candidate_failover_decision_stream(
status_code: u16,
response_text: Option<&str>,
) -> LocalFailoverDecision {
resolve_local_failover_decision(state, plan, report_context, status_code, response_text).await
resolve_local_candidate_failover_analysis_stream(
state,
plan,
report_context,
status_code,
response_text,
)
.await
.decision
}
pub(crate) fn local_failover_response_text(
@@ -255,304 +263,6 @@ pub(crate) fn local_failover_response_text(
.map(ToOwned::to_owned)
}
async fn resolve_local_failover_decision(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
status_code: u16,
response_text: Option<&str>,
) -> LocalFailoverDecision {
let Some(candidate_index) = local_candidate_index(report_context) else {
return LocalFailoverDecision::UseDefault;
};
let policy = resolve_local_failover_policy(state, plan, report_context).await;
let response_text = response_text
.map(str::trim)
.filter(|value| !value.is_empty());
if policy.stop_status_codes.contains(&status_code) {
return LocalFailoverDecision::StopLocalFailover;
}
if status_code >= 400
&& response_text.is_some_and(|text| {
policy
.error_stop_patterns
.iter()
.any(|rule| local_failover_regex_rule_matches(rule, text, status_code))
})
{
return LocalFailoverDecision::StopLocalFailover;
}
if policy
.max_retries
.is_some_and(|max_retries| candidate_index >= max_retries)
{
return LocalFailoverDecision::UseDefault;
}
if status_code == 200
&& response_text.is_some_and(|text| {
policy
.success_failover_patterns
.iter()
.any(|rule| local_failover_regex_rule_matches(rule, text, status_code))
})
{
return LocalFailoverDecision::RetryNextCandidate;
}
if policy.continue_status_codes.contains(&status_code) {
return LocalFailoverDecision::RetryNextCandidate;
}
if should_failover_local_upstream_status(status_code) {
return LocalFailoverDecision::RetryNextCandidate;
}
LocalFailoverDecision::UseDefault
}
async fn resolve_local_failover_policy(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> LocalFailoverPolicy {
if let Some(policy) = local_failover_policy_from_report_context(report_context) {
debug!(
event_name = "local_failover_policy_loaded",
log_type = "debug",
request_id = %plan.request_id,
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
source = "report_context",
max_retries = ?policy.max_retries,
stop_status_code_count = policy.stop_status_codes.len(),
continue_status_code_count = policy.continue_status_codes.len(),
success_failover_pattern_count = policy.success_failover_patterns.len(),
error_stop_pattern_count = policy.error_stop_patterns.len(),
"gateway loaded local failover policy from report context"
);
return policy;
}
let transport = match state
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
.await
{
Ok(Some(transport)) => transport,
Ok(None) | Err(_) => return LocalFailoverPolicy::default(),
};
let policy = local_failover_policy_from_transport(&transport);
debug!(
event_name = "local_failover_policy_loaded",
log_type = "debug",
request_id = %plan.request_id,
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
source = "transport_snapshot",
max_retries = ?policy.max_retries,
stop_status_code_count = policy.stop_status_codes.len(),
continue_status_code_count = policy.continue_status_codes.len(),
success_failover_pattern_count = policy.success_failover_patterns.len(),
error_stop_pattern_count = policy.error_stop_patterns.len(),
"gateway loaded local failover policy from transport snapshot"
);
policy
}
fn local_failover_policy_from_transport(
transport: &GatewayProviderTransportSnapshot,
) -> LocalFailoverPolicy {
let rules = transport
.provider
.config
.as_ref()
.and_then(|config| config.get("failover_rules"))
.and_then(serde_json::Value::as_object);
let max_retries = rules
.and_then(|value| value.get("max_retries"))
.and_then(parse_u64_value)
.or_else(|| {
transport
.endpoint
.max_retries
.and_then(|value| u64::try_from(value).ok())
})
.or_else(|| {
transport
.provider
.max_retries
.and_then(|value| u64::try_from(value).ok())
});
LocalFailoverPolicy {
max_retries,
stop_status_codes: rules
.map(|value| {
parse_status_code_set(
value,
&[
"stop_on_status_codes",
"early_stop_status_codes",
"non_retryable_status_codes",
"stop_status_codes",
],
)
})
.unwrap_or_default(),
continue_status_codes: rules
.map(|value| {
parse_status_code_set(
value,
&[
"continue_on_status_codes",
"retryable_status_codes",
"retry_on_status_codes",
"continue_status_codes",
],
)
})
.unwrap_or_default(),
success_failover_patterns: rules
.map(|value| parse_regex_rules(value, "success_failover_patterns"))
.unwrap_or_default(),
error_stop_patterns: rules
.map(|value| parse_regex_rules(value, "error_stop_patterns"))
.unwrap_or_default(),
}
}
fn local_failover_policy_from_report_context(
report_context: Option<&Value>,
) -> Option<LocalFailoverPolicy> {
let object = report_context
.and_then(Value::as_object)?
.get("local_failover_policy")?
.as_object()?;
Some(LocalFailoverPolicy {
max_retries: object.get("max_retries").and_then(parse_u64_value),
stop_status_codes: object
.get("stop_status_codes")
.map(parse_status_code_list)
.unwrap_or_default(),
continue_status_codes: object
.get("continue_status_codes")
.map(parse_status_code_list)
.unwrap_or_default(),
success_failover_patterns: parse_regex_rules(object, "success_failover_patterns"),
error_stop_patterns: parse_regex_rules(object, "error_stop_patterns"),
})
}
fn parse_status_code_list(value: &Value) -> BTreeSet<u16> {
value
.as_array()
.into_iter()
.flat_map(|values| values.iter())
.filter_map(|value| parse_u64_value(value).and_then(|value| u16::try_from(value).ok()))
.collect()
}
fn local_failover_policy_to_value(policy: &LocalFailoverPolicy) -> Value {
json!({
"max_retries": policy.max_retries,
"stop_status_codes": policy.stop_status_codes.iter().copied().collect::<Vec<_>>(),
"continue_status_codes": policy.continue_status_codes.iter().copied().collect::<Vec<_>>(),
"success_failover_patterns": policy.success_failover_patterns.iter().map(local_failover_regex_rule_to_value).collect::<Vec<_>>(),
"error_stop_patterns": policy.error_stop_patterns.iter().map(local_failover_regex_rule_to_value).collect::<Vec<_>>(),
})
}
fn local_failover_regex_rule_to_value(rule: &LocalFailoverRegexRule) -> Value {
json!({
"pattern": rule.pattern,
"status_codes": rule.status_codes.iter().copied().collect::<Vec<_>>(),
})
}
pub(crate) fn append_local_failover_policy_to_value(
value: Value,
transport: &GatewayProviderTransportSnapshot,
) -> Value {
let Value::Object(mut object) = value else {
return value;
};
object.insert(
"local_failover_policy".to_string(),
local_failover_policy_to_value(&local_failover_policy_from_transport(transport)),
);
Value::Object(object)
}
fn parse_regex_rules(
rules: &serde_json::Map<String, serde_json::Value>,
key: &str,
) -> Vec<LocalFailoverRegexRule> {
rules
.get(key)
.and_then(serde_json::Value::as_array)
.into_iter()
.flat_map(|items| items.iter())
.filter_map(parse_regex_rule)
.collect()
}
fn parse_regex_rule(value: &serde_json::Value) -> Option<LocalFailoverRegexRule> {
let object = value.as_object()?;
let pattern = object
.get("pattern")
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
Some(LocalFailoverRegexRule {
pattern: pattern.to_string(),
status_codes: object
.get("status_codes")
.and_then(serde_json::Value::as_array)
.into_iter()
.flat_map(|values| values.iter())
.filter_map(|value| parse_u64_value(value).and_then(|value| u16::try_from(value).ok()))
.collect(),
})
}
fn local_failover_regex_rule_matches(
rule: &LocalFailoverRegexRule,
response_text: &str,
status_code: u16,
) -> bool {
if !rule.status_codes.is_empty() && !rule.status_codes.contains(&status_code) {
return false;
}
Regex::new(&rule.pattern)
.ok()
.is_some_and(|regex| regex.is_match(response_text))
}
fn parse_status_code_set(
rules: &serde_json::Map<String, serde_json::Value>,
keys: &[&str],
) -> BTreeSet<u16> {
keys.iter()
.filter_map(|key| rules.get(*key))
.filter_map(serde_json::Value::as_array)
.flat_map(|values| values.iter())
.filter_map(|value| parse_u64_value(value).and_then(|value| u16::try_from(value).ok()))
.collect()
}
fn parse_u64_value(value: &serde_json::Value) -> Option<u64> {
value
.as_u64()
.or_else(|| value.as_i64().and_then(|value| u64::try_from(value).ok()))
}
pub(crate) fn should_fallback_to_control_stream(
plan_kind: &str,
status_code: u16,
@@ -623,13 +333,15 @@ mod tests {
use super::{
resolve_core_stream_error_finalize_report_kind,
resolve_core_sync_error_finalize_report_kind, resolve_local_failover_policy,
should_fallback_to_control_stream, should_fallback_to_control_sync,
should_retry_next_local_candidate_stream, should_retry_next_local_candidate_sync,
should_stop_local_candidate_failover_stream, should_stop_local_candidate_failover_sync,
LocalFailoverPolicy, LocalFailoverRegexRule,
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_stream,
should_fallback_to_control_sync, should_retry_next_local_candidate_stream,
should_retry_next_local_candidate_sync, should_stop_local_candidate_failover_stream,
should_stop_local_candidate_failover_sync,
};
use crate::data::GatewayDataState;
use crate::orchestration::{
resolve_local_failover_policy, LocalFailoverPolicy, LocalFailoverRegexRule,
};
use crate::AppState;
fn sample_plan() -> aether_contracts::ExecutionPlan {
@@ -1122,7 +834,7 @@ mod tests {
.await
);
assert!(
!should_retry_next_local_candidate_stream(
should_retry_next_local_candidate_stream(
&state,
&plan,
"openai_chat_stream",

10
apps/aether-gateway/src/execution_runtime/mod.rs
View File

@@ -6,7 +6,6 @@ use serde_json::{Map, Value};
mod constants;
mod fallback;
pub(crate) mod ndjson;
mod pool_feedback;
#[cfg(test)]
pub(crate) mod remote_compat;
mod server;
@@ -20,18 +19,17 @@ pub(crate) use self::constants::{
MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES,
};
pub(crate) use self::fallback::{
append_local_failover_policy_to_value, local_failover_response_text,
analyze_local_candidate_failover_sync, local_failover_response_text,
resolve_core_stream_direct_finalize_report_kind,
resolve_core_stream_error_finalize_report_kind, resolve_core_sync_error_finalize_report_kind,
resolve_local_candidate_failover_analysis_stream,
resolve_local_candidate_failover_decision_stream, should_fallback_to_control_stream,
should_fallback_to_control_sync, should_finalize_sync_response,
should_retry_next_local_candidate_stream, should_retry_next_local_candidate_sync,
should_stop_local_candidate_failover_stream, should_stop_local_candidate_failover_sync,
LocalFailoverDecision,
};
pub(crate) use pool_feedback::{
record_pool_error_feedback, record_pool_stream_timeout_feedback,
record_stream_pool_success_feedback, record_sync_pool_success_feedback,
pub(crate) use crate::orchestration::{
append_local_failover_policy_to_value, LocalFailoverAnalysis, LocalFailoverDecision,
};
pub use server::{
build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit,

188
apps/aether-gateway/src/execution_runtime/pool_feedback.rs
View File

@@ -1,188 +0,0 @@
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
use aether_usage_runtime::{
build_stream_terminal_usage_outcome, build_sync_terminal_usage_outcome, TerminalUsageOutcome,
};
use serde_json::Value;
use std::collections::BTreeMap;
use tracing::warn;
use crate::ai_pipeline::extract_pool_sticky_session_token;
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
use crate::handlers::shared::provider_pool::{
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
record_admin_provider_pool_success, AdminProviderPoolConfig,
};
use crate::usage::{GatewayStreamReportRequest, GatewaySyncReportRequest};
use crate::AppState;
struct PoolFeedbackContext {
runner: aether_data::redis::RedisKvRunner,
pool_config: AdminProviderPoolConfig,
sticky_session_token: Option<String>,
}
fn pool_feedback_request_body<'a>(
plan: &'a ExecutionPlan,
report_context: Option<&'a Value>,
) -> Option<&'a Value> {
report_context
.and_then(Value::as_object)
.and_then(|object| object.get("original_request_body"))
.filter(|value| !value.is_null())
.or(plan.body.json_body.as_ref())
}
async fn resolve_pool_feedback_context(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
) -> Option<PoolFeedbackContext> {
let Some(runner) = state.redis_kv_runner() else {
return None;
};
let transport = match state
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
.await
{
Ok(Some(transport)) => transport,
Ok(None) => return None,
Err(err) => {
warn!(
"gateway execution runtime pool feedback: failed to read transport snapshot for provider {} endpoint {} key {}: {:?}",
plan.provider_id, plan.endpoint_id, plan.key_id, err
);
return None;
}
};
let Some(pool_config) =
admin_provider_pool_config_from_config_value(transport.provider.config.as_ref())
else {
return None;
};
let sticky_session_token = pool_feedback_request_body(plan, report_context)
.and_then(extract_pool_sticky_session_token);
Some(PoolFeedbackContext {
runner,
pool_config,
sticky_session_token,
})
}
fn total_tokens_used(outcome: &TerminalUsageOutcome) -> u64 {
outcome
.standardized_usage
.as_ref()
.map(|usage| {
usage
.input_tokens
.saturating_add(usage.output_tokens)
.max(0) as u64
})
.unwrap_or(0)
}
fn resolve_ttfb_ms(telemetry: Option<&ExecutionTelemetry>) -> Option<u64> {
telemetry.and_then(|telemetry| telemetry.ttfb_ms.or(telemetry.elapsed_ms))
}
pub(crate) async fn record_sync_pool_success_feedback(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
payload: &GatewaySyncReportRequest,
) {
let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
return;
};
let usage_outcome = build_sync_terminal_usage_outcome(plan, report_context, payload);
record_admin_provider_pool_success(
&context.runner,
&plan.provider_id,
&plan.key_id,
&context.pool_config,
context.sticky_session_token.as_deref(),
total_tokens_used(&usage_outcome),
resolve_ttfb_ms(payload.telemetry.as_ref()),
)
.await;
}
pub(crate) async fn record_stream_pool_success_feedback(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
payload: &GatewayStreamReportRequest,
) {
let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
return;
};
let usage_outcome = build_stream_terminal_usage_outcome(plan, report_context, payload);
record_admin_provider_pool_success(
&context.runner,
&plan.provider_id,
&plan.key_id,
&context.pool_config,
context.sticky_session_token.as_deref(),
total_tokens_used(&usage_outcome),
resolve_ttfb_ms(payload.telemetry.as_ref()),
)
.await;
}
pub(crate) async fn record_pool_error_feedback(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
status_code: u16,
headers: &BTreeMap<String, String>,
error_body: Option<&str>,
) {
if status_code < 400 {
return;
}
let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
return;
};
if status_code == 401 {
let _ = state
.invalidate_local_oauth_refresh_entry(&plan.key_id)
.await;
}
record_admin_provider_pool_error(
&context.runner,
&plan.provider_id,
&plan.key_id,
&context.pool_config,
status_code,
error_body,
Some(headers),
)
.await;
}
pub(crate) async fn record_pool_stream_timeout_feedback(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
) {
let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
return;
};
record_admin_provider_pool_stream_timeout(
&context.runner,
&plan.provider_id,
&plan.key_id,
&context.pool_config,
)
.await;
}

106
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -53,14 +53,18 @@ use crate::execution_runtime::transport::{
DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
};
use crate::execution_runtime::{
local_failover_response_text, record_pool_error_feedback, record_stream_pool_success_feedback,
resolve_core_stream_direct_finalize_report_kind,
local_failover_response_text, resolve_core_stream_direct_finalize_report_kind,
resolve_core_stream_error_finalize_report_kind,
resolve_local_candidate_failover_decision_stream, should_fallback_to_control_stream,
resolve_local_candidate_failover_analysis_stream, should_fallback_to_control_stream,
should_retry_next_local_candidate_stream, LocalFailoverDecision,
};
use crate::execution_runtime::{MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES};
use crate::log_ids::short_request_id;
use crate::orchestration::{
apply_local_execution_effect, LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect,
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
};
use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, record_local_request_candidate_status,
};
@@ -566,16 +570,7 @@ async fn execute_stream_from_frame_stream(
let (body_json, body_base64) = decode_stream_error_body(&headers, &error_body);
let error_response_text =
local_failover_response_text(body_json.as_ref(), &error_body, None);
record_pool_error_feedback(
state,
&plan,
report_context.as_ref(),
status_code,
&headers,
error_response_text.as_deref(),
)
.await;
let failover_decision = resolve_local_candidate_failover_decision_stream(
let failover_analysis = resolve_local_candidate_failover_analysis_stream(
state,
&plan,
report_context.as_ref(),
@@ -583,6 +578,70 @@ async fn execute_stream_from_frame_stream(
error_response_text.as_deref(),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code,
classification: failover_analysis.classification,
}),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect {
status_code,
classification: failover_analysis.classification,
headers: Some(&headers),
}),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
status_code,
classification: failover_analysis.classification,
}),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
status_code,
response_text: error_response_text.as_deref(),
}),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::PoolError(LocalPoolErrorEffect {
status_code,
classification: failover_analysis.classification,
headers: &headers,
error_body: error_response_text.as_deref(),
}),
)
.await;
let failover_decision = failover_analysis.decision;
debug!(
event_name = "execution_runtime_stream_failover_decided",
log_type = "debug",
@@ -1502,11 +1561,24 @@ async fn execute_stream_from_frame_stream(
}),
telemetry: telemetry.clone(),
};
record_stream_pool_success_feedback(
apply_local_execution_effect(
&state_for_report,
&plan_for_report,
report_context_owned.as_ref(),
&usage_payload,
LocalExecutionEffectContext {
plan: &plan_for_report,
report_context: report_context_owned.as_ref(),
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
apply_local_execution_effect(
&state_for_report,
LocalExecutionEffectContext {
plan: &plan_for_report,
report_context: report_context_owned.as_ref(),
},
LocalExecutionEffect::PoolSuccessStream {
payload: &usage_payload,
},
)
.await;
record_stream_terminal_usage(

93
apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs
View File

@@ -16,8 +16,13 @@ use crate::control::GatewayControlDecision;
use crate::execution_runtime::submission::{
resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize,
};
use crate::execution_runtime::{record_pool_error_feedback, record_pool_stream_timeout_feedback};
use crate::log_ids::short_request_id;
use crate::orchestration::{
apply_local_execution_effect, resolve_local_failover_analysis_for_attempt,
LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect, LocalExecutionEffect,
LocalExecutionEffectContext, LocalHealthFailureEffect, LocalOAuthInvalidationEffect,
LocalPoolErrorEffect,
};
use crate::request_candidate_runtime::record_report_request_candidate_status;
use crate::usage::submit_sync_report;
use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
@@ -123,22 +128,92 @@ async fn record_stream_sync_failure(
failure: &StreamFailureReport,
started_at_unix_ms: Option<u64>,
) {
if matches!(
failure.error_type.as_str(),
"first_byte_timeout" | "read_timeout"
) {
record_pool_stream_timeout_feedback(state, plan, report_context).await;
}
let error_body = serde_json::to_string(&failure.body_json).ok();
record_pool_error_feedback(
let failure_analysis = resolve_local_failover_analysis_for_attempt(
state,
plan,
report_context,
failure.status_code,
&payload.headers,
error_body.as_deref(),
)
.await;
if matches!(
failure.error_type.as_str(),
"first_byte_timeout" | "read_timeout"
) {
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan,
report_context,
},
LocalExecutionEffect::PoolStreamTimeout,
)
.await;
}
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan,
report_context,
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: failure.status_code,
classification: failure_analysis.classification,
}),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan,
report_context,
},
LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect {
status_code: failure.status_code,
classification: failure_analysis.classification,
headers: Some(&payload.headers),
}),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan,
report_context,
},
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
status_code: failure.status_code,
classification: failure_analysis.classification,
}),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan,
report_context,
},
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
status_code: failure.status_code,
response_text: error_body.as_deref(),
}),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan,
report_context,
},
LocalExecutionEffect::PoolError(LocalPoolErrorEffect {
status_code: failure.status_code,
classification: failure_analysis.classification,
headers: &payload.headers,
error_body: error_body.as_deref(),
}),
)
.await;
let context_seed = build_terminal_usage_context_seed(plan, report_context);
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
state

135
apps/aether-gateway/src/execution_runtime/sync/execution.rs
View File

@@ -30,12 +30,16 @@ use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_
use crate::execution_runtime::submission::submit_local_core_error_or_sync_finalize;
use crate::execution_runtime::transport::DirectSyncExecutionRuntime;
use crate::execution_runtime::{
local_failover_response_text, record_pool_error_feedback, record_sync_pool_success_feedback,
analyze_local_candidate_failover_sync, local_failover_response_text,
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_sync,
should_finalize_sync_response, should_retry_next_local_candidate_sync,
should_stop_local_candidate_failover_sync,
should_finalize_sync_response, LocalFailoverDecision,
};
use crate::log_ids::short_request_id;
use crate::orchestration::{
apply_local_execution_effect, LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect,
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
};
use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, record_local_request_candidate_status,
};
@@ -247,7 +251,7 @@ pub(crate) async fn execute_execution_runtime_sync(
&body_bytes,
result.error.as_ref().map(|error| error.message.as_str()),
);
let stop_local_failover = should_stop_local_candidate_failover_sync(
let local_failover_analysis = analyze_local_candidate_failover_sync(
state,
&plan,
plan_kind,
@@ -257,27 +261,74 @@ pub(crate) async fn execute_execution_runtime_sync(
)
.await;
if result.status_code >= 400 {
record_pool_error_feedback(
apply_local_execution_effect(
state,
&plan,
report_context.as_ref(),
result.status_code,
&headers,
local_failover_response_text.as_deref(),
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: result.status_code,
classification: local_failover_analysis.classification,
}),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect {
status_code: result.status_code,
classification: local_failover_analysis.classification,
headers: Some(&headers),
}),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
status_code: result.status_code,
classification: local_failover_analysis.classification,
}),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
status_code: result.status_code,
response_text: local_failover_response_text.as_deref(),
}),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::PoolError(LocalPoolErrorEffect {
status_code: result.status_code,
classification: local_failover_analysis.classification,
headers: &headers,
error_body: local_failover_response_text.as_deref(),
}),
)
.await;
}
if should_retry_next_local_candidate_sync(
state,
&plan,
plan_kind,
report_context.as_ref(),
&result,
local_failover_response_text.as_deref(),
)
.await
&& !stop_local_failover
{
if matches!(
local_failover_analysis.decision,
LocalFailoverDecision::RetryNextCandidate
) {
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
@@ -341,16 +392,17 @@ pub(crate) async fn execute_execution_runtime_sync(
mapped_error_finalize_kind.clone()
};
if !stop_local_failover
&& should_fallback_to_control_sync(
plan_kind,
&result,
body_json.as_ref(),
has_body_bytes,
explicit_finalize || implicit_finalize.is_some(),
mapped_error_finalize_kind.is_some(),
)
{
if !matches!(
local_failover_analysis.decision,
LocalFailoverDecision::StopLocalFailover
) && should_fallback_to_control_sync(
plan_kind,
&result,
body_json.as_ref(),
has_body_bytes,
explicit_finalize || implicit_finalize.is_some(),
mapped_error_finalize_kind.is_some(),
) {
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
@@ -406,11 +458,24 @@ pub(crate) async fn execute_execution_runtime_sync(
telemetry: result.telemetry.clone(),
};
if result.status_code < 400 {
record_sync_pool_success_feedback(
apply_local_execution_effect(
state,
&plan,
report_context.as_ref(),
&base_usage_payload,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::PoolSuccessSync {
payload: &base_usage_payload,
},
)
.await;
}

2
apps/aether-gateway/src/handlers/admin/request/features.rs
View File

@@ -128,7 +128,7 @@ impl<'a> AdminAppState<'a> {
display_name: Option<&str>,
mime_type: Option<&str>,
) -> Result<(), GatewayError> {
crate::usage::reporting::store_local_gemini_file_mapping(
crate::orchestration::store_local_gemini_file_mapping(
self.app,
file_name,
key_id,

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

@@ -48,6 +48,7 @@ mod log_ids;
mod maintenance;
pub(crate) mod middleware;
mod model_fetch;
mod orchestration;
mod provider_key_auth;
pub(crate) use aether_provider_transport as provider_transport;
mod query;

335
apps/aether-gateway/src/orchestration/adaptive.rs Normal file
View File

@@ -0,0 +1,335 @@
use std::collections::BTreeMap;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::{json, Value};
use super::LocalFailoverClassification;
use crate::handlers::shared::default_provider_key_status_snapshot;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct LocalAdaptiveRateLimitProjection {
pub(crate) rpm_429_count: u32,
pub(crate) last_429_at_unix_secs: u64,
pub(crate) last_429_type: String,
pub(crate) status_snapshot: Value,
}
pub(crate) fn project_local_adaptive_rate_limit(
current_key: &StoredProviderCatalogKey,
classification: LocalFailoverClassification,
status_code: u16,
headers: Option<&BTreeMap<String, String>>,
observed_at_unix_secs: u64,
) -> Option<LocalAdaptiveRateLimitProjection> {
if current_key.rpm_limit.is_some() {
return None;
}
if !local_candidate_failure_should_record_adaptive_rate_limit(classification, status_code) {
return None;
}
let latest_upstream_limit = parse_latest_upstream_limit(headers);
Some(LocalAdaptiveRateLimitProjection {
rpm_429_count: current_key
.rpm_429_count
.unwrap_or_default()
.saturating_add(1),
last_429_at_unix_secs: observed_at_unix_secs,
last_429_type: "rpm".to_string(),
status_snapshot: project_local_adaptive_status_snapshot(current_key, latest_upstream_limit),
})
}
fn local_candidate_failure_should_record_adaptive_rate_limit(
classification: LocalFailoverClassification,
status_code: u16,
) -> bool {
status_code == 429
|| matches!(
classification,
LocalFailoverClassification::RetrySemanticRateLimit
)
}
fn project_local_adaptive_status_snapshot(
current_key: &StoredProviderCatalogKey,
latest_upstream_limit: Option<u64>,
) -> Value {
let default_snapshot = default_provider_key_status_snapshot();
let mut snapshot = current_key
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.cloned()
.or_else(|| default_snapshot.as_object().cloned())
.unwrap_or_default();
let observation_count = snapshot
.get("observation_count")
.and_then(Value::as_u64)
.unwrap_or(0)
.saturating_add(1);
snapshot.insert("observation_count".to_string(), json!(observation_count));
if let Some(limit) = latest_upstream_limit {
let header_observation_count = snapshot
.get("header_observation_count")
.and_then(Value::as_u64)
.unwrap_or(0)
.saturating_add(1);
snapshot.insert(
"header_observation_count".to_string(),
json!(header_observation_count),
);
snapshot.insert("latest_upstream_limit".to_string(), json!(limit));
}
let header_observation_count = snapshot
.get("header_observation_count")
.and_then(Value::as_u64)
.unwrap_or(0);
let effective_upstream_limit = latest_upstream_limit.or_else(|| {
snapshot
.get("latest_upstream_limit")
.and_then(Value::as_u64)
});
let learning_confidence = projected_learning_confidence(
observation_count,
header_observation_count,
current_key.learned_rpm_limit.is_some(),
effective_upstream_limit.is_some(),
);
snapshot.insert(
"learning_confidence".to_string(),
json!(learning_confidence),
);
snapshot.insert(
"enforcement_active".to_string(),
json!(adaptive_enforcement_active(
learning_confidence,
current_key.learned_rpm_limit.is_some(),
effective_upstream_limit.is_some(),
)),
);
Value::Object(snapshot)
}
fn projected_learning_confidence(
observation_count: u64,
header_observation_count: u64,
has_learned_limit: bool,
has_upstream_limit: bool,
) -> f64 {
let base = if has_learned_limit || has_upstream_limit {
0.1
} else {
0.0
};
let observation_score = (observation_count.min(8) as f64) * 0.05;
let header_score = (header_observation_count.min(3) as f64) * (0.4 / 3.0);
((base + observation_score + header_score).min(1.0) * 1000.0).round() / 1000.0
}
fn adaptive_enforcement_active(
learning_confidence: f64,
has_learned_limit: bool,
has_upstream_limit: bool,
) -> bool {
(has_learned_limit || has_upstream_limit) && learning_confidence >= 0.5
}
fn parse_latest_upstream_limit(headers: Option<&BTreeMap<String, String>>) -> Option<u64> {
let normalized = headers?
.iter()
.map(|(key, value)| (key.trim().to_ascii_lowercase(), value.trim().to_string()))
.collect::<BTreeMap<_, _>>();
const CANDIDATE_KEYS: &[&str] = &[
"x-ratelimit-limit-requests",
"x-ratelimit-limit-request",
"x-ratelimit-limit",
"x-rate-limit-limit",
"ratelimit-limit",
];
for key in CANDIDATE_KEYS {
if let Some(limit) = normalized
.get(*key)
.and_then(|value| parse_limit_header_value(value))
{
return Some(limit);
}
}
normalized.iter().find_map(|(key, value)| {
if !key.contains("ratelimit") || !key.contains("limit") {
return None;
}
if key.contains("token") {
return None;
}
parse_limit_header_value(value)
})
}
fn parse_limit_header_value(raw: &str) -> Option<u64> {
raw.split([',', ';'])
.find_map(|part| {
let digits = part
.trim()
.chars()
.take_while(|ch| ch.is_ascii_digit())
.collect::<String>();
(!digits.is_empty())
.then(|| digits.parse::<u64>().ok())
.flatten()
})
.filter(|value| *value > 0)
}
#[cfg(test)]
mod tests {
use std::collections::BTreeMap;
use super::project_local_adaptive_rate_limit;
use crate::orchestration::LocalFailoverClassification;
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::json;
fn sample_adaptive_key() -> StoredProviderCatalogKey {
let mut key = StoredProviderCatalogKey::new(
"key-1".to_string(),
"provider-1".to_string(),
"adaptive".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("key should build");
key.rpm_limit = None;
key.learned_rpm_limit = Some(12);
key.rpm_429_count = Some(2);
key
}
#[test]
fn rate_limit_projection_increments_adaptive_rpm_observation() {
let key = sample_adaptive_key();
let projection = project_local_adaptive_rate_limit(
&key,
LocalFailoverClassification::RetrySemanticRateLimit,
429,
None,
1_760_000_000,
)
.expect("projection should exist");
assert_eq!(projection.rpm_429_count, 3);
assert_eq!(projection.last_429_at_unix_secs, 1_760_000_000);
assert_eq!(projection.last_429_type, "rpm");
assert_eq!(projection.status_snapshot["observation_count"], json!(1));
assert_eq!(
projection.status_snapshot["learning_confidence"],
json!(0.15)
);
assert_eq!(
projection.status_snapshot["enforcement_active"],
json!(false)
);
}
#[test]
fn rate_limit_projection_ignores_fixed_limit_keys() {
let mut key = sample_adaptive_key();
key.rpm_limit = Some(20);
assert!(project_local_adaptive_rate_limit(
&key,
LocalFailoverClassification::RetrySemanticRateLimit,
429,
None,
1_760_000_000,
)
.is_none());
}
#[test]
fn rate_limit_projection_ignores_non_rate_limit_failures() {
let key = sample_adaptive_key();
assert!(project_local_adaptive_rate_limit(
&key,
LocalFailoverClassification::RetryUpstreamFailure,
503,
None,
1_760_000_000,
)
.is_none());
}
#[test]
fn rate_limit_projection_records_header_observation_and_limit() {
let mut key = sample_adaptive_key();
key.status_snapshot = Some(json!({
"oauth": { "code": "ok" },
"observation_count": 4,
"header_observation_count": 1,
"latest_upstream_limit": 20
}));
let headers =
BTreeMap::from([("x-ratelimit-limit-requests".to_string(), "60".to_string())]);
let projection = project_local_adaptive_rate_limit(
&key,
LocalFailoverClassification::RetrySemanticRateLimit,
429,
Some(&headers),
1_760_000_000,
)
.expect("projection should exist");
assert_eq!(projection.status_snapshot["observation_count"], json!(5));
assert_eq!(
projection.status_snapshot["header_observation_count"],
json!(2)
);
assert_eq!(
projection.status_snapshot["latest_upstream_limit"],
json!(60)
);
assert_eq!(projection.status_snapshot["oauth"]["code"], json!("ok"));
}
#[test]
fn rate_limit_projection_derives_confidence_and_enforcement_from_evidence() {
let mut key = sample_adaptive_key();
key.status_snapshot = Some(json!({
"observation_count": 7,
"header_observation_count": 2,
"latest_upstream_limit": 24
}));
let headers =
BTreeMap::from([("x-ratelimit-limit-requests".to_string(), "60".to_string())]);
let projection = project_local_adaptive_rate_limit(
&key,
LocalFailoverClassification::RetrySemanticRateLimit,
429,
Some(&headers),
1_760_000_000,
)
.expect("projection should exist");
assert_eq!(
projection.status_snapshot["learning_confidence"],
json!(0.9)
);
assert_eq!(
projection.status_snapshot["enforcement_active"],
json!(true)
);
}
}

234
apps/aether-gateway/src/orchestration/attempt.rs Normal file
View File

@@ -0,0 +1,234 @@
use aether_scheduler_core::parse_request_candidate_report_context;
use serde_json::Value;
use crate::provider_transport::GatewayProviderTransportSnapshot;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct ExecutionAttemptIdentity {
pub(crate) candidate_index: u32,
pub(crate) retry_index: u32,
pub(crate) pool_key_index: Option<u32>,
}
impl ExecutionAttemptIdentity {
pub(crate) const fn new(candidate_index: u32, retry_index: u32) -> Self {
Self {
candidate_index,
retry_index,
pool_key_index: None,
}
}
pub(crate) const fn with_pool_key_index(mut self, pool_key_index: Option<u32>) -> Self {
self.pool_key_index = pool_key_index;
self
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct LocalExecutionCandidateMetadata {
pub(crate) candidate_group_id: Option<String>,
pub(crate) pool_key_index: Option<u32>,
}
pub(crate) fn attempt_identity_from_report_context(
report_context: Option<&Value>,
) -> Option<ExecutionAttemptIdentity> {
let metadata = parse_request_candidate_report_context(report_context)?;
let candidate_metadata = local_execution_candidate_metadata_from_report_context(report_context);
Some(ExecutionAttemptIdentity {
candidate_index: metadata.candidate_index?,
retry_index: metadata.retry_index,
pool_key_index: candidate_metadata.pool_key_index,
})
}
pub(crate) fn local_execution_candidate_metadata_from_report_context(
report_context: Option<&Value>,
) -> LocalExecutionCandidateMetadata {
LocalExecutionCandidateMetadata {
candidate_group_id: report_context
.and_then(Value::as_object)
.and_then(|value| value.get("candidate_group_id"))
.and_then(Value::as_str)
.map(ToOwned::to_owned),
pool_key_index: report_context
.and_then(|value| value.get("pool_key_index"))
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok()),
}
}
pub(crate) fn build_local_attempt_identities(
candidate_index: u32,
transport: &GatewayProviderTransportSnapshot,
) -> Vec<ExecutionAttemptIdentity> {
let attempt_slots = resolve_local_attempt_slot_count(transport);
(0..attempt_slots)
.map(|retry_index| ExecutionAttemptIdentity::new(candidate_index, retry_index))
.collect()
}
fn resolve_local_attempt_slot_count(transport: &GatewayProviderTransportSnapshot) -> u32 {
local_attempt_slots_from_transport(transport).unwrap_or(1)
}
fn local_attempt_slots_from_transport(transport: &GatewayProviderTransportSnapshot) -> Option<u32> {
transport
.provider
.config
.as_ref()
.and_then(|config| config.get("failover_rules"))
.and_then(Value::as_object)
.and_then(|value| value.get("max_retries"))
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok())
.map(|value| value.max(1))
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{
attempt_identity_from_report_context, build_local_attempt_identities,
local_execution_candidate_metadata_from_report_context, ExecutionAttemptIdentity,
LocalExecutionCandidateMetadata,
};
use crate::provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
fn sample_transport(
provider_max_retries: Option<i32>,
endpoint_max_retries: Option<i32>,
provider_config: Option<serde_json::Value>,
) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "OpenAI".to_string(),
provider_type: "llm".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
concurrent_limit: None,
max_retries: provider_max_retries,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: provider_config,
},
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: endpoint_max_retries,
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: "primary".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_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,
},
}
}
#[test]
fn build_local_attempt_identities_defaults_to_single_attempt() {
let identities = build_local_attempt_identities(3, &sample_transport(None, None, None));
assert_eq!(identities, vec![ExecutionAttemptIdentity::new(3, 0)]);
}
#[test]
fn build_local_attempt_identities_prefer_failover_rules_over_endpoint_and_provider() {
let identities = build_local_attempt_identities(
1,
&sample_transport(
Some(5),
Some(4),
Some(json!({
"failover_rules": {
"max_retries": 2
}
})),
),
);
assert_eq!(
identities,
vec![
ExecutionAttemptIdentity::new(1, 0),
ExecutionAttemptIdentity::new(1, 1),
]
);
}
#[test]
fn build_local_attempt_identities_require_explicit_failover_rule_for_expansion() {
let identities =
build_local_attempt_identities(2, &sample_transport(Some(5), Some(3), None));
assert_eq!(identities, vec![ExecutionAttemptIdentity::new(2, 0)]);
}
#[test]
fn parse_attempt_identity_from_report_context_reads_candidate_and_retry_indices() {
let identity = attempt_identity_from_report_context(Some(&json!({
"candidate_index": 4,
"retry_index": 1,
"pool_key_index": 7,
})))
.expect("attempt identity should parse");
assert_eq!(
identity,
ExecutionAttemptIdentity {
candidate_index: 4,
retry_index: 1,
pool_key_index: Some(7),
}
);
}
#[test]
fn parse_candidate_metadata_from_report_context_reads_group_and_pool_metadata() {
let metadata = local_execution_candidate_metadata_from_report_context(Some(&json!({
"candidate_group_id": "group-1",
"pool_key_index": 3,
})));
assert_eq!(
metadata,
LocalExecutionCandidateMetadata {
candidate_group_id: Some("group-1".to_string()),
pool_key_index: Some(3),
}
);
}
}

510
apps/aether-gateway/src/orchestration/classifier.rs Normal file
View File

@@ -0,0 +1,510 @@
use regex::Regex;
use serde_json::Value;
use super::{LocalFailoverPolicy, LocalFailoverRegexRule};
const CLIENT_ERROR_TYPES: &[&str] = &[
"invalid_request_error",
"invalid_argument",
"failed_precondition",
"validation_error",
"bad_request",
];
const CLIENT_ERROR_REASONS: &[&str] = &[
"CONTENT_LENGTH_EXCEEDS_THRESHOLD",
"CONTEXT_LENGTH_EXCEEDED",
"MAX_TOKENS_EXCEEDED",
"INVALID_CONTENT",
"CONTENT_POLICY_VIOLATION",
];
const CLIENT_ERROR_PATTERNS: &[&str] = &[
"could not process image",
"image too large",
"invalid image",
"unsupported image",
"content_policy_violation",
"context_length_exceeded",
"content_length_limit",
"content_length_exceeds",
"invalid_prompt",
"content too long",
"input is too long",
"message is too long",
"prompt is too long",
"image exceeds",
"pdf too large",
"file too large",
"tool_use_id",
"validationexception",
];
const COMPATIBILITY_ERROR_PATTERNS: &[&str] = &[
"unsupported parameter",
"unsupported model",
"unsupported feature",
"not supported with this model",
"model does not support",
"parameter is not supported",
"feature is not supported",
"not available for this model",
];
const THINKING_ERROR_PATTERNS: &[&str] = &[
"invalid `signature` in `thinking` block",
"invalid signature in thinking block",
"thinking.signature: field required",
"thinking.signature:",
"signature verification failed",
"must start with a thinking block",
"expected thinking or redacted_thinking",
"expected `thinking`",
"expected thinking, found",
"expected `thinking`, found",
"expected redacted_thinking, found",
"expected `redacted_thinking`, found",
"thoughtsignature",
"thought_signature",
];
const RETRYABLE_RATE_LIMIT_PATTERNS: &[&str] = &[
"rate_limit",
"rate limited",
"resource_exhausted",
"throttl",
"too many requests",
"quota reached",
"quota exceeded",
"quota hit",
];
#[derive(Debug, Clone, PartialEq, Eq, Default)]
struct ParsedLocalErrorResponse {
type_name: Option<String>,
message: Option<String>,
reason: Option<String>,
raw: Option<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct LocalFailoverInput<'a> {
pub(crate) status_code: u16,
pub(crate) response_text: Option<&'a str>,
}
impl<'a> LocalFailoverInput<'a> {
pub(crate) fn new(status_code: u16, response_text: Option<&'a str>) -> Self {
Self {
status_code,
response_text: response_text
.map(str::trim)
.filter(|value| !value.is_empty()),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum LocalFailoverClassification {
UseDefault,
StopStatusCode,
StopErrorPattern,
StopSemanticClientError,
RetrySuccessPattern,
RetrySemanticCompatibilityError,
RetrySemanticRateLimit,
RetrySemanticThinkingError,
RetryStatusCode,
RetryUpstreamFailure,
}
pub(crate) fn classify_local_failover(
policy: &LocalFailoverPolicy,
input: LocalFailoverInput<'_>,
) -> LocalFailoverClassification {
if policy.stop_status_codes.contains(&input.status_code) {
return LocalFailoverClassification::StopStatusCode;
}
if input.status_code >= 400
&& input.response_text.is_some_and(|text| {
policy
.error_stop_patterns
.iter()
.any(|rule| local_failover_regex_rule_matches(rule, text, input.status_code))
})
{
return LocalFailoverClassification::StopErrorPattern;
}
if input.status_code == 200
&& input.response_text.is_some_and(|text| {
policy
.success_failover_patterns
.iter()
.any(|rule| local_failover_regex_rule_matches(rule, text, input.status_code))
})
{
return LocalFailoverClassification::RetrySuccessPattern;
}
let parsed_error = parse_local_error_response(input.response_text);
if is_semantic_thinking_error(input.status_code, &parsed_error) {
return LocalFailoverClassification::RetrySemanticThinkingError;
}
if is_semantic_compatibility_error(input.status_code, &parsed_error) {
return LocalFailoverClassification::RetrySemanticCompatibilityError;
}
if is_semantic_rate_limit_error(input.status_code, &parsed_error) {
return LocalFailoverClassification::RetrySemanticRateLimit;
}
if is_semantic_client_error(input.status_code, &parsed_error) {
return LocalFailoverClassification::StopSemanticClientError;
}
if policy.continue_status_codes.contains(&input.status_code) {
return LocalFailoverClassification::RetryStatusCode;
}
if should_failover_local_upstream_status(input.status_code) {
return LocalFailoverClassification::RetryUpstreamFailure;
}
LocalFailoverClassification::UseDefault
}
pub(crate) fn local_failover_error_message(response_text: Option<&str>) -> Option<String> {
let parsed = parse_local_error_response(response_text);
parsed
.message
.or(parsed.reason)
.or(parsed.raw)
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
fn should_failover_local_upstream_status(status_code: u16) -> bool {
status_code >= 400
}
fn parse_local_error_response(response_text: Option<&str>) -> ParsedLocalErrorResponse {
let raw = response_text
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned);
let Some(raw_text) = raw.clone() else {
return ParsedLocalErrorResponse::default();
};
let mut parsed = ParsedLocalErrorResponse {
raw: Some(raw_text.clone()),
..ParsedLocalErrorResponse::default()
};
let Ok(value) = serde_json::from_str::<Value>(&raw_text) else {
parsed.message = Some(raw_text);
return parsed;
};
let body_object = value.as_object();
let error_object = body_object
.and_then(|object| object.get("error"))
.and_then(Value::as_object);
parsed.type_name = first_non_empty_json_text(error_object, &["type", "__type"])
.or_else(|| first_non_empty_json_text(body_object, &["type", "__type"]));
parsed.message = first_non_empty_json_text(error_object, &["message", "detail", "reason"])
.or_else(|| first_non_empty_json_text(body_object, &["errorMessage"]))
.or_else(|| {
body_object
.and_then(|object| object.get("error"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
})
.or_else(|| first_non_empty_json_text(body_object, &["message", "detail", "reason"]));
parsed.reason = first_non_empty_json_text(error_object, &["reason", "code", "status"])
.or_else(|| first_non_empty_json_text(body_object, &["reason", "code", "status"]));
let Some(message) = parsed.message.clone() else {
return parsed;
};
if !message.starts_with('{') {
return parsed;
}
let Ok(nested) = serde_json::from_str::<Value>(&message) else {
return parsed;
};
let nested_object = nested.as_object();
let nested_error_object = nested_object
.and_then(|object| object.get("error"))
.and_then(Value::as_object);
parsed.type_name = parsed
.type_name
.or_else(|| first_non_empty_json_text(nested_error_object, &["type", "__type"]))
.or_else(|| first_non_empty_json_text(nested_object, &["type", "__type"]));
parsed.message =
first_non_empty_json_text(nested_error_object, &["message", "detail", "reason"])
.or_else(|| first_non_empty_json_text(nested_object, &["message", "detail", "reason"]))
.or(parsed.message);
parsed.reason = parsed
.reason
.or_else(|| first_non_empty_json_text(nested_error_object, &["reason", "code", "status"]))
.or_else(|| first_non_empty_json_text(nested_object, &["reason", "code", "status"]));
parsed
}
fn first_non_empty_json_text(
object: Option<&serde_json::Map<String, Value>>,
keys: &[&str],
) -> Option<String> {
let object = object?;
for key in keys {
let Some(value) = object.get(*key) else {
continue;
};
match value {
Value::String(text) if !text.trim().is_empty() => return Some(text.trim().to_string()),
Value::Number(number) => return Some(number.to_string()),
_ => {}
}
}
None
}
fn semantic_search_text(parsed: &ParsedLocalErrorResponse) -> String {
[
parsed.type_name.as_deref(),
parsed.reason.as_deref(),
parsed.message.as_deref(),
parsed.raw.as_deref(),
]
.into_iter()
.flatten()
.map(str::trim)
.filter(|value| !value.is_empty())
.map(str::to_ascii_lowercase)
.collect::<Vec<_>>()
.join(" ")
}
fn is_semantic_client_error(status_code: u16, parsed: &ParsedLocalErrorResponse) -> bool {
if status_code < 400 {
return false;
}
if parsed.type_name.as_deref().is_some_and(|type_name| {
let type_name = type_name.to_ascii_lowercase();
CLIENT_ERROR_TYPES
.iter()
.any(|pattern| type_name.contains(pattern))
}) {
return true;
}
if parsed.reason.as_deref().is_some_and(|reason| {
let reason = reason.to_ascii_uppercase();
CLIENT_ERROR_REASONS
.iter()
.any(|pattern| reason.contains(pattern))
}) {
return true;
}
let search_text = semantic_search_text(parsed);
!search_text.is_empty()
&& CLIENT_ERROR_PATTERNS
.iter()
.any(|pattern| search_text.contains(&pattern.to_ascii_lowercase()))
}
fn is_semantic_compatibility_error(status_code: u16, parsed: &ParsedLocalErrorResponse) -> bool {
if status_code < 400 {
return false;
}
let search_text = semantic_search_text(parsed);
!search_text.is_empty()
&& COMPATIBILITY_ERROR_PATTERNS
.iter()
.any(|pattern| search_text.contains(&pattern.to_ascii_lowercase()))
}
fn is_semantic_thinking_error(status_code: u16, parsed: &ParsedLocalErrorResponse) -> bool {
if status_code != 400 {
return false;
}
let search_text = semantic_search_text(parsed);
!search_text.is_empty()
&& THINKING_ERROR_PATTERNS
.iter()
.any(|pattern| search_text.contains(&pattern.to_ascii_lowercase()))
}
fn is_semantic_rate_limit_error(status_code: u16, parsed: &ParsedLocalErrorResponse) -> bool {
if status_code < 400 {
return false;
}
let search_text = semantic_search_text(parsed);
!search_text.is_empty()
&& RETRYABLE_RATE_LIMIT_PATTERNS
.iter()
.any(|pattern| search_text.contains(&pattern.to_ascii_lowercase()))
}
fn local_failover_regex_rule_matches(
rule: &LocalFailoverRegexRule,
response_text: &str,
status_code: u16,
) -> bool {
if !rule.status_codes.is_empty() && !rule.status_codes.contains(&status_code) {
return false;
}
Regex::new(&rule.pattern)
.ok()
.is_some_and(|regex| regex.is_match(response_text))
}
#[cfg(test)]
mod tests {
use std::collections::BTreeSet;
use super::{classify_local_failover, LocalFailoverClassification, LocalFailoverInput};
use crate::orchestration::{LocalFailoverPolicy, LocalFailoverRegexRule};
#[test]
fn classifier_honors_explicit_stop_before_default_retryable_status() {
let policy = LocalFailoverPolicy {
stop_status_codes: [503].into_iter().collect(),
..LocalFailoverPolicy::default()
};
assert_eq!(
classify_local_failover(&policy, LocalFailoverInput::new(503, None)),
LocalFailoverClassification::StopStatusCode
);
}
#[test]
fn classifier_detects_success_failover_pattern() {
let policy = LocalFailoverPolicy {
success_failover_patterns: vec![LocalFailoverRegexRule {
pattern: "relay:.*格式错误".to_string(),
status_codes: BTreeSet::new(),
}],
..LocalFailoverPolicy::default()
};
assert_eq!(
classify_local_failover(
&policy,
LocalFailoverInput::new(200, Some("{\"error\":\"relay: 返回格式错误\"}"))
),
LocalFailoverClassification::RetrySuccessPattern
);
}
#[test]
fn classifier_detects_error_stop_pattern() {
let policy = LocalFailoverPolicy {
error_stop_patterns: vec![LocalFailoverRegexRule {
pattern: "content_policy_violation".to_string(),
status_codes: [400, 403].into_iter().collect(),
}],
..LocalFailoverPolicy::default()
};
assert_eq!(
classify_local_failover(
&policy,
LocalFailoverInput::new(400, Some("{\"error\":\"content_policy_violation\"}"))
),
LocalFailoverClassification::StopErrorPattern
);
}
#[test]
fn classifier_stops_semantic_client_errors_without_custom_rule() {
assert_eq!(
classify_local_failover(
&LocalFailoverPolicy::default(),
LocalFailoverInput::new(
400,
Some(
"{\"error\":{\"type\":\"invalid_request_error\",\"message\":\"prompt is too long\"}}"
)
)
),
LocalFailoverClassification::StopSemanticClientError
);
}
#[test]
fn classifier_retries_semantic_compatibility_errors() {
assert_eq!(
classify_local_failover(
&LocalFailoverPolicy::default(),
LocalFailoverInput::new(
400,
Some("{\"error\":{\"message\":\"Unsupported parameter: max_tokens is not supported with this model\"}}")
)
),
LocalFailoverClassification::RetrySemanticCompatibilityError
);
}
#[test]
fn classifier_retries_semantic_thinking_errors() {
assert_eq!(
classify_local_failover(
&LocalFailoverPolicy::default(),
LocalFailoverInput::new(
400,
Some(
"{\"error\":{\"message\":\"invalid `signature` in `thinking` block: signature is for a different request\"}}"
)
)
),
LocalFailoverClassification::RetrySemanticThinkingError
);
}
#[test]
fn classifier_retries_semantic_rate_limit_errors_even_when_status_is_not_429() {
assert_eq!(
classify_local_failover(
&LocalFailoverPolicy::default(),
LocalFailoverInput::new(
400,
Some("{\"error\":{\"message\":\"resource_exhausted: quota reached\"}}")
)
),
LocalFailoverClassification::RetrySemanticRateLimit
);
}
#[test]
fn classifier_keeps_embedded_rate_limit_error_in_success_response_on_default_path() {
assert_eq!(
classify_local_failover(
&LocalFailoverPolicy::default(),
LocalFailoverInput::new(
200,
Some(
"{\"error\":{\"message\":\"quota reached\",\"type\":\"rate_limit_error\"}}"
)
)
),
LocalFailoverClassification::UseDefault
);
}
}

1215
apps/aether-gateway/src/orchestration/effects.rs Normal file
View File

File diff suppressed because it is too large Load Diff

188
apps/aether-gateway/src/orchestration/health.rs Normal file
View File

@@ -0,0 +1,188 @@
use serde_json::{json, Value};
use super::LocalFailoverClassification;
use crate::handlers::shared::unix_secs_to_rfc3339;
const LOCAL_HEALTH_SCORE_FLOOR: f64 = 0.2;
pub(crate) fn project_local_failure_health(
current_health_by_format: Option<&Value>,
api_format: &str,
classification: LocalFailoverClassification,
status_code: u16,
observed_at_unix_secs: u64,
) -> Option<Value> {
if !local_candidate_failure_should_project_health(classification, status_code) {
return None;
}
let api_format = api_format.trim();
if api_format.is_empty() {
return None;
}
let mut health_by_format = current_health_by_format
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
let current = health_by_format
.get(api_format)
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
let previous_failures = current
.get("consecutive_failures")
.and_then(Value::as_i64)
.unwrap_or(0)
.max(0) as u64;
let consecutive_failures = previous_failures.saturating_add(1);
health_by_format.insert(
api_format.to_string(),
json!({
"health_score": projected_failure_health_score(classification, status_code, consecutive_failures),
"consecutive_failures": consecutive_failures,
"last_failure_at": unix_secs_to_rfc3339(observed_at_unix_secs),
}),
);
Some(Value::Object(health_by_format))
}
pub(crate) fn project_local_success_health(
current_health_by_format: Option<&Value>,
api_format: &str,
) -> Option<Value> {
let api_format = api_format.trim();
if api_format.is_empty() {
return None;
}
let mut health_by_format = current_health_by_format
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
health_by_format.insert(
api_format.to_string(),
json!({
"health_score": 1.0,
"consecutive_failures": 0,
"last_failure_at": Value::Null,
}),
);
Some(Value::Object(health_by_format))
}
fn local_candidate_failure_should_project_health(
classification: LocalFailoverClassification,
status_code: u16,
) -> bool {
if status_code < 400 {
return false;
}
match classification {
LocalFailoverClassification::RetrySuccessPattern
| LocalFailoverClassification::RetrySemanticCompatibilityError
| LocalFailoverClassification::RetrySemanticRateLimit
| LocalFailoverClassification::RetrySemanticThinkingError
| LocalFailoverClassification::RetryStatusCode
| LocalFailoverClassification::RetryUpstreamFailure => true,
LocalFailoverClassification::UseDefault | LocalFailoverClassification::StopStatusCode => {
status_code >= 500
}
LocalFailoverClassification::StopErrorPattern
| LocalFailoverClassification::StopSemanticClientError => false,
}
}
fn projected_failure_health_score(
classification: LocalFailoverClassification,
status_code: u16,
consecutive_failures: u64,
) -> f64 {
let base_score = match classification {
LocalFailoverClassification::RetrySemanticRateLimit => 0.7,
LocalFailoverClassification::RetrySemanticCompatibilityError
| LocalFailoverClassification::RetrySemanticThinkingError => 0.8,
LocalFailoverClassification::RetrySuccessPattern => 0.75,
_ if status_code >= 500 => 0.6,
_ => 0.7,
};
let penalty = consecutive_failures.saturating_sub(1) as f64 * 0.15;
let normalized = (base_score - penalty).max(LOCAL_HEALTH_SCORE_FLOOR);
(normalized * 1000.0).round() / 1000.0
}
#[cfg(test)]
mod tests {
use serde_json::{json, Value};
use super::{project_local_failure_health, project_local_success_health};
use crate::orchestration::LocalFailoverClassification;
#[test]
fn failure_projection_tracks_consecutive_failures_and_degrades_score() {
let projected = project_local_failure_health(
Some(&json!({
"openai:chat": {
"health_score": 0.7,
"consecutive_failures": 1,
"last_failure_at": "2026-01-01T00:00:00+00:00"
}
})),
"openai:chat",
LocalFailoverClassification::RetryUpstreamFailure,
503,
1_760_000_000,
)
.expect("projection should exist");
assert_eq!(projected["openai:chat"]["consecutive_failures"], json!(2));
assert_eq!(projected["openai:chat"]["health_score"], json!(0.45));
assert!(projected["openai:chat"]["last_failure_at"].is_string());
}
#[test]
fn failure_projection_ignores_semantic_client_error() {
assert!(project_local_failure_health(
None,
"openai:chat",
LocalFailoverClassification::StopSemanticClientError,
400,
1_760_000_000,
)
.is_none());
}
#[test]
fn success_projection_resets_only_target_format() {
let projected = project_local_success_health(
Some(&json!({
"openai:chat": {
"health_score": 0.4,
"consecutive_failures": 3,
"last_failure_at": "2026-01-01T00:00:00+00:00"
},
"openai:responses": {
"health_score": 0.8,
"consecutive_failures": 1,
"last_failure_at": "2026-01-02T00:00:00+00:00"
}
})),
"openai:chat",
)
.expect("projection should exist");
assert_eq!(
projected["openai:chat"],
json!({
"health_score": 1.0,
"consecutive_failures": 0,
"last_failure_at": Value::Null,
})
);
assert_eq!(projected["openai:responses"]["health_score"], json!(0.8));
}
}

78
apps/aether-gateway/src/orchestration/mod.rs Normal file
View File

@@ -0,0 +1,78 @@
use aether_contracts::ExecutionPlan;
use crate::AppState;
mod adaptive;
mod attempt;
mod classifier;
mod effects;
mod health;
mod policy;
mod recovery;
mod report_effects;
pub(crate) use self::adaptive::{
project_local_adaptive_rate_limit, LocalAdaptiveRateLimitProjection,
};
pub(crate) use self::attempt::{
attempt_identity_from_report_context, build_local_attempt_identities,
local_execution_candidate_metadata_from_report_context, ExecutionAttemptIdentity,
LocalExecutionCandidateMetadata,
};
pub(crate) use self::classifier::{
classify_local_failover, local_failover_error_message, LocalFailoverClassification,
LocalFailoverInput,
};
pub(crate) use self::effects::{
apply_local_execution_effect, LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect,
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
};
pub(crate) use self::health::{project_local_failure_health, project_local_success_health};
pub(crate) use self::policy::{
append_local_failover_policy_to_value, local_failover_policy_from_report_context,
local_failover_policy_from_transport, resolve_local_failover_policy, LocalFailoverPolicy,
LocalFailoverRegexRule,
};
pub(crate) use self::recovery::{
analyze_local_failover, recover_local_failover_decision, LocalFailoverAnalysis,
LocalFailoverDecision,
};
#[cfg(test)]
pub(crate) use self::report_effects::clear_local_report_effect_caches_for_tests;
pub(crate) use self::report_effects::{
apply_local_report_effect, store_local_gemini_file_mapping, LocalReportEffect,
};
pub(crate) async fn resolve_local_failover_analysis_for_attempt(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
status_code: u16,
response_text: Option<&str>,
) -> LocalFailoverAnalysis {
if attempt_identity_from_report_context(report_context).is_none() {
return LocalFailoverAnalysis::use_default();
}
let policy = resolve_local_failover_policy(state, plan, report_context).await;
analyze_local_failover(&policy, LocalFailoverInput::new(status_code, response_text))
}
pub(crate) async fn resolve_local_failover_decision_for_attempt(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
status_code: u16,
response_text: Option<&str>,
) -> LocalFailoverDecision {
resolve_local_failover_analysis_for_attempt(
state,
plan,
report_context,
status_code,
response_text,
)
.await
.decision
}

358
apps/aether-gateway/src/orchestration/policy.rs Normal file
View File

@@ -0,0 +1,358 @@
use std::collections::BTreeSet;
use aether_contracts::ExecutionPlan;
use serde_json::{json, Value};
use tracing::debug;
use crate::provider_transport::GatewayProviderTransportSnapshot;
use crate::AppState;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub(crate) struct LocalFailoverPolicy {
pub(crate) max_retries: Option<u64>,
pub(crate) stop_status_codes: BTreeSet<u16>,
pub(crate) continue_status_codes: BTreeSet<u16>,
pub(crate) success_failover_patterns: Vec<LocalFailoverRegexRule>,
pub(crate) error_stop_patterns: Vec<LocalFailoverRegexRule>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct LocalFailoverRegexRule {
pub(crate) pattern: String,
pub(crate) status_codes: BTreeSet<u16>,
}
pub(crate) async fn resolve_local_failover_policy(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
) -> LocalFailoverPolicy {
if let Some(policy) = local_failover_policy_from_report_context(report_context) {
debug!(
event_name = "local_failover_policy_loaded",
log_type = "debug",
request_id = %plan.request_id,
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
source = "report_context",
max_retries = ?policy.max_retries,
stop_status_code_count = policy.stop_status_codes.len(),
continue_status_code_count = policy.continue_status_codes.len(),
success_failover_pattern_count = policy.success_failover_patterns.len(),
error_stop_pattern_count = policy.error_stop_patterns.len(),
"gateway loaded local failover policy from report context"
);
return policy;
}
let transport = match state
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
.await
{
Ok(Some(transport)) => transport,
Ok(None) | Err(_) => return LocalFailoverPolicy::default(),
};
let policy = local_failover_policy_from_transport(&transport);
debug!(
event_name = "local_failover_policy_loaded",
log_type = "debug",
request_id = %plan.request_id,
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
source = "transport_snapshot",
max_retries = ?policy.max_retries,
stop_status_code_count = policy.stop_status_codes.len(),
continue_status_code_count = policy.continue_status_codes.len(),
success_failover_pattern_count = policy.success_failover_patterns.len(),
error_stop_pattern_count = policy.error_stop_patterns.len(),
"gateway loaded local failover policy from transport snapshot"
);
policy
}
pub(crate) fn local_failover_policy_from_transport(
transport: &GatewayProviderTransportSnapshot,
) -> LocalFailoverPolicy {
let rules = transport
.provider
.config
.as_ref()
.and_then(|config| config.get("failover_rules"))
.and_then(Value::as_object);
let max_retries = rules
.and_then(|value| value.get("max_retries"))
.and_then(parse_u64_value)
.or_else(|| {
transport
.endpoint
.max_retries
.and_then(|value| u64::try_from(value).ok())
})
.or_else(|| {
transport
.provider
.max_retries
.and_then(|value| u64::try_from(value).ok())
});
LocalFailoverPolicy {
max_retries,
stop_status_codes: rules
.map(|value| {
parse_status_code_set(
value,
&[
"stop_on_status_codes",
"early_stop_status_codes",
"non_retryable_status_codes",
"stop_status_codes",
],
)
})
.unwrap_or_default(),
continue_status_codes: rules
.map(|value| {
parse_status_code_set(
value,
&[
"continue_on_status_codes",
"retryable_status_codes",
"retry_on_status_codes",
"continue_status_codes",
],
)
})
.unwrap_or_default(),
success_failover_patterns: rules
.map(|value| parse_regex_rules(value, "success_failover_patterns"))
.unwrap_or_default(),
error_stop_patterns: rules
.map(|value| parse_regex_rules(value, "error_stop_patterns"))
.unwrap_or_default(),
}
}
pub(crate) fn local_failover_policy_from_report_context(
report_context: Option<&Value>,
) -> Option<LocalFailoverPolicy> {
let object = report_context
.and_then(Value::as_object)?
.get("local_failover_policy")?
.as_object()?;
Some(LocalFailoverPolicy {
max_retries: object.get("max_retries").and_then(parse_u64_value),
stop_status_codes: object
.get("stop_status_codes")
.map(parse_status_code_list)
.unwrap_or_default(),
continue_status_codes: object
.get("continue_status_codes")
.map(parse_status_code_list)
.unwrap_or_default(),
success_failover_patterns: parse_regex_rules(object, "success_failover_patterns"),
error_stop_patterns: parse_regex_rules(object, "error_stop_patterns"),
})
}
pub(crate) fn append_local_failover_policy_to_value(
value: Value,
transport: &GatewayProviderTransportSnapshot,
) -> Value {
let Value::Object(mut object) = value else {
return value;
};
object.insert(
"local_failover_policy".to_string(),
local_failover_policy_to_value(&local_failover_policy_from_transport(transport)),
);
Value::Object(object)
}
fn parse_status_code_list(value: &Value) -> BTreeSet<u16> {
value
.as_array()
.into_iter()
.flat_map(|values| values.iter())
.filter_map(|value| parse_u64_value(value).and_then(|value| u16::try_from(value).ok()))
.collect()
}
fn local_failover_policy_to_value(policy: &LocalFailoverPolicy) -> Value {
json!({
"max_retries": policy.max_retries,
"stop_status_codes": policy.stop_status_codes.iter().copied().collect::<Vec<_>>(),
"continue_status_codes": policy.continue_status_codes.iter().copied().collect::<Vec<_>>(),
"success_failover_patterns": policy.success_failover_patterns.iter().map(local_failover_regex_rule_to_value).collect::<Vec<_>>(),
"error_stop_patterns": policy.error_stop_patterns.iter().map(local_failover_regex_rule_to_value).collect::<Vec<_>>(),
})
}
fn local_failover_regex_rule_to_value(rule: &LocalFailoverRegexRule) -> Value {
json!({
"pattern": rule.pattern,
"status_codes": rule.status_codes.iter().copied().collect::<Vec<_>>(),
})
}
fn parse_regex_rules(
rules: &serde_json::Map<String, serde_json::Value>,
key: &str,
) -> Vec<LocalFailoverRegexRule> {
rules
.get(key)
.and_then(Value::as_array)
.into_iter()
.flat_map(|items| items.iter())
.filter_map(parse_regex_rule)
.collect()
}
fn parse_regex_rule(value: &serde_json::Value) -> Option<LocalFailoverRegexRule> {
let object = value.as_object()?;
let pattern = object
.get("pattern")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
Some(LocalFailoverRegexRule {
pattern: pattern.to_string(),
status_codes: object
.get("status_codes")
.and_then(Value::as_array)
.into_iter()
.flat_map(|values| values.iter())
.filter_map(|value| parse_u64_value(value).and_then(|value| u16::try_from(value).ok()))
.collect(),
})
}
fn parse_status_code_set(
rules: &serde_json::Map<String, serde_json::Value>,
keys: &[&str],
) -> BTreeSet<u16> {
keys.iter()
.filter_map(|key| rules.get(*key))
.filter_map(Value::as_array)
.flat_map(|values| values.iter())
.filter_map(|value| parse_u64_value(value).and_then(|value| u16::try_from(value).ok()))
.collect()
}
fn parse_u64_value(value: &serde_json::Value) -> Option<u64> {
value
.as_u64()
.or_else(|| value.as_i64().and_then(|value| u64::try_from(value).ok()))
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::{
append_local_failover_policy_to_value, local_failover_policy_from_report_context,
LocalFailoverPolicy, LocalFailoverRegexRule,
};
use crate::provider_transport::snapshot::{
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
fn sample_transport(
provider_max_retries: Option<i32>,
endpoint_max_retries: Option<i32>,
provider_config: Option<serde_json::Value>,
) -> GatewayProviderTransportSnapshot {
GatewayProviderTransportSnapshot {
provider: GatewayProviderTransportProvider {
id: "provider-1".to_string(),
name: "OpenAI".to_string(),
provider_type: "llm".to_string(),
website: None,
is_active: true,
keep_priority_on_conversion: false,
enable_format_conversion: true,
concurrent_limit: None,
max_retries: provider_max_retries,
proxy: None,
request_timeout_secs: None,
stream_first_byte_timeout_secs: None,
config: provider_config,
},
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: endpoint_max_retries,
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: "primary".to_string(),
auth_type: "bearer".to_string(),
is_active: true,
api_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,
},
}
}
#[test]
fn append_local_failover_policy_to_value_round_trips_policy_shape() {
let report_context = append_local_failover_policy_to_value(
json!({
"request_id": "req-1",
}),
&sample_transport(
Some(5),
Some(4),
Some(json!({
"failover_rules": {
"max_retries": 2,
"continue_status_codes": [429],
"stop_status_codes": [400],
"success_failover_patterns": [{"pattern": "quota", "status_codes": [200]}],
"error_stop_patterns": [{"pattern": "validation", "status_codes": [422]}]
}
})),
),
);
assert_eq!(
local_failover_policy_from_report_context(Some(&report_context)),
Some(LocalFailoverPolicy {
max_retries: Some(2),
stop_status_codes: [400].into_iter().collect(),
continue_status_codes: [429].into_iter().collect(),
success_failover_patterns: vec![LocalFailoverRegexRule {
pattern: "quota".to_string(),
status_codes: [200].into_iter().collect(),
}],
error_stop_patterns: vec![LocalFailoverRegexRule {
pattern: "validation".to_string(),
status_codes: [422].into_iter().collect(),
}],
})
);
}
}

150
apps/aether-gateway/src/orchestration/recovery.rs Normal file
View File

@@ -0,0 +1,150 @@
use super::classifier::{classify_local_failover, LocalFailoverClassification, LocalFailoverInput};
use super::LocalFailoverPolicy;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum LocalFailoverDecision {
UseDefault,
RetryNextCandidate,
StopLocalFailover,
}
impl LocalFailoverDecision {
pub(crate) const fn as_str(self) -> &'static str {
match self {
Self::UseDefault => "use_default",
Self::RetryNextCandidate => "retry_next_candidate",
Self::StopLocalFailover => "stop_local_failover",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct LocalFailoverAnalysis {
pub(crate) classification: LocalFailoverClassification,
pub(crate) decision: LocalFailoverDecision,
}
impl LocalFailoverAnalysis {
pub(crate) const fn use_default() -> Self {
Self {
classification: LocalFailoverClassification::UseDefault,
decision: LocalFailoverDecision::UseDefault,
}
}
}
pub(crate) fn analyze_local_failover(
policy: &LocalFailoverPolicy,
input: LocalFailoverInput<'_>,
) -> LocalFailoverAnalysis {
let classification = classify_local_failover(policy, input);
LocalFailoverAnalysis {
classification,
decision: decision_from_classification(classification),
}
}
pub(crate) fn recover_local_failover_decision(
policy: &LocalFailoverPolicy,
input: LocalFailoverInput<'_>,
) -> LocalFailoverDecision {
analyze_local_failover(policy, input).decision
}
const fn decision_from_classification(
classification: LocalFailoverClassification,
) -> LocalFailoverDecision {
match classification {
LocalFailoverClassification::UseDefault => LocalFailoverDecision::UseDefault,
LocalFailoverClassification::StopStatusCode
| LocalFailoverClassification::StopErrorPattern
| LocalFailoverClassification::StopSemanticClientError => {
LocalFailoverDecision::StopLocalFailover
}
LocalFailoverClassification::RetrySuccessPattern
| LocalFailoverClassification::RetrySemanticCompatibilityError
| LocalFailoverClassification::RetrySemanticRateLimit
| LocalFailoverClassification::RetrySemanticThinkingError
| LocalFailoverClassification::RetryStatusCode
| LocalFailoverClassification::RetryUpstreamFailure => {
LocalFailoverDecision::RetryNextCandidate
}
}
}
#[cfg(test)]
mod tests {
use super::{analyze_local_failover, recover_local_failover_decision, LocalFailoverDecision};
use crate::orchestration::{
LocalFailoverClassification, LocalFailoverInput, LocalFailoverPolicy,
};
#[test]
fn recovery_maps_retryable_status_to_retry_next_candidate() {
let policy = LocalFailoverPolicy {
continue_status_codes: [429].into_iter().collect(),
..LocalFailoverPolicy::default()
};
assert_eq!(
recover_local_failover_decision(&policy, LocalFailoverInput::new(429, None)),
LocalFailoverDecision::RetryNextCandidate
);
}
#[test]
fn recovery_maps_neutral_status_to_use_default() {
assert_eq!(
recover_local_failover_decision(
&LocalFailoverPolicy::default(),
LocalFailoverInput::new(200, None)
),
LocalFailoverDecision::UseDefault
);
}
#[test]
fn recovery_maps_semantic_client_error_to_stop_failover() {
assert_eq!(
recover_local_failover_decision(
&LocalFailoverPolicy::default(),
LocalFailoverInput::new(
400,
Some("{\"error\":{\"type\":\"invalid_request_error\",\"message\":\"prompt is too long\"}}")
)
),
LocalFailoverDecision::StopLocalFailover
);
}
#[test]
fn recovery_maps_semantic_thinking_error_to_retry_next_candidate() {
assert_eq!(
recover_local_failover_decision(
&LocalFailoverPolicy::default(),
LocalFailoverInput::new(
400,
Some("{\"error\":{\"message\":\"invalid `signature` in `thinking` block\"}}")
)
),
LocalFailoverDecision::RetryNextCandidate
);
}
#[test]
fn analysis_keeps_classification_and_decision_together() {
let analysis = analyze_local_failover(
&LocalFailoverPolicy::default(),
LocalFailoverInput::new(
400,
Some("{\"error\":{\"message\":\"Unsupported parameter: stream_options is not supported with this model\"}}"),
),
);
assert_eq!(analysis.decision, LocalFailoverDecision::RetryNextCandidate);
assert_eq!(
analysis.classification,
LocalFailoverClassification::RetrySemanticCompatibilityError
);
}
}

424
apps/aether-gateway/src/orchestration/report_effects.rs Normal file
View File

@@ -0,0 +1,424 @@
use std::collections::{BTreeMap, HashMap};
use std::sync::{Mutex, OnceLock};
use std::time::{Duration, Instant};
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_usage_runtime::{
extract_gemini_file_mapping_entries, gemini_file_mapping_cache_key, normalize_gemini_file_name,
report_request_id, GatewayStreamReportRequest, GatewaySyncReportRequest,
GEMINI_FILE_MAPPING_TTL_SECONDS,
};
use serde_json::Value;
use tracing::warn;
use uuid::Uuid;
use crate::clock::current_unix_secs;
use crate::handlers::shared::sync_provider_key_quota_status_snapshot;
use crate::log_ids::short_request_id;
use crate::{AppState, GatewayError};
const CODEX_QUOTA_CACHE_TTL_SECONDS: u64 = 30;
const CODEX_QUOTA_CACHE_MAX_ENTRIES: usize = 4096;
type HeaderFingerprintCache = Mutex<HashMap<String, (String, Instant)>>;
static CODEX_QUOTA_HEADER_FINGERPRINT_CACHE: OnceLock<HeaderFingerprintCache> = OnceLock::new();
#[derive(Debug, Clone, Copy)]
pub(crate) enum LocalReportEffect<'a> {
Sync {
payload: &'a GatewaySyncReportRequest,
},
Stream {
payload: &'a GatewayStreamReportRequest,
},
}
pub(crate) async fn apply_local_report_effect(state: &AppState, effect: LocalReportEffect<'_>) {
match effect {
LocalReportEffect::Sync { payload } => {
apply_local_sync_report_effect(state, payload).await;
}
LocalReportEffect::Stream { payload } => {
apply_local_stream_report_effect(state, payload).await;
}
}
}
fn codex_quota_header_fingerprint_cache() -> &'static HeaderFingerprintCache {
CODEX_QUOTA_HEADER_FINGERPRINT_CACHE.get_or_init(|| Mutex::new(HashMap::new()))
}
fn report_context_key_id(report_context: Option<&Value>) -> Option<String> {
report_context
.and_then(|context| context.get("key_id"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn is_volatile_compare_field(key: &str) -> bool {
key == "updated_at" || key.ends_with("_reset_seconds") || key.ends_with("_reset_after_seconds")
}
fn canonicalize_value(value: &Value) -> Value {
match value {
Value::Array(items) => Value::Array(items.iter().map(canonicalize_value).collect()),
Value::Object(object) => {
let mut entries = object.iter().collect::<Vec<_>>();
entries.sort_by(|left, right| left.0.cmp(right.0));
let mut normalized = serde_json::Map::new();
for (key, value) in entries {
normalized.insert(key.clone(), canonicalize_value(value));
}
Value::Object(normalized)
}
_ => value.clone(),
}
}
fn fingerprint_codex_payload(value: &Value) -> Option<String> {
let object = value.as_object()?;
let mut entries = object
.iter()
.filter(|(key, _)| !is_volatile_compare_field(key))
.collect::<Vec<_>>();
entries.sort_by(|left, right| left.0.cmp(right.0));
let mut normalized = serde_json::Map::new();
for (key, value) in entries {
normalized.insert(key.clone(), canonicalize_value(value));
}
serde_json::to_string(&Value::Object(normalized)).ok()
}
fn get_cached_codex_quota_fingerprint(key_id: &str, now: Instant) -> Option<String> {
let mut cache = codex_quota_header_fingerprint_cache()
.lock()
.expect("codex realtime quota cache should lock");
match cache.get(key_id) {
Some((fingerprint, expires_at)) if *expires_at > now => Some(fingerprint.clone()),
Some(_) => {
cache.remove(key_id);
None
}
None => None,
}
}
fn set_cached_codex_quota_fingerprint(key_id: &str, fingerprint: String, now: Instant) {
let mut cache = codex_quota_header_fingerprint_cache()
.lock()
.expect("codex realtime quota cache should lock");
cache.insert(
key_id.to_string(),
(
fingerprint,
now.checked_add(Duration::from_secs(CODEX_QUOTA_CACHE_TTL_SECONDS))
.unwrap_or(now),
),
);
cache.retain(|_, (_, expires_at)| *expires_at > now);
if cache.len() <= CODEX_QUOTA_CACHE_MAX_ENTRIES {
return;
}
let mut entries = cache
.iter()
.map(|(key, (_, expires_at))| (key.clone(), *expires_at))
.collect::<Vec<_>>();
entries.sort_by_key(|entry| entry.1);
for (key, _) in entries
.into_iter()
.take(cache.len() - CODEX_QUOTA_CACHE_MAX_ENTRIES)
{
cache.remove(&key);
}
}
fn merge_metadata_object(
current: Option<&Value>,
section_key: &str,
section_value: Value,
) -> Option<Value> {
let mut merged = current
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
merged.insert(section_key.to_string(), section_value);
Some(Value::Object(merged))
}
async fn apply_local_sync_report_effect(state: &AppState, payload: &GatewaySyncReportRequest) {
apply_local_gemini_file_mapping_report_effect(state, payload).await;
if let Err(err) = sync_codex_quota_from_response_headers(
state,
payload.report_context.as_ref(),
&payload.headers,
)
.await
{
warn!(
event_name = "codex_realtime_quota_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist codex realtime quota from sync response headers"
);
}
}
async fn apply_local_stream_report_effect(state: &AppState, payload: &GatewayStreamReportRequest) {
if let Err(err) = sync_codex_quota_from_response_headers(
state,
payload.report_context.as_ref(),
&payload.headers,
)
.await
{
warn!(
event_name = "codex_realtime_quota_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist codex realtime quota from stream response headers"
);
}
}
async fn apply_local_gemini_file_mapping_report_effect(
state: &AppState,
payload: &GatewaySyncReportRequest,
) {
match payload.report_kind.as_str() {
"gemini_files_store_mapping" => {
if payload.status_code >= 300 {
return;
}
let key_id = payload
.report_context
.as_ref()
.and_then(|context| context.get("file_key_id"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let user_id = payload
.report_context
.as_ref()
.and_then(|context| context.get("user_id"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let Some(key_id) = key_id else {
return;
};
for entry in extract_gemini_file_mapping_entries(payload) {
if let Err(err) = store_local_gemini_file_mapping(
state,
entry.file_name.as_str(),
key_id,
user_id,
entry.display_name.as_deref(),
entry.mime_type.as_deref(),
)
.await
{
warn!(
event_name = "gemini_file_mapping_store_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
file_name = %entry.file_name,
error = ?err,
"gateway failed to persist gemini file mapping locally"
);
}
}
}
"gemini_files_delete_mapping" if payload.status_code < 300 => {
let file_name = payload
.report_context
.as_ref()
.and_then(|context| context.get("file_name"))
.and_then(Value::as_str)
.and_then(normalize_gemini_file_name);
let Some(file_name) = file_name else {
return;
};
if let Err(err) = delete_local_gemini_file_mapping(state, file_name.as_str()).await {
warn!(
event_name = "gemini_file_mapping_delete_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
file_name = %file_name,
error = ?err,
"gateway failed to delete gemini file mapping locally"
);
}
}
_ => {}
}
}
pub(crate) async fn store_local_gemini_file_mapping(
state: &AppState,
file_name: &str,
key_id: &str,
user_id: Option<&str>,
display_name: Option<&str>,
mime_type: Option<&str>,
) -> Result<(), GatewayError> {
let Some(file_name) = normalize_gemini_file_name(file_name) else {
return Ok(());
};
let expires_at_unix_secs = current_unix_secs().saturating_add(GEMINI_FILE_MAPPING_TTL_SECONDS);
let _stored = state
.upsert_gemini_file_mapping(
aether_data::repository::gemini_file_mappings::UpsertGeminiFileMappingRecord {
id: Uuid::new_v4().to_string(),
file_name: file_name.clone(),
key_id: key_id.to_string(),
user_id: user_id.map(ToOwned::to_owned),
display_name: display_name.map(ToOwned::to_owned),
mime_type: mime_type.map(ToOwned::to_owned),
source_hash: None,
expires_at_unix_secs,
},
)
.await?;
state
.cache_set_string_with_ttl(
gemini_file_mapping_cache_key(file_name.as_str()).as_str(),
key_id,
GEMINI_FILE_MAPPING_TTL_SECONDS,
)
.await?;
Ok(())
}
async fn delete_local_gemini_file_mapping(
state: &AppState,
file_name: &str,
) -> Result<(), GatewayError> {
let Some(file_name) = normalize_gemini_file_name(file_name) else {
return Ok(());
};
let _deleted = state
.delete_gemini_file_mapping_by_file_name(file_name.as_str())
.await?;
state
.cache_delete_key(gemini_file_mapping_cache_key(file_name.as_str()).as_str())
.await?;
Ok(())
}
async fn sync_codex_quota_from_response_headers(
state: &AppState,
report_context: Option<&Value>,
headers: &BTreeMap<String, String>,
) -> Result<bool, GatewayError> {
let key_id = match report_context_key_id(report_context) {
Some(value) => value,
None => return Ok(false),
};
let now_unix_secs = current_unix_secs();
let Some(parsed) = admin_provider_quota_pure::parse_codex_usage_headers(headers, now_unix_secs)
else {
return Ok(false);
};
let Some(incoming_fingerprint) = fingerprint_codex_payload(&parsed) else {
return Ok(false);
};
let now = Instant::now();
if get_cached_codex_quota_fingerprint(&key_id, now).as_deref()
== Some(incoming_fingerprint.as_str())
{
return Ok(false);
}
let Some(key) = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?
.into_iter()
.next()
else {
set_cached_codex_quota_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
};
let Some(provider) = state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id))
.await?
.into_iter()
.next()
else {
set_cached_codex_quota_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
};
if !provider.provider_type.trim().eq_ignore_ascii_case("codex") {
set_cached_codex_quota_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
}
let current_codex = key
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(Value::as_object)
.cloned()
.unwrap_or_else(serde_json::Map::new);
let current_codex = Value::Object(current_codex);
let Some(current_fingerprint) = fingerprint_codex_payload(&current_codex) else {
set_cached_codex_quota_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
};
if current_fingerprint == incoming_fingerprint {
set_cached_codex_quota_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
}
let updated_upstream_metadata =
merge_metadata_object(key.upstream_metadata.as_ref(), "codex", parsed);
let updated_status_snapshot = sync_provider_key_quota_status_snapshot(
key.status_snapshot.as_ref(),
provider.provider_type.as_str(),
updated_upstream_metadata.as_ref(),
"response_headers",
);
let mut updated_key = key;
updated_key.upstream_metadata = updated_upstream_metadata;
updated_key.status_snapshot = updated_status_snapshot;
updated_key.updated_at_unix_secs = Some(now_unix_secs);
let updated = state
.update_provider_catalog_key(&updated_key)
.await?
.is_some();
if updated {
set_cached_codex_quota_fingerprint(&key_id, incoming_fingerprint, now);
}
Ok(updated)
}
#[cfg(test)]
pub(crate) fn clear_local_report_effect_caches_for_tests() {
if let Some(cache) = CODEX_QUOTA_HEADER_FINGERPRINT_CACHE.get() {
cache
.lock()
.expect("codex realtime quota cache should lock")
.clear();
}
}

15
apps/aether-gateway/src/request_candidate_runtime.rs
View File

@@ -363,6 +363,7 @@ pub(crate) async fn persist_available_local_candidate(
api_key_id: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
retry_index: u32,
candidate_id: &str,
required_capabilities: Option<&Value>,
extra_data: Option<serde_json::Value>,
@@ -378,7 +379,7 @@ pub(crate) async fn persist_available_local_candidate(
username: None,
api_key_name: None,
candidate_index,
retry_index: 0,
retry_index,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
@@ -405,7 +406,7 @@ pub(crate) async fn persist_available_local_candidate(
request_id = %short_request_id(trace_id),
candidate_id = %stored.id,
candidate_index,
retry_index = 0,
retry_index,
status = "available",
source = "planner_available",
provider_id = %candidate.provider_id,
@@ -423,7 +424,7 @@ pub(crate) async fn persist_available_local_candidate(
request_id = %short_request_id(trace_id),
candidate_id = %candidate_id,
candidate_index,
retry_index = 0,
retry_index,
status = "available",
source = "planner_available",
provider_id = %candidate.provider_id,
@@ -452,6 +453,7 @@ pub(crate) async fn persist_skipped_local_candidate(
api_key_id: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
candidate_index: u32,
retry_index: u32,
candidate_id: &str,
required_capabilities: Option<&Value>,
skip_reason: &str,
@@ -468,7 +470,7 @@ pub(crate) async fn persist_skipped_local_candidate(
username: None,
api_key_name: None,
candidate_index,
retry_index: 0,
retry_index,
provider_id: Some(candidate.provider_id.clone()),
endpoint_id: Some(candidate.endpoint_id.clone()),
key_id: Some(candidate.key_id.clone()),
@@ -495,7 +497,7 @@ pub(crate) async fn persist_skipped_local_candidate(
request_id = %short_request_id(trace_id),
candidate_id = %stored.id,
candidate_index,
retry_index = 0,
retry_index,
status = "skipped",
skip_reason,
source = "planner_skipped",
@@ -513,7 +515,7 @@ pub(crate) async fn persist_skipped_local_candidate(
request_id = %short_request_id(trace_id),
candidate_id = %candidate_id,
candidate_index,
retry_index = 0,
retry_index,
status = "skipped",
skip_reason,
source = "planner_skipped",
@@ -892,6 +894,7 @@ mod tests {
"api-key-1",
&sample_minimal_candidate(),
0,
0,
"cand-runtime-cap-123",
Some(&required_capabilities),
None,

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

@@ -688,6 +688,39 @@ impl AppState {
.map_err(|err| GatewayError::Internal(err.to_string()))
}
pub(crate) async fn update_provider_catalog_key_format_health(
&self,
key_id: &str,
api_format: &str,
health_by_format: &serde_json::Value,
) -> Result<bool, GatewayError> {
let api_format = api_format.trim();
if api_format.is_empty() {
return Ok(false);
}
let Some(current_key) = self
.read_provider_catalog_keys_by_ids(&[key_id.to_string()])
.await?
.into_iter()
.next()
else {
return Ok(false);
};
if current_key.health_by_format.as_ref() == Some(health_by_format) {
return Ok(false);
}
self.update_provider_catalog_key_health_state(
key_id,
current_key.is_active,
Some(health_by_format),
current_key.circuit_breaker_by_format.as_ref(),
)
.await
}
pub(crate) async fn update_provider_catalog_key_health_state(
&self,
key_id: &str,

98
apps/aether-gateway/src/state/oauth.rs
View File

@@ -7,6 +7,8 @@ use crate::handlers::shared::default_provider_key_status_snapshot;
use crate::provider_transport::LocalOAuthHttpExecutor;
use super::super::provider_transport;
use crate::provider_key_auth::provider_key_is_oauth_managed;
use aether_admin::provider::quota as admin_provider_quota_pure;
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use base64::{engine::general_purpose::STANDARD, Engine as _};
@@ -676,6 +678,65 @@ impl AppState {
self.oauth_refresh.invalidate_cached_entry(key_id).await
}
pub(crate) async fn mark_provider_catalog_key_oauth_invalid(
&self,
key_id: &str,
provider_type: &str,
invalid_reason: &str,
) -> Result<bool, GatewayError> {
let invalid_reason = invalid_reason.trim();
if invalid_reason.is_empty() {
return Ok(false);
}
let Some(mut latest_key) = self
.data
.list_provider_catalog_keys_by_ids(&[key_id.to_string()])
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?
.into_iter()
.next()
else {
return Ok(false);
};
if !provider_key_is_oauth_managed(&latest_key, provider_type) {
return Ok(false);
}
let now_unix_secs = SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|duration| duration.as_secs())
.unwrap_or(0);
let (oauth_invalid_at_unix_secs, oauth_invalid_reason) = merge_runtime_oauth_invalid_state(
provider_type,
&latest_key,
invalid_reason,
now_unix_secs,
);
if oauth_invalid_at_unix_secs == latest_key.oauth_invalid_at_unix_secs
&& oauth_invalid_reason == latest_key.oauth_invalid_reason
{
return Ok(false);
}
latest_key.oauth_invalid_at_unix_secs = oauth_invalid_at_unix_secs;
latest_key.oauth_invalid_reason = oauth_invalid_reason;
latest_key.updated_at_unix_secs = Some(now_unix_secs);
let current_status_snapshot = latest_key.status_snapshot.take();
latest_key.status_snapshot =
sync_provider_key_oauth_status_snapshot(current_status_snapshot, &latest_key);
let updated = self
.update_provider_catalog_key(&latest_key)
.await?
.is_some();
if updated {
let _ = self.invalidate_local_oauth_refresh_entry(key_id).await;
}
Ok(updated)
}
async fn persist_local_oauth_refresh_entry(
&self,
transport: &provider_transport::GatewayProviderTransportSnapshot,
@@ -858,6 +919,43 @@ impl AppState {
}
}
fn merge_runtime_oauth_invalid_state(
provider_type: &str,
key: &StoredProviderCatalogKey,
invalid_reason: &str,
now_unix_secs: u64,
) -> (Option<u64>, Option<String>) {
let candidate_reason = invalid_reason.trim();
if candidate_reason.is_empty() {
return (
key.oauth_invalid_at_unix_secs,
key.oauth_invalid_reason.clone(),
);
}
if provider_type.trim().eq_ignore_ascii_case("codex") {
return admin_provider_quota_pure::codex_build_invalid_state(
key,
candidate_reason.to_string(),
now_unix_secs,
);
}
let current_reason = key
.oauth_invalid_reason
.as_deref()
.map(str::trim)
.unwrap_or_default();
if current_reason == candidate_reason {
return (
key.oauth_invalid_at_unix_secs,
(!current_reason.is_empty()).then_some(current_reason.to_string()),
);
}
(Some(now_unix_secs), Some(candidate_reason.to_string()))
}
fn local_oauth_execution_body_text(result: &aether_contracts::ExecutionResult) -> String {
result
.body

33
apps/aether-gateway/src/tests/architecture/usage.rs
View File

@@ -177,13 +177,11 @@ fn usage_runtime_paths_depend_on_shared_crates_not_app_runtime_shims() {
"is_local_ai_sync_report_kind",
"is_local_ai_stream_report_kind",
"sync_report_represents_failure",
"extract_gemini_file_mapping_entries",
"gemini_file_mapping_cache_key",
"normalize_gemini_file_name",
"report_request_id",
"should_handle_local_sync_report",
"should_handle_local_stream_report",
"GEMINI_FILE_MAPPING_TTL_SECONDS",
"apply_local_report_effect",
"LocalReportEffect",
] {
assert!(
usage_reporting_mod.contains(pattern),
@@ -206,12 +204,39 @@ fn usage_runtime_paths_depend_on_shared_crates_not_app_runtime_shims() {
"fn should_handle_local_sync_report(",
"fn should_handle_local_stream_report(",
"\"openai_video_delete_sync_success\" && payload.status_code == 404",
"sync_codex_quota_from_response_headers(",
"apply_local_gemini_file_mapping_report_effect(",
"pub(crate) async fn store_local_gemini_file_mapping(",
] {
assert!(
!usage_reporting_mod.contains(pattern),
"usage/reporting/mod.rs should not own local report classification logic {pattern}"
);
}
let report_effects =
read_workspace_file("apps/aether-gateway/src/orchestration/report_effects.rs");
assert!(
report_effects.contains("aether_usage_runtime"),
"orchestration/report_effects.rs should depend on aether_usage_runtime"
);
for pattern in [
"extract_gemini_file_mapping_entries",
"gemini_file_mapping_cache_key",
"normalize_gemini_file_name",
"report_request_id",
"GEMINI_FILE_MAPPING_TTL_SECONDS",
"sync_codex_quota_from_response_headers",
"store_local_gemini_file_mapping",
"delete_local_gemini_file_mapping",
"GatewaySyncReportRequest",
"GatewayStreamReportRequest",
] {
assert!(
report_effects.contains(pattern),
"orchestration/report_effects.rs should own local report effect detail {pattern}"
);
}
}
#[test]

61
apps/aether-gateway/src/tests/control/admin/oauth.rs
View File

@@ -2768,6 +2768,49 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
}),
);
#[derive(Debug, Clone)]
struct SeenExecutionRuntimeRequest {
url: String,
authorization: String,
}
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
let execution_runtime = Router::new().route(
"/v1/execute/sync",
any(move |request: Request| {
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
async move {
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
&to_bytes(request.into_body(), usize::MAX)
.await
.expect("body should read"),
)
.expect("plan should parse");
*seen_execution_runtime_inner
.lock()
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
url: plan.url.clone(),
authorization: plan
.headers
.get("authorization")
.cloned()
.unwrap_or_default(),
});
let result = aether_contracts::ExecutionResult {
request_id: plan.request_id,
candidate_id: None,
status_code: 401,
headers: std::collections::BTreeMap::new(),
body: None,
telemetry: None,
error: None,
};
(StatusCode::OK, Json(result))
}
}),
);
let mut provider = sample_provider("provider-codex", "codex", 10);
provider.provider_type = "codex".to_string();
let endpoint = sample_endpoint(
@@ -2832,6 +2875,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
let (upstream_url, upstream_handle) = start_server(upstream).await;
let (token_url, token_handle) = start_server(token_server).await;
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
let oauth_refresh =
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
Arc::new(
@@ -2840,8 +2884,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
),
]);
let gateway = build_router_with_state(
AppState::new()
.expect("gateway should build")
build_state_with_execution_runtime_override(execution_runtime_url)
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(),
@@ -2889,6 +2932,19 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
}
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
let seen_execution_runtime_request = seen_execution_runtime
.lock()
.expect("mutex should lock")
.clone()
.expect("execution runtime request should be captured");
assert_eq!(
seen_execution_runtime_request.url,
"https://chatgpt.com/backend-api/wham/usage"
);
assert_eq!(
seen_execution_runtime_request.authorization,
"Bearer refreshed-codex-access-token"
);
let stored_key = provider_catalog_repository
.list_keys_by_ids(&["key-codex-oauth-refresh".to_string()])
@@ -2974,6 +3030,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
);
gateway_handle.abort();
execution_runtime_handle.abort();
token_handle.abort();
upstream_handle.abort();
}

625
apps/aether-gateway/src/usage/reporting/codex_realtime_quota.rs
View File

@@ -1,625 +0,0 @@
use std::collections::{BTreeMap, HashMap};
use std::sync::{Mutex, OnceLock};
use std::time::{Duration, Instant};
use crate::clock::current_unix_secs;
use crate::handlers::shared::sync_provider_key_quota_status_snapshot;
use crate::{AppState, GatewayError};
use aether_admin::provider::quota as admin_provider_quota_pure;
use serde_json::Value;
const CACHE_TTL_SECONDS: u64 = 30;
const CACHE_MAX_ENTRIES: usize = 4096;
type HeaderFingerprintCache = Mutex<HashMap<String, (String, Instant)>>;
static HEADER_FINGERPRINT_CACHE: OnceLock<HeaderFingerprintCache> = OnceLock::new();
fn header_fingerprint_cache() -> &'static HeaderFingerprintCache {
HEADER_FINGERPRINT_CACHE.get_or_init(|| Mutex::new(HashMap::new()))
}
fn report_context_key_id(report_context: Option<&Value>) -> Option<String> {
report_context
.and_then(|context| context.get("key_id"))
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn is_volatile_compare_field(key: &str) -> bool {
key == "updated_at" || key.ends_with("_reset_seconds") || key.ends_with("_reset_after_seconds")
}
fn canonicalize_value(value: &Value) -> Value {
match value {
Value::Array(items) => Value::Array(items.iter().map(canonicalize_value).collect()),
Value::Object(object) => {
let mut entries = object.iter().collect::<Vec<_>>();
entries.sort_by(|left, right| left.0.cmp(right.0));
let mut normalized = serde_json::Map::new();
for (key, value) in entries {
normalized.insert(key.clone(), canonicalize_value(value));
}
Value::Object(normalized)
}
_ => value.clone(),
}
}
fn fingerprint_codex_payload(value: &Value) -> Option<String> {
let object = value.as_object()?;
let mut entries = object
.iter()
.filter(|(key, _)| !is_volatile_compare_field(key))
.collect::<Vec<_>>();
entries.sort_by(|left, right| left.0.cmp(right.0));
let mut normalized = serde_json::Map::new();
for (key, value) in entries {
normalized.insert(key.clone(), canonicalize_value(value));
}
serde_json::to_string(&Value::Object(normalized)).ok()
}
fn get_cached_fingerprint(key_id: &str, now: Instant) -> Option<String> {
let mut cache = header_fingerprint_cache()
.lock()
.expect("codex realtime quota cache should lock");
match cache.get(key_id) {
Some((fingerprint, expires_at)) if *expires_at > now => Some(fingerprint.clone()),
Some(_) => {
cache.remove(key_id);
None
}
None => None,
}
}
fn set_cached_fingerprint(key_id: &str, fingerprint: String, now: Instant) {
let mut cache = header_fingerprint_cache()
.lock()
.expect("codex realtime quota cache should lock");
cache.insert(
key_id.to_string(),
(
fingerprint,
now.checked_add(Duration::from_secs(CACHE_TTL_SECONDS))
.unwrap_or(now),
),
);
cache.retain(|_, (_, expires_at)| *expires_at > now);
if cache.len() <= CACHE_MAX_ENTRIES {
return;
}
let mut entries = cache
.iter()
.map(|(key, (_, expires_at))| (key.clone(), *expires_at))
.collect::<Vec<_>>();
entries.sort_by_key(|entry| entry.1);
for (key, _) in entries.into_iter().take(cache.len() - CACHE_MAX_ENTRIES) {
cache.remove(&key);
}
}
fn merge_metadata_object(
current: Option<&Value>,
section_key: &str,
section_value: Value,
) -> Option<Value> {
let mut merged = current
.and_then(Value::as_object)
.cloned()
.unwrap_or_default();
merged.insert(section_key.to_string(), section_value);
Some(Value::Object(merged))
}
pub(super) async fn sync_codex_quota_from_response_headers(
state: &AppState,
report_context: Option<&Value>,
headers: &BTreeMap<String, String>,
) -> Result<bool, GatewayError> {
let key_id = match report_context_key_id(report_context) {
Some(value) => value,
None => return Ok(false),
};
let now_unix_secs = current_unix_secs();
let Some(parsed) = admin_provider_quota_pure::parse_codex_usage_headers(headers, now_unix_secs)
else {
return Ok(false);
};
let Some(incoming_fingerprint) = fingerprint_codex_payload(&parsed) else {
return Ok(false);
};
let now = Instant::now();
if get_cached_fingerprint(&key_id, now).as_deref() == Some(incoming_fingerprint.as_str()) {
return Ok(false);
}
let Some(key) = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
.await?
.into_iter()
.next()
else {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
};
let Some(provider) = state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id))
.await?
.into_iter()
.next()
else {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
};
if !provider.provider_type.trim().eq_ignore_ascii_case("codex") {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
}
let current_codex = key
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(Value::as_object)
.cloned()
.unwrap_or_else(serde_json::Map::new);
let current_codex = Value::Object(current_codex);
let Some(current_fingerprint) = fingerprint_codex_payload(&current_codex) else {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
};
if current_fingerprint == incoming_fingerprint {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
return Ok(false);
}
let updated_upstream_metadata =
merge_metadata_object(key.upstream_metadata.as_ref(), "codex", parsed);
let updated_status_snapshot = sync_provider_key_quota_status_snapshot(
key.status_snapshot.as_ref(),
provider.provider_type.as_str(),
updated_upstream_metadata.as_ref(),
"response_headers",
);
let mut updated_key = key;
updated_key.upstream_metadata = updated_upstream_metadata;
updated_key.status_snapshot = updated_status_snapshot;
updated_key.updated_at_unix_secs = Some(now_unix_secs);
let updated = state
.update_provider_catalog_key(&updated_key)
.await?
.is_some();
if updated {
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
}
Ok(updated)
}
#[cfg(test)]
pub(super) fn clear_codex_quota_fingerprint_cache() {
if let Some(cache) = HEADER_FINGERPRINT_CACHE.get() {
cache
.lock()
.expect("codex realtime quota cache should lock")
.clear();
}
}
#[cfg(test)]
mod tests {
use super::{clear_codex_quota_fingerprint_cache, sync_codex_quota_from_response_headers};
use crate::data::GatewayDataState;
use crate::AppState;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use serde_json::{json, Value};
use std::collections::BTreeMap;
use std::sync::Arc;
fn build_state(repository: Arc<InMemoryProviderCatalogReadRepository>) -> AppState {
AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(
GatewayDataState::with_provider_catalog_repository_for_tests(repository),
)
}
fn sample_provider(provider_id: &str, provider_type: &str) -> StoredProviderCatalogProvider {
StoredProviderCatalogProvider::new(
provider_id.to_string(),
provider_type.to_string(),
None,
provider_type.to_string(),
)
.expect("provider should build")
}
fn sample_key(
key_id: &str,
provider_id: &str,
upstream_metadata: Option<Value>,
) -> StoredProviderCatalogKey {
let mut key = StoredProviderCatalogKey::new(
key_id.to_string(),
provider_id.to_string(),
"default".to_string(),
"bearer".to_string(),
None,
true,
)
.expect("key should build")
.with_transport_fields(
Some(json!(["openai:cli"])),
"sk-codex-test".to_string(),
None,
None,
None,
None,
None,
None,
None,
)
.expect("key transport should build");
key.upstream_metadata = upstream_metadata;
key
}
fn quota_snapshot<'a>(key: &'a StoredProviderCatalogKey) -> &'a serde_json::Map<String, Value> {
key.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("quota"))
.and_then(Value::as_object)
.expect("quota snapshot should exist")
}
fn paid_headers(
primary_used_percent: &str,
secondary_used_percent: &str,
primary_reset_after_seconds: &str,
secondary_reset_after_seconds: &str,
) -> BTreeMap<String, String> {
BTreeMap::from([
("x-codex-plan-type".to_string(), "team".to_string()),
(
"x-codex-primary-used-percent".to_string(),
primary_used_percent.to_string(),
),
(
"x-codex-secondary-used-percent".to_string(),
secondary_used_percent.to_string(),
),
(
"x-codex-primary-window-minutes".to_string(),
"300".to_string(),
),
(
"x-codex-secondary-window-minutes".to_string(),
"10080".to_string(),
),
(
"x-codex-primary-reset-after-seconds".to_string(),
primary_reset_after_seconds.to_string(),
),
(
"x-codex-secondary-reset-after-seconds".to_string(),
secondary_reset_after_seconds.to_string(),
),
(
"x-codex-primary-reset-at".to_string(),
"1776148929".to_string(),
),
(
"x-codex-secondary-reset-at".to_string(),
"1776657828".to_string(),
),
(
"x-codex-credits-has-credits".to_string(),
"False".to_string(),
),
("x-codex-credits-balance".to_string(), "".to_string()),
("x-codex-credits-unlimited".to_string(), "False".to_string()),
])
}
#[tokio::test]
async fn sync_codex_quota_replaces_existing_codex_fields_and_preserves_other_sections() {
clear_codex_quota_fingerprint_cache();
let mut key = sample_key(
"key-codex-1",
"provider-codex",
Some(json!({
"codex": {
"legacy_marker": "drop-me",
"secondary_used_percent": 2.0,
"credits_balance": 42.0,
"account_disabled": true,
"reason": "deactivated_workspace"
},
"other": {
"value": true
}
})),
);
key.status_snapshot = Some(json!({
"oauth": {
"code": "valid",
"label": "有效",
"requires_reauth": false,
"expiring_soon": false
},
"account": {
"code": "ok",
"blocked": false,
"recoverable": false
},
"quota": {
"code": "unknown",
"exhausted": false
}
}));
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-codex", "codex")],
Vec::new(),
vec![key],
));
let state = build_state(Arc::clone(&repository));
let updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-codex-realtime-1",
"key_id": "key-codex-1"
})),
&paid_headers("100", "31", "15160", "524059"),
)
.await
.expect("codex realtime sync should succeed");
assert!(updated);
let reloaded = repository
.list_keys_by_ids(&["key-codex-1".to_string()])
.await
.expect("keys should list");
let codex = reloaded[0]
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(Value::as_object)
.expect("codex metadata should exist");
assert_eq!(codex.get("plan_type"), Some(&json!("team")));
assert_eq!(codex.get("primary_used_percent"), Some(&json!(31.0)));
assert_eq!(codex.get("secondary_used_percent"), Some(&json!(100.0)));
assert_eq!(codex.get("has_credits"), Some(&json!(false)));
assert_eq!(codex.get("credits_unlimited"), Some(&json!(false)));
assert!(codex.get("legacy_marker").is_none());
assert!(codex.get("credits_balance").is_none());
assert!(codex.get("account_disabled").is_none());
assert!(codex.get("reason").is_none());
assert!(codex.get("updated_at").and_then(Value::as_u64).is_some());
assert_eq!(
reloaded[0]
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("other")),
Some(&json!({"value": true}))
);
let quota = quota_snapshot(&reloaded[0]);
assert_eq!(quota.get("version"), Some(&json!(2)));
assert_eq!(quota.get("provider_type"), Some(&json!("codex")));
assert_eq!(quota.get("source"), Some(&json!("response_headers")));
assert_eq!(quota.get("code"), Some(&json!("exhausted")));
assert_eq!(quota.get("exhausted"), Some(&json!(true)));
assert_eq!(quota.get("plan_type"), Some(&json!("team")));
assert_eq!(quota.get("usage_ratio"), Some(&json!(1.0)));
assert_eq!(quota.get("updated_at"), quota.get("observed_at"));
let windows = quota
.get("windows")
.and_then(Value::as_array)
.expect("windows should be array");
assert_eq!(windows.len(), 2);
assert_eq!(windows[0].get("code"), Some(&json!("weekly")));
assert_eq!(windows[1].get("code"), Some(&json!("5h")));
let oauth = reloaded[0]
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("oauth"))
.and_then(Value::as_object)
.expect("oauth snapshot should exist");
assert_eq!(oauth.get("code"), Some(&json!("valid")));
let account = reloaded[0]
.status_snapshot
.as_ref()
.and_then(Value::as_object)
.and_then(|snapshot| snapshot.get("account"))
.and_then(Value::as_object)
.expect("account snapshot should exist");
assert_eq!(account.get("code"), Some(&json!("ok")));
}
#[tokio::test]
async fn sync_codex_quota_skips_non_codex_provider() {
clear_codex_quota_fingerprint_cache();
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-openai", "openai")],
Vec::new(),
vec![sample_key("key-openai-1", "provider-openai", None)],
));
let state = build_state(Arc::clone(&repository));
let updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-openai-realtime-1",
"key_id": "key-openai-1"
})),
&paid_headers("100", "31", "15160", "524059"),
)
.await
.expect("non-codex realtime sync should not fail");
assert!(!updated);
let reloaded = repository
.list_keys_by_ids(&["key-openai-1".to_string()])
.await
.expect("keys should list");
assert_eq!(reloaded[0].upstream_metadata, None);
}
#[tokio::test]
async fn sync_codex_quota_skips_when_headers_do_not_contain_codex_metadata() {
clear_codex_quota_fingerprint_cache();
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-codex", "codex")],
Vec::new(),
vec![sample_key("key-codex-2", "provider-codex", None)],
));
let state = build_state(Arc::clone(&repository));
let updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-codex-realtime-2",
"key_id": "key-codex-2"
})),
&BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
)
.await
.expect("empty codex headers should not fail");
assert!(!updated);
let reloaded = repository
.list_keys_by_ids(&["key-codex-2".to_string()])
.await
.expect("keys should list");
assert_eq!(reloaded[0].upstream_metadata, None);
}
#[tokio::test]
async fn sync_codex_quota_cache_hit_skips_when_only_reset_countdown_changes() {
clear_codex_quota_fingerprint_cache();
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-codex", "codex")],
Vec::new(),
vec![sample_key("key-codex-3", "provider-codex", None)],
));
let state = build_state(Arc::clone(&repository));
let key_ids = ["key-codex-3".to_string()];
let first_updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-codex-realtime-3",
"key_id": "key-codex-3"
})),
&paid_headers("100", "31", "15160", "524059"),
)
.await
.expect("first realtime sync should succeed");
assert!(first_updated);
let first_snapshot = repository
.list_keys_by_ids(&key_ids)
.await
.expect("keys should list")
.into_iter()
.next()
.expect("key should exist");
let second_updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-codex-realtime-3",
"key_id": "key-codex-3"
})),
&paid_headers("100", "31", "42", "84"),
)
.await
.expect("second realtime sync should succeed");
assert!(!second_updated);
let second_snapshot = repository
.list_keys_by_ids(&key_ids)
.await
.expect("keys should list")
.into_iter()
.next()
.expect("key should exist");
assert_eq!(
first_snapshot.upstream_metadata,
second_snapshot.upstream_metadata
);
assert_eq!(
first_snapshot.updated_at_unix_secs,
second_snapshot.updated_at_unix_secs
);
}
#[tokio::test]
async fn sync_codex_quota_updates_when_usage_percent_changes() {
clear_codex_quota_fingerprint_cache();
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-codex", "codex")],
Vec::new(),
vec![sample_key("key-codex-4", "provider-codex", None)],
));
let state = build_state(Arc::clone(&repository));
let first_updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-codex-realtime-4",
"key_id": "key-codex-4"
})),
&paid_headers("98", "31", "15160", "524059"),
)
.await
.expect("first realtime sync should succeed");
assert!(first_updated);
let second_updated = sync_codex_quota_from_response_headers(
&state,
Some(&json!({
"request_id": "req-codex-realtime-4",
"key_id": "key-codex-4"
})),
&paid_headers("100", "31", "15160", "524059"),
)
.await
.expect("second realtime sync should succeed");
assert!(second_updated);
let reloaded = repository
.list_keys_by_ids(&["key-codex-4".to_string()])
.await
.expect("keys should list");
let codex = reloaded[0]
.upstream_metadata
.as_ref()
.and_then(Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(Value::as_object)
.expect("codex metadata should exist");
assert_eq!(codex.get("secondary_used_percent"), Some(&json!(100.0)));
}
}

183
apps/aether-gateway/src/usage/reporting/mod.rs
View File

@@ -4,22 +4,20 @@ use aether_contracts::ExecutionError;
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::{execution_error_details, SchedulerRequestCandidateStatusUpdate};
use tracing::{debug, warn};
use uuid::Uuid;
use crate::clock::{current_unix_ms, current_unix_secs};
use crate::clock::current_unix_ms;
use crate::log_ids::short_request_id;
use crate::orchestration::{apply_local_report_effect, LocalReportEffect};
use crate::request_candidate_runtime::record_report_request_candidate_status;
use crate::{AppState, GatewayError};
mod codex_realtime_quota;
mod context;
use context::{report_context_is_locally_actionable, resolve_locally_actionable_report_context};
use aether_usage_runtime::{
extract_gemini_file_mapping_entries, gemini_file_mapping_cache_key,
is_local_ai_stream_report_kind, is_local_ai_sync_report_kind, normalize_gemini_file_name,
report_request_id, should_handle_local_stream_report, should_handle_local_sync_report,
sync_report_represents_failure, GEMINI_FILE_MAPPING_TTL_SECONDS,
is_local_ai_stream_report_kind, is_local_ai_sync_report_kind, report_request_id,
should_handle_local_stream_report, should_handle_local_sync_report,
sync_report_represents_failure,
};
pub(crate) use aether_usage_runtime::{GatewayStreamReportRequest, GatewaySyncReportRequest};
@@ -188,7 +186,6 @@ pub(crate) async fn submit_stream_report(
}
async fn handle_local_sync_report(state: &AppState, payload: &GatewaySyncReportRequest) {
apply_local_gemini_file_mapping_side_effect(state, payload).await;
let terminal_unix_ms = current_unix_ms();
let (error_type, error_message) =
execution_error_details(None::<&ExecutionError>, payload.body_json.as_ref());
@@ -215,22 +212,7 @@ async fn handle_local_sync_report(state: &AppState, payload: &GatewaySyncReportR
},
)
.await;
if let Err(err) = codex_realtime_quota::sync_codex_quota_from_response_headers(
state,
payload.report_context.as_ref(),
&payload.headers,
)
.await
{
warn!(
event_name = "codex_realtime_quota_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist codex realtime quota from sync response headers"
);
}
apply_local_report_effect(state, LocalReportEffect::Sync { payload }).await;
}
async fn handle_local_stream_report(state: &AppState, payload: &GatewayStreamReportRequest) {
@@ -253,154 +235,7 @@ async fn handle_local_stream_report(state: &AppState, payload: &GatewayStreamRep
},
)
.await;
if let Err(err) = codex_realtime_quota::sync_codex_quota_from_response_headers(
state,
payload.report_context.as_ref(),
&payload.headers,
)
.await
{
warn!(
event_name = "codex_realtime_quota_sync_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
error = ?err,
"gateway failed to persist codex realtime quota from stream response headers"
);
}
}
async fn apply_local_gemini_file_mapping_side_effect(
state: &AppState,
payload: &GatewaySyncReportRequest,
) {
match payload.report_kind.as_str() {
"gemini_files_store_mapping" => {
if payload.status_code >= 300 {
return;
}
let key_id = payload
.report_context
.as_ref()
.and_then(|context| context.get("file_key_id"))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let user_id = payload
.report_context
.as_ref()
.and_then(|context| context.get("user_id"))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let Some(key_id) = key_id else {
return;
};
for entry in extract_gemini_file_mapping_entries(payload) {
if let Err(err) = store_local_gemini_file_mapping(
state,
entry.file_name.as_str(),
key_id,
user_id,
entry.display_name.as_deref(),
entry.mime_type.as_deref(),
)
.await
{
warn!(
event_name = "gemini_file_mapping_store_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
file_name = %entry.file_name,
error = ?err,
"gateway failed to persist gemini file mapping locally"
);
}
}
}
"gemini_files_delete_mapping" if payload.status_code < 300 => {
let file_name = payload
.report_context
.as_ref()
.and_then(|context| context.get("file_name"))
.and_then(serde_json::Value::as_str)
.and_then(normalize_gemini_file_name);
let Some(file_name) = file_name else {
return;
};
if let Err(err) = delete_local_gemini_file_mapping(state, file_name.as_str()).await {
warn!(
event_name = "gemini_file_mapping_delete_failed",
log_type = "ops",
report_kind = %payload.report_kind,
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
file_name = %file_name,
error = ?err,
"gateway failed to delete gemini file mapping locally"
);
}
}
_ => {}
}
}
pub(crate) async fn store_local_gemini_file_mapping(
state: &AppState,
file_name: &str,
key_id: &str,
user_id: Option<&str>,
display_name: Option<&str>,
mime_type: Option<&str>,
) -> Result<(), GatewayError> {
let Some(file_name) = normalize_gemini_file_name(file_name) else {
return Ok(());
};
let expires_at_unix_secs = current_unix_secs().saturating_add(GEMINI_FILE_MAPPING_TTL_SECONDS);
let _stored = state
.upsert_gemini_file_mapping(
aether_data::repository::gemini_file_mappings::UpsertGeminiFileMappingRecord {
id: Uuid::new_v4().to_string(),
file_name: file_name.clone(),
key_id: key_id.to_string(),
user_id: user_id.map(ToOwned::to_owned),
display_name: display_name.map(ToOwned::to_owned),
mime_type: mime_type.map(ToOwned::to_owned),
source_hash: None,
expires_at_unix_secs,
},
)
.await?;
state
.cache_set_string_with_ttl(
gemini_file_mapping_cache_key(file_name.as_str()).as_str(),
key_id,
GEMINI_FILE_MAPPING_TTL_SECONDS,
)
.await?;
Ok(())
}
async fn delete_local_gemini_file_mapping(
state: &AppState,
file_name: &str,
) -> Result<(), GatewayError> {
let Some(file_name) = normalize_gemini_file_name(file_name) else {
return Ok(());
};
let _deleted = state
.delete_gemini_file_mapping_by_file_name(file_name.as_str())
.await?;
state
.cache_delete_key(gemini_file_mapping_cache_key(file_name.as_str()).as_str())
.await?;
Ok(())
apply_local_report_effect(state, LocalReportEffect::Stream { payload }).await;
}
#[cfg(test)]
@@ -824,7 +659,7 @@ mod tests {
#[tokio::test]
async fn submit_sync_report_updates_codex_quota_from_response_headers() {
super::codex_realtime_quota::clear_codex_quota_fingerprint_cache();
crate::orchestration::clear_local_report_effect_caches_for_tests();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider(
@@ -888,7 +723,7 @@ mod tests {
#[tokio::test]
async fn submit_stream_report_updates_codex_quota_from_response_headers() {
super::codex_realtime_quota::clear_codex_quota_fingerprint_cache();
crate::orchestration::clear_local_report_effect_caches_for_tests();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider(