i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-03 01:40:21 +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

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