mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-12 14:10:19 +08:00
refactor gateway orchestration and failover effects
This commit is contained in:
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user