refactor gateway orchestration and failover effects

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
@@ -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",