mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
refactor gateway orchestration and failover effects
This commit is contained in:
@@ -4,6 +4,7 @@ use aether_scheduler_core::SchedulerMinimalCandidateSelectionCandidate;
|
|||||||
use std::collections::BTreeSet;
|
use std::collections::BTreeSet;
|
||||||
|
|
||||||
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
|
use crate::ai_pipeline::{GatewayProviderTransportSnapshot, PlannerAppState};
|
||||||
|
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||||
|
|
||||||
use super::candidate_affinity::rank_eligible_local_execution_candidates;
|
use super::candidate_affinity::rank_eligible_local_execution_candidates;
|
||||||
use super::pool_scheduler::apply_local_execution_pool_scheduler;
|
use super::pool_scheduler::apply_local_execution_pool_scheduler;
|
||||||
@@ -13,6 +14,7 @@ pub(crate) struct EligibleLocalExecutionCandidate {
|
|||||||
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
|
pub(crate) candidate: SchedulerMinimalCandidateSelectionCandidate,
|
||||||
pub(crate) transport: GatewayProviderTransportSnapshot,
|
pub(crate) transport: GatewayProviderTransportSnapshot,
|
||||||
pub(crate) provider_api_format: String,
|
pub(crate) provider_api_format: String,
|
||||||
|
pub(crate) orchestration: LocalExecutionCandidateMetadata,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, PartialEq)]
|
#[derive(Debug, Clone, PartialEq)]
|
||||||
@@ -122,6 +124,7 @@ where
|
|||||||
provider_api_format: transport.endpoint.api_format.trim().to_ascii_lowercase(),
|
provider_api_format: transport.endpoint.api_format.trim().to_ascii_lowercase(),
|
||||||
candidate,
|
candidate,
|
||||||
transport,
|
transport,
|
||||||
|
orchestration: LocalExecutionCandidateMetadata::default(),
|
||||||
}),
|
}),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -9,15 +9,26 @@ use crate::ai_pipeline::planner::candidate_eligibility::{
|
|||||||
use crate::ai_pipeline::planner::runtime_miss::record_local_runtime_candidate_skip_reason;
|
use crate::ai_pipeline::planner::runtime_miss::record_local_runtime_candidate_skip_reason;
|
||||||
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
use crate::ai_pipeline::{GatewayAuthApiKeySnapshot, PlannerAppState};
|
||||||
use crate::clock::current_unix_ms;
|
use crate::clock::current_unix_ms;
|
||||||
|
use crate::orchestration::{build_local_attempt_identities, ExecutionAttemptIdentity};
|
||||||
use crate::AppState;
|
use crate::AppState;
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub(crate) struct LocalExecutionCandidateAttempt {
|
pub(crate) struct LocalExecutionCandidateAttempt {
|
||||||
pub(crate) eligible: EligibleLocalExecutionCandidate,
|
pub(crate) eligible: EligibleLocalExecutionCandidate,
|
||||||
pub(crate) candidate_index: u32,
|
pub(crate) candidate_index: u32,
|
||||||
|
pub(crate) retry_index: u32,
|
||||||
|
pub(crate) pool_key_index: Option<u32>,
|
||||||
|
pub(crate) candidate_group_id: Option<String>,
|
||||||
pub(crate) candidate_id: String,
|
pub(crate) candidate_id: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl LocalExecutionCandidateAttempt {
|
||||||
|
pub(crate) fn attempt_identity(&self) -> ExecutionAttemptIdentity {
|
||||||
|
ExecutionAttemptIdentity::new(self.candidate_index, self.retry_index)
|
||||||
|
.with_pool_key_index(self.pool_key_index)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Copy)]
|
#[derive(Debug, Clone, Copy)]
|
||||||
pub(crate) struct LocalAvailableCandidatePersistenceContext<'a> {
|
pub(crate) struct LocalAvailableCandidatePersistenceContext<'a> {
|
||||||
pub(crate) user_id: &'a str,
|
pub(crate) user_id: &'a str,
|
||||||
@@ -73,30 +84,43 @@ where
|
|||||||
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value>,
|
F: Fn(&EligibleLocalExecutionCandidate) -> Option<Value>,
|
||||||
{
|
{
|
||||||
let created_at_unix_ms = current_unix_ms();
|
let created_at_unix_ms = current_unix_ms();
|
||||||
let mut materialized = Vec::with_capacity(candidates.len());
|
let mut materialized = Vec::new();
|
||||||
|
|
||||||
for (candidate_index, eligible) in candidates.into_iter().enumerate() {
|
for (candidate_index, eligible) in candidates.into_iter().enumerate() {
|
||||||
let generated_candidate_id = Uuid::new_v4().to_string();
|
let candidate_index = candidate_index as u32;
|
||||||
let candidate_id = state
|
let attempt_identities =
|
||||||
.persist_available_local_candidate(
|
build_local_attempt_identities(candidate_index, &eligible.transport)
|
||||||
trace_id,
|
.into_iter()
|
||||||
user_id,
|
.map(|identity| identity.with_pool_key_index(eligible.orchestration.pool_key_index))
|
||||||
api_key_id,
|
.collect::<Vec<_>>();
|
||||||
&eligible.candidate,
|
|
||||||
candidate_index as u32,
|
|
||||||
&generated_candidate_id,
|
|
||||||
required_capabilities,
|
|
||||||
build_extra_data(&eligible),
|
|
||||||
created_at_unix_ms,
|
|
||||||
error_context,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
|
|
||||||
materialized.push(LocalExecutionCandidateAttempt {
|
for attempt_identity in attempt_identities {
|
||||||
eligible,
|
let generated_candidate_id = Uuid::new_v4().to_string();
|
||||||
candidate_index: candidate_index as u32,
|
let candidate_id = state
|
||||||
candidate_id,
|
.persist_available_local_candidate(
|
||||||
});
|
trace_id,
|
||||||
|
user_id,
|
||||||
|
api_key_id,
|
||||||
|
&eligible.candidate,
|
||||||
|
attempt_identity.candidate_index,
|
||||||
|
attempt_identity.retry_index,
|
||||||
|
&generated_candidate_id,
|
||||||
|
required_capabilities,
|
||||||
|
build_extra_data(&eligible),
|
||||||
|
created_at_unix_ms,
|
||||||
|
error_context,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
materialized.push(LocalExecutionCandidateAttempt {
|
||||||
|
eligible: eligible.clone(),
|
||||||
|
candidate_index: attempt_identity.candidate_index,
|
||||||
|
retry_index: attempt_identity.retry_index,
|
||||||
|
pool_key_index: attempt_identity.pool_key_index,
|
||||||
|
candidate_group_id: eligible.orchestration.candidate_group_id.clone(),
|
||||||
|
candidate_id,
|
||||||
|
});
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
materialized
|
materialized
|
||||||
@@ -151,6 +175,7 @@ pub(crate) async fn persist_skipped_local_execution_candidate(
|
|||||||
api_key_id,
|
api_key_id,
|
||||||
candidate,
|
candidate,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
|
0,
|
||||||
candidate_id,
|
candidate_id,
|
||||||
required_capabilities,
|
required_capabilities,
|
||||||
skip_reason,
|
skip_reason,
|
||||||
|
|||||||
@@ -155,6 +155,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
|||||||
Some(&input.requested_model),
|
Some(&input.requested_model),
|
||||||
&candidates,
|
&candidates,
|
||||||
);
|
);
|
||||||
|
let available_candidate_count = candidates.len() as u32;
|
||||||
let attempts = persist_available_local_execution_candidates_with_context(
|
let attempts = persist_available_local_execution_candidates_with_context(
|
||||||
planner_state,
|
planner_state,
|
||||||
trace_id,
|
trace_id,
|
||||||
@@ -180,7 +181,7 @@ pub(crate) async fn materialize_local_same_format_provider_candidate_attempts(
|
|||||||
state,
|
state,
|
||||||
trace_id,
|
trace_id,
|
||||||
persistence_policy.skipped,
|
persistence_policy.skipped,
|
||||||
attempts.len() as u32,
|
available_candidate_count,
|
||||||
skipped_candidates,
|
skipped_candidates,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|||||||
@@ -41,7 +41,9 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
|||||||
let LocalSameFormatProviderCandidateAttempt {
|
let LocalSameFormatProviderCandidateAttempt {
|
||||||
eligible,
|
eligible,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
|
candidate_group_id,
|
||||||
candidate_id,
|
candidate_id,
|
||||||
|
..
|
||||||
} = &attempt;
|
} = &attempt;
|
||||||
let candidate = &eligible.candidate;
|
let candidate = &eligible.candidate;
|
||||||
let resolved = resolve_local_same_format_provider_candidate_payload_parts(
|
let resolved = resolve_local_same_format_provider_candidate_payload_parts(
|
||||||
@@ -83,8 +85,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
|||||||
auth_context: &input.auth_context,
|
auth_context: &input.auth_context,
|
||||||
request_id: trace_id,
|
request_id: trace_id,
|
||||||
candidate_id,
|
candidate_id,
|
||||||
candidate_index: *candidate_index,
|
attempt_identity: attempt.attempt_identity(),
|
||||||
retry_index: 0,
|
|
||||||
model: &input.requested_model,
|
model: &input.requested_model,
|
||||||
provider_name: &resolved.transport.provider.name,
|
provider_name: &resolved.transport.provider.name,
|
||||||
provider_id: &candidate.provider_id,
|
provider_id: &candidate.provider_id,
|
||||||
@@ -94,6 +95,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
|
|||||||
provider_api_format: spec_metadata.api_format,
|
provider_api_format: spec_metadata.api_format,
|
||||||
client_api_format: spec_metadata.api_format,
|
client_api_format: spec_metadata.api_format,
|
||||||
mapped_model: Some(&resolved.mapped_model),
|
mapped_model: Some(&resolved.mapped_model),
|
||||||
|
candidate_group_id: candidate_group_id.as_deref(),
|
||||||
upstream_url: Some(&resolved.upstream_url),
|
upstream_url: Some(&resolved.upstream_url),
|
||||||
provider_request_method: Some(serde_json::Value::Null),
|
provider_request_method: Some(serde_json::Value::Null),
|
||||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||||
|
|||||||
@@ -20,6 +20,7 @@ use crate::handlers::shared::{
|
|||||||
parse_catalog_auth_config_json, provider_key_health_summary,
|
parse_catalog_auth_config_json, provider_key_health_summary,
|
||||||
provider_key_status_snapshot_payload,
|
provider_key_status_snapshot_payload,
|
||||||
};
|
};
|
||||||
|
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||||
use crate::provider_key_auth::provider_key_auth_semantics;
|
use crate::provider_key_auth::provider_key_auth_semantics;
|
||||||
|
|
||||||
const POOL_ACCOUNT_BLOCKED_SKIP_REASON: &str = "pool_account_blocked";
|
const POOL_ACCOUNT_BLOCKED_SKIP_REASON: &str = "pool_account_blocked";
|
||||||
@@ -337,17 +338,27 @@ fn apply_local_execution_pool_scheduler_with_runtime_map(
|
|||||||
let Some(group) = groups.remove(&group_key) else {
|
let Some(group) = groups.remove(&group_key) else {
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
|
let candidate_group_id = local_execution_candidate_group_id(&group_key);
|
||||||
let Some(pool_config) =
|
let Some(pool_config) =
|
||||||
pool_config_for_candidate(group.first().expect("group should exist"))
|
pool_config_for_candidate(group.first().expect("group should exist"))
|
||||||
else {
|
else {
|
||||||
reordered.extend(group);
|
reordered.extend(annotate_local_execution_group_candidates(
|
||||||
|
group,
|
||||||
|
candidate_group_id.as_str(),
|
||||||
|
false,
|
||||||
|
));
|
||||||
continue;
|
continue;
|
||||||
};
|
};
|
||||||
let runtime = runtime_by_provider
|
let runtime = runtime_by_provider
|
||||||
.get(&group_key.provider_id)
|
.get(&group_key.provider_id)
|
||||||
.unwrap_or(&default_runtime);
|
.unwrap_or(&default_runtime);
|
||||||
let (group_candidates, group_skipped) =
|
let (group_candidates, group_skipped) = schedule_pool_group(
|
||||||
schedule_pool_group(group, pool_config, runtime, key_context_by_id);
|
group,
|
||||||
|
pool_config,
|
||||||
|
runtime,
|
||||||
|
key_context_by_id,
|
||||||
|
candidate_group_id.as_str(),
|
||||||
|
);
|
||||||
reordered.extend(group_candidates);
|
reordered.extend(group_candidates);
|
||||||
skipped.extend(group_skipped);
|
skipped.extend(group_skipped);
|
||||||
}
|
}
|
||||||
@@ -366,6 +377,18 @@ fn pool_group_key(candidate: &EligibleLocalExecutionCandidate, pool_enabled: boo
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn local_execution_candidate_group_id(group_key: &PoolGroupKey) -> String {
|
||||||
|
format!(
|
||||||
|
"provider={}|endpoint={}|model={}|selected_model={}|api_format={}|singleton_key={}",
|
||||||
|
group_key.provider_id,
|
||||||
|
group_key.endpoint_id,
|
||||||
|
group_key.model_id,
|
||||||
|
group_key.selected_provider_model_name,
|
||||||
|
group_key.provider_api_format,
|
||||||
|
group_key.singleton_key_id.as_deref().unwrap_or("*"),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
fn pool_config_for_candidate(
|
fn pool_config_for_candidate(
|
||||||
candidate: &EligibleLocalExecutionCandidate,
|
candidate: &EligibleLocalExecutionCandidate,
|
||||||
) -> Option<AdminProviderPoolConfig> {
|
) -> Option<AdminProviderPoolConfig> {
|
||||||
@@ -377,6 +400,7 @@ fn schedule_pool_group(
|
|||||||
pool_config: AdminProviderPoolConfig,
|
pool_config: AdminProviderPoolConfig,
|
||||||
runtime: &AdminProviderPoolRuntimeState,
|
runtime: &AdminProviderPoolRuntimeState,
|
||||||
key_context_by_id: &BTreeMap<String, PoolCatalogKeyContext>,
|
key_context_by_id: &BTreeMap<String, PoolCatalogKeyContext>,
|
||||||
|
candidate_group_id: &str,
|
||||||
) -> (
|
) -> (
|
||||||
Vec<EligibleLocalExecutionCandidate>,
|
Vec<EligibleLocalExecutionCandidate>,
|
||||||
Vec<SkippedLocalExecutionCandidate>,
|
Vec<SkippedLocalExecutionCandidate>,
|
||||||
@@ -403,6 +427,7 @@ fn schedule_pool_group(
|
|||||||
candidate,
|
candidate,
|
||||||
transport,
|
transport,
|
||||||
provider_api_format,
|
provider_api_format,
|
||||||
|
orchestration,
|
||||||
} = eligible;
|
} = eligible;
|
||||||
let key_id = candidate.key_id.clone();
|
let key_id = candidate.key_id.clone();
|
||||||
let mut key_context = key_context_by_id.get(&key_id).cloned().unwrap_or_default();
|
let mut key_context = key_context_by_id.get(&key_id).cloned().unwrap_or_default();
|
||||||
@@ -460,6 +485,7 @@ fn schedule_pool_group(
|
|||||||
candidate,
|
candidate,
|
||||||
transport,
|
transport,
|
||||||
provider_api_format,
|
provider_api_format,
|
||||||
|
orchestration,
|
||||||
},
|
},
|
||||||
key_context,
|
key_context,
|
||||||
original_index,
|
original_index,
|
||||||
@@ -516,7 +542,28 @@ fn schedule_pool_group(
|
|||||||
}
|
}
|
||||||
ordered.extend(available.into_iter().map(|item| item.eligible));
|
ordered.extend(available.into_iter().map(|item| item.eligible));
|
||||||
|
|
||||||
(ordered, skipped)
|
(
|
||||||
|
annotate_local_execution_group_candidates(ordered, candidate_group_id, true),
|
||||||
|
skipped,
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn annotate_local_execution_group_candidates(
|
||||||
|
candidates: Vec<EligibleLocalExecutionCandidate>,
|
||||||
|
candidate_group_id: &str,
|
||||||
|
pool_enabled: bool,
|
||||||
|
) -> Vec<EligibleLocalExecutionCandidate> {
|
||||||
|
candidates
|
||||||
|
.into_iter()
|
||||||
|
.enumerate()
|
||||||
|
.map(|(index, mut candidate)| {
|
||||||
|
candidate.orchestration = LocalExecutionCandidateMetadata {
|
||||||
|
candidate_group_id: Some(candidate_group_id.to_string()),
|
||||||
|
pool_key_index: pool_enabled.then_some(index as u32),
|
||||||
|
};
|
||||||
|
candidate
|
||||||
|
})
|
||||||
|
.collect()
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
@@ -980,6 +1027,7 @@ mod tests {
|
|||||||
use crate::handlers::shared::provider_pool::{
|
use crate::handlers::shared::provider_pool::{
|
||||||
AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,
|
AdminProviderPoolRuntimeState, AdminProviderPoolSchedulingPreset,
|
||||||
};
|
};
|
||||||
|
use crate::orchestration::LocalExecutionCandidateMetadata;
|
||||||
use crate::AppState;
|
use crate::AppState;
|
||||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
||||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||||
@@ -1039,6 +1087,72 @@ mod tests {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn pool_scheduler_attaches_group_and_pool_metadata_to_ranked_candidates() {
|
||||||
|
let pool_first = sample_eligible_candidate(
|
||||||
|
"provider-pool",
|
||||||
|
"endpoint-1",
|
||||||
|
"key-pool-a",
|
||||||
|
10,
|
||||||
|
Some(json!({ "pool_advanced": { "lru_enabled": true } })),
|
||||||
|
);
|
||||||
|
let other =
|
||||||
|
sample_eligible_candidate("provider-other", "endpoint-2", "key-other", 10, None);
|
||||||
|
let pool_second = sample_eligible_candidate(
|
||||||
|
"provider-pool",
|
||||||
|
"endpoint-1",
|
||||||
|
"key-pool-b",
|
||||||
|
10,
|
||||||
|
Some(json!({ "pool_advanced": { "lru_enabled": true } })),
|
||||||
|
);
|
||||||
|
|
||||||
|
let mut runtime_by_provider = BTreeMap::new();
|
||||||
|
runtime_by_provider.insert(
|
||||||
|
"provider-pool".to_string(),
|
||||||
|
AdminProviderPoolRuntimeState {
|
||||||
|
lru_score_by_key: BTreeMap::from([
|
||||||
|
("key-pool-a".to_string(), 20.0),
|
||||||
|
("key-pool-b".to_string(), 10.0),
|
||||||
|
]),
|
||||||
|
..AdminProviderPoolRuntimeState::default()
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
let (reordered, skipped) = apply_local_execution_pool_scheduler_with_runtime_map(
|
||||||
|
vec![pool_first, other, pool_second],
|
||||||
|
&runtime_by_provider,
|
||||||
|
&BTreeMap::new(),
|
||||||
|
);
|
||||||
|
|
||||||
|
assert!(skipped.is_empty());
|
||||||
|
assert_eq!(reordered.len(), 3);
|
||||||
|
assert_eq!(
|
||||||
|
reordered[0].orchestration,
|
||||||
|
LocalExecutionCandidateMetadata {
|
||||||
|
candidate_group_id: Some(
|
||||||
|
"provider=provider-pool|endpoint=endpoint-1|model=model-1|selected_model=gpt-5|api_format=openai:chat|singleton_key=*"
|
||||||
|
.to_string(),
|
||||||
|
),
|
||||||
|
pool_key_index: Some(0),
|
||||||
|
}
|
||||||
|
);
|
||||||
|
assert_eq!(reordered[1].orchestration.pool_key_index, Some(1));
|
||||||
|
assert_eq!(
|
||||||
|
reordered[1].orchestration.candidate_group_id,
|
||||||
|
reordered[0].orchestration.candidate_group_id
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
reordered[2].orchestration,
|
||||||
|
LocalExecutionCandidateMetadata {
|
||||||
|
candidate_group_id: Some(
|
||||||
|
"provider=provider-other|endpoint=endpoint-2|model=model-1|selected_model=gpt-5|api_format=openai:chat|singleton_key=key-other"
|
||||||
|
.to_string(),
|
||||||
|
),
|
||||||
|
pool_key_index: None,
|
||||||
|
}
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn pool_scheduler_promotes_sticky_hit_before_other_sorted_keys() {
|
fn pool_scheduler_promotes_sticky_hit_before_other_sorted_keys() {
|
||||||
let key_a = sample_eligible_candidate(
|
let key_a = sample_eligible_candidate(
|
||||||
@@ -1394,6 +1508,7 @@ mod tests {
|
|||||||
mapping_matched_model: None,
|
mapping_matched_model: None,
|
||||||
},
|
},
|
||||||
provider_api_format: "openai:chat".to_string(),
|
provider_api_format: "openai:chat".to_string(),
|
||||||
|
orchestration: LocalExecutionCandidateMetadata::default(),
|
||||||
transport: crate::ai_pipeline::GatewayProviderTransportSnapshot {
|
transport: crate::ai_pipeline::GatewayProviderTransportSnapshot {
|
||||||
provider: GatewayProviderTransportProvider {
|
provider: GatewayProviderTransportProvider {
|
||||||
id: provider_id.to_string(),
|
id: provider_id.to_string(),
|
||||||
|
|||||||
@@ -3,13 +3,13 @@ use std::collections::BTreeMap;
|
|||||||
use serde_json::{Map, Value};
|
use serde_json::{Map, Value};
|
||||||
|
|
||||||
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
use crate::ai_pipeline::contracts::ExecutionRuntimeAuthContext;
|
||||||
|
use crate::orchestration::ExecutionAttemptIdentity;
|
||||||
|
|
||||||
pub(crate) struct LocalExecutionReportContextParts<'a> {
|
pub(crate) struct LocalExecutionReportContextParts<'a> {
|
||||||
pub(crate) auth_context: &'a ExecutionRuntimeAuthContext,
|
pub(crate) auth_context: &'a ExecutionRuntimeAuthContext,
|
||||||
pub(crate) request_id: &'a str,
|
pub(crate) request_id: &'a str,
|
||||||
pub(crate) candidate_id: &'a str,
|
pub(crate) candidate_id: &'a str,
|
||||||
pub(crate) candidate_index: u32,
|
pub(crate) attempt_identity: ExecutionAttemptIdentity,
|
||||||
pub(crate) retry_index: u32,
|
|
||||||
pub(crate) model: &'a str,
|
pub(crate) model: &'a str,
|
||||||
pub(crate) provider_name: &'a str,
|
pub(crate) provider_name: &'a str,
|
||||||
pub(crate) provider_id: &'a str,
|
pub(crate) provider_id: &'a str,
|
||||||
@@ -19,6 +19,7 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
|
|||||||
pub(crate) provider_api_format: &'a str,
|
pub(crate) provider_api_format: &'a str,
|
||||||
pub(crate) client_api_format: &'a str,
|
pub(crate) client_api_format: &'a str,
|
||||||
pub(crate) mapped_model: Option<&'a str>,
|
pub(crate) mapped_model: Option<&'a str>,
|
||||||
|
pub(crate) candidate_group_id: Option<&'a str>,
|
||||||
pub(crate) upstream_url: Option<&'a str>,
|
pub(crate) upstream_url: Option<&'a str>,
|
||||||
pub(crate) provider_request_method: Option<Value>,
|
pub(crate) provider_request_method: Option<Value>,
|
||||||
pub(crate) provider_request_headers: Option<&'a BTreeMap<String, String>>,
|
pub(crate) provider_request_headers: Option<&'a BTreeMap<String, String>>,
|
||||||
@@ -69,11 +70,11 @@ pub(crate) fn build_local_execution_report_context(
|
|||||||
);
|
);
|
||||||
object.insert(
|
object.insert(
|
||||||
"candidate_index".to_string(),
|
"candidate_index".to_string(),
|
||||||
Value::Number(parts.candidate_index.into()),
|
Value::Number(parts.attempt_identity.candidate_index.into()),
|
||||||
);
|
);
|
||||||
object.insert(
|
object.insert(
|
||||||
"retry_index".to_string(),
|
"retry_index".to_string(),
|
||||||
Value::Number(parts.retry_index.into()),
|
Value::Number(parts.attempt_identity.retry_index.into()),
|
||||||
);
|
);
|
||||||
object.insert("model".to_string(), Value::String(parts.model.to_string()));
|
object.insert("model".to_string(), Value::String(parts.model.to_string()));
|
||||||
object.insert(
|
object.insert(
|
||||||
@@ -127,6 +128,12 @@ pub(crate) fn build_local_execution_report_context(
|
|||||||
Value::String(mapped_model.to_string()),
|
Value::String(mapped_model.to_string()),
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
if let Some(candidate_group_id) = parts.candidate_group_id {
|
||||||
|
object.insert(
|
||||||
|
"candidate_group_id".to_string(),
|
||||||
|
Value::String(candidate_group_id.to_string()),
|
||||||
|
);
|
||||||
|
}
|
||||||
if let Some(upstream_url) = parts.upstream_url {
|
if let Some(upstream_url) = parts.upstream_url {
|
||||||
object.insert(
|
object.insert(
|
||||||
"upstream_url".to_string(),
|
"upstream_url".to_string(),
|
||||||
@@ -146,6 +153,12 @@ pub(crate) fn build_local_execution_report_context(
|
|||||||
.expect("provider request headers should serialize"),
|
.expect("provider request headers should serialize"),
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
if let Some(pool_key_index) = parts.attempt_identity.pool_key_index {
|
||||||
|
object.insert(
|
||||||
|
"pool_key_index".to_string(),
|
||||||
|
Value::Number(pool_key_index.into()),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
object.extend(parts.extra_fields);
|
object.extend(parts.extra_fields);
|
||||||
Value::Object(object)
|
Value::Object(object)
|
||||||
|
|||||||
@@ -34,6 +34,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
|||||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||||
let spec_metadata = local_gemini_files_spec_metadata(spec);
|
let spec_metadata = local_gemini_files_spec_metadata(spec);
|
||||||
let planner_state = PlannerAppState::new(state);
|
let planner_state = PlannerAppState::new(state);
|
||||||
|
let attempt_identity = attempt.attempt_identity();
|
||||||
let resolved = resolve_local_gemini_files_candidate_payload_parts(
|
let resolved = resolve_local_gemini_files_candidate_payload_parts(
|
||||||
state,
|
state,
|
||||||
parts,
|
parts,
|
||||||
@@ -48,8 +49,9 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
|||||||
.await?;
|
.await?;
|
||||||
let LocalGeminiFilesCandidateAttempt {
|
let LocalGeminiFilesCandidateAttempt {
|
||||||
eligible,
|
eligible,
|
||||||
candidate_index,
|
candidate_group_id,
|
||||||
candidate_id,
|
candidate_id,
|
||||||
|
..
|
||||||
} = attempt;
|
} = attempt;
|
||||||
let candidate = eligible.candidate;
|
let candidate = eligible.candidate;
|
||||||
let transport = resolved.transport;
|
let transport = resolved.transport;
|
||||||
@@ -107,8 +109,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
|||||||
auth_context: &input.auth_context,
|
auth_context: &input.auth_context,
|
||||||
request_id: trace_id,
|
request_id: trace_id,
|
||||||
candidate_id: &candidate_id,
|
candidate_id: &candidate_id,
|
||||||
candidate_index,
|
attempt_identity,
|
||||||
retry_index: 0,
|
|
||||||
model: "gemini-files",
|
model: "gemini-files",
|
||||||
provider_name: &transport.provider.name,
|
provider_name: &transport.provider.name,
|
||||||
provider_id: &candidate.provider_id,
|
provider_id: &candidate.provider_id,
|
||||||
@@ -118,6 +119,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
|
|||||||
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
|
provider_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
|
||||||
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
|
client_api_format: GEMINI_FILES_CLIENT_API_FORMAT,
|
||||||
mapped_model: None,
|
mapped_model: None,
|
||||||
|
candidate_group_id: candidate_group_id.as_deref(),
|
||||||
upstream_url: None,
|
upstream_url: None,
|
||||||
provider_request_method: None,
|
provider_request_method: None,
|
||||||
provider_request_headers: None,
|
provider_request_headers: None,
|
||||||
|
|||||||
@@ -105,6 +105,7 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
|
|||||||
None,
|
None,
|
||||||
&candidates,
|
&candidates,
|
||||||
);
|
);
|
||||||
|
let available_candidate_count = candidates.len() as u32;
|
||||||
let attempts = persist_available_local_execution_candidates_with_context(
|
let attempts = persist_available_local_execution_candidates_with_context(
|
||||||
planner_state,
|
planner_state,
|
||||||
trace_id,
|
trace_id,
|
||||||
@@ -132,7 +133,7 @@ pub(super) async fn materialize_local_gemini_files_candidate_attempts(
|
|||||||
state,
|
state,
|
||||||
trace_id,
|
trace_id,
|
||||||
persistence_policy.skipped,
|
persistence_policy.skipped,
|
||||||
attempts.len() as u32,
|
available_candidate_count,
|
||||||
skipped_candidates
|
skipped_candidates
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|mut skipped_candidate| {
|
.map(|mut skipped_candidate| {
|
||||||
|
|||||||
@@ -27,14 +27,16 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
|||||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||||
let spec_metadata = local_video_create_spec_metadata(spec);
|
let spec_metadata = local_video_create_spec_metadata(spec);
|
||||||
let planner_state = PlannerAppState::new(state);
|
let planner_state = PlannerAppState::new(state);
|
||||||
|
let attempt_identity = attempt.attempt_identity();
|
||||||
let resolved = resolve_local_video_create_candidate_payload_parts(
|
let resolved = resolve_local_video_create_candidate_payload_parts(
|
||||||
state, parts, body_json, trace_id, input, &attempt, spec,
|
state, parts, body_json, trace_id, input, &attempt, spec,
|
||||||
)
|
)
|
||||||
.await?;
|
.await?;
|
||||||
let LocalVideoCreateCandidateAttempt {
|
let LocalVideoCreateCandidateAttempt {
|
||||||
eligible,
|
eligible,
|
||||||
candidate_index,
|
candidate_group_id,
|
||||||
candidate_id,
|
candidate_id,
|
||||||
|
..
|
||||||
} = attempt;
|
} = attempt;
|
||||||
let candidate = eligible.candidate;
|
let candidate = eligible.candidate;
|
||||||
let transport = resolved.transport;
|
let transport = resolved.transport;
|
||||||
@@ -90,8 +92,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
|||||||
auth_context: &input.auth_context,
|
auth_context: &input.auth_context,
|
||||||
request_id: trace_id,
|
request_id: trace_id,
|
||||||
candidate_id: &candidate_id,
|
candidate_id: &candidate_id,
|
||||||
candidate_index,
|
attempt_identity,
|
||||||
retry_index: 0,
|
|
||||||
model: &input.requested_model,
|
model: &input.requested_model,
|
||||||
provider_name: &transport.provider.name,
|
provider_name: &transport.provider.name,
|
||||||
provider_id: &candidate.provider_id,
|
provider_id: &candidate.provider_id,
|
||||||
@@ -101,6 +102,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
|
|||||||
provider_api_format: spec_metadata.api_format,
|
provider_api_format: spec_metadata.api_format,
|
||||||
client_api_format: spec_metadata.api_format,
|
client_api_format: spec_metadata.api_format,
|
||||||
mapped_model: Some(&resolved.mapped_model),
|
mapped_model: Some(&resolved.mapped_model),
|
||||||
|
candidate_group_id: candidate_group_id.as_deref(),
|
||||||
upstream_url: None,
|
upstream_url: None,
|
||||||
provider_request_method: None,
|
provider_request_method: None,
|
||||||
provider_request_headers: None,
|
provider_request_headers: None,
|
||||||
|
|||||||
@@ -182,6 +182,7 @@ async fn materialize_local_video_create_candidate_attempts(
|
|||||||
Some(&input.requested_model),
|
Some(&input.requested_model),
|
||||||
&candidates,
|
&candidates,
|
||||||
);
|
);
|
||||||
|
let available_candidate_count = candidates.len() as u32;
|
||||||
let attempts = persist_available_local_execution_candidates_with_context(
|
let attempts = persist_available_local_execution_candidates_with_context(
|
||||||
state,
|
state,
|
||||||
trace_id,
|
trace_id,
|
||||||
@@ -204,7 +205,7 @@ async fn materialize_local_video_create_candidate_attempts(
|
|||||||
state.app(),
|
state.app(),
|
||||||
trace_id,
|
trace_id,
|
||||||
persistence_policy.skipped,
|
persistence_policy.skipped,
|
||||||
attempts.len() as u32,
|
available_candidate_count,
|
||||||
skipped_candidates
|
skipped_candidates
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|mut skipped_candidate| {
|
.map(|mut skipped_candidate| {
|
||||||
|
|||||||
@@ -234,6 +234,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
|||||||
Some(&input.requested_model),
|
Some(&input.requested_model),
|
||||||
&candidates,
|
&candidates,
|
||||||
);
|
);
|
||||||
|
let available_candidate_count = candidates.len() as u32;
|
||||||
let attempts = persist_available_local_execution_candidates_with_context(
|
let attempts = persist_available_local_execution_candidates_with_context(
|
||||||
planner_state,
|
planner_state,
|
||||||
trace_id,
|
trace_id,
|
||||||
@@ -273,7 +274,7 @@ pub(super) async fn materialize_local_standard_candidate_attempts(
|
|||||||
state,
|
state,
|
||||||
trace_id,
|
trace_id,
|
||||||
persistence_policy.skipped,
|
persistence_policy.skipped,
|
||||||
attempts.len() as u32,
|
available_candidate_count,
|
||||||
skipped_candidates,
|
skipped_candidates,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|||||||
@@ -35,7 +35,9 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
|||||||
let LocalStandardCandidateAttempt {
|
let LocalStandardCandidateAttempt {
|
||||||
eligible,
|
eligible,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
|
candidate_group_id,
|
||||||
candidate_id,
|
candidate_id,
|
||||||
|
..
|
||||||
} = &attempt;
|
} = &attempt;
|
||||||
let candidate = &eligible.candidate;
|
let candidate = &eligible.candidate;
|
||||||
let resolved = resolve_local_standard_candidate_payload_parts(
|
let resolved = resolve_local_standard_candidate_payload_parts(
|
||||||
@@ -89,8 +91,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
|||||||
auth_context: &input.auth_context,
|
auth_context: &input.auth_context,
|
||||||
request_id: trace_id,
|
request_id: trace_id,
|
||||||
candidate_id,
|
candidate_id,
|
||||||
candidate_index: *candidate_index,
|
attempt_identity: attempt.attempt_identity(),
|
||||||
retry_index: 0,
|
|
||||||
model: &input.requested_model,
|
model: &input.requested_model,
|
||||||
provider_name: &candidate.provider_name,
|
provider_name: &candidate.provider_name,
|
||||||
provider_id: &candidate.provider_id,
|
provider_id: &candidate.provider_id,
|
||||||
@@ -100,6 +101,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
|
|||||||
provider_api_format: &resolved.provider_api_format,
|
provider_api_format: &resolved.provider_api_format,
|
||||||
client_api_format: spec_metadata.api_format,
|
client_api_format: spec_metadata.api_format,
|
||||||
mapped_model: Some(&resolved.mapped_model),
|
mapped_model: Some(&resolved.mapped_model),
|
||||||
|
candidate_group_id: candidate_group_id.as_deref(),
|
||||||
upstream_url: Some(&resolved.upstream_url),
|
upstream_url: Some(&resolved.upstream_url),
|
||||||
provider_request_method: Some(serde_json::Value::Null),
|
provider_request_method: Some(serde_json::Value::Null),
|
||||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||||
|
|||||||
@@ -28,10 +28,13 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
|||||||
report_kind: &str,
|
report_kind: &str,
|
||||||
upstream_is_stream: bool,
|
upstream_is_stream: bool,
|
||||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||||
|
let attempt_identity = attempt.attempt_identity();
|
||||||
let LocalOpenAiChatCandidateAttempt {
|
let LocalOpenAiChatCandidateAttempt {
|
||||||
eligible,
|
eligible,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
|
candidate_group_id,
|
||||||
candidate_id,
|
candidate_id,
|
||||||
|
..
|
||||||
} = attempt;
|
} = attempt;
|
||||||
let resolved = resolve_local_openai_chat_candidate_payload_parts(
|
let resolved = resolve_local_openai_chat_candidate_payload_parts(
|
||||||
state,
|
state,
|
||||||
@@ -105,8 +108,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
|||||||
auth_context: &input.auth_context,
|
auth_context: &input.auth_context,
|
||||||
request_id: trace_id,
|
request_id: trace_id,
|
||||||
candidate_id: &candidate_id,
|
candidate_id: &candidate_id,
|
||||||
candidate_index,
|
attempt_identity,
|
||||||
retry_index: 0,
|
|
||||||
model: &input.requested_model,
|
model: &input.requested_model,
|
||||||
provider_name: &resolved.transport.provider.name,
|
provider_name: &resolved.transport.provider.name,
|
||||||
provider_id: &candidate.provider_id,
|
provider_id: &candidate.provider_id,
|
||||||
@@ -116,6 +118,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
|
|||||||
provider_api_format: &resolved.provider_api_format,
|
provider_api_format: &resolved.provider_api_format,
|
||||||
client_api_format: "openai:chat",
|
client_api_format: "openai:chat",
|
||||||
mapped_model: Some(&resolved.mapped_model),
|
mapped_model: Some(&resolved.mapped_model),
|
||||||
|
candidate_group_id: candidate_group_id.as_deref(),
|
||||||
upstream_url: Some(&resolved.upstream_url),
|
upstream_url: Some(&resolved.upstream_url),
|
||||||
provider_request_method: Some(serde_json::Value::Null),
|
provider_request_method: Some(serde_json::Value::Null),
|
||||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||||
|
|||||||
@@ -122,6 +122,7 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
|
|||||||
Some(&input.requested_model),
|
Some(&input.requested_model),
|
||||||
&candidates,
|
&candidates,
|
||||||
);
|
);
|
||||||
|
let available_candidate_count = candidates.len() as u32;
|
||||||
let attempts = persist_available_local_execution_candidates_with_context(
|
let attempts = persist_available_local_execution_candidates_with_context(
|
||||||
planner_state,
|
planner_state,
|
||||||
trace_id,
|
trace_id,
|
||||||
@@ -156,7 +157,7 @@ pub(crate) async fn materialize_local_openai_chat_candidate_attempts(
|
|||||||
state,
|
state,
|
||||||
trace_id,
|
trace_id,
|
||||||
persistence_policy.skipped,
|
persistence_policy.skipped,
|
||||||
attempts.len() as u32,
|
available_candidate_count,
|
||||||
skipped_candidates,
|
skipped_candidates,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|||||||
@@ -31,10 +31,13 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
|
|||||||
spec: LocalOpenAiCliSpec,
|
spec: LocalOpenAiCliSpec,
|
||||||
) -> Option<GatewayControlSyncDecisionResponse> {
|
) -> Option<GatewayControlSyncDecisionResponse> {
|
||||||
let spec_metadata = local_openai_cli_spec_metadata(spec);
|
let spec_metadata = local_openai_cli_spec_metadata(spec);
|
||||||
|
let attempt_identity = attempt.attempt_identity();
|
||||||
let LocalOpenAiCliCandidateAttempt {
|
let LocalOpenAiCliCandidateAttempt {
|
||||||
eligible,
|
eligible,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
|
candidate_group_id,
|
||||||
candidate_id,
|
candidate_id,
|
||||||
|
..
|
||||||
} = attempt;
|
} = attempt;
|
||||||
let resolved = resolve_local_openai_cli_candidate_payload_parts(
|
let resolved = resolve_local_openai_cli_candidate_payload_parts(
|
||||||
state,
|
state,
|
||||||
@@ -133,8 +136,7 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
|
|||||||
auth_context: &input.auth_context,
|
auth_context: &input.auth_context,
|
||||||
request_id: trace_id,
|
request_id: trace_id,
|
||||||
candidate_id: &candidate_id,
|
candidate_id: &candidate_id,
|
||||||
candidate_index,
|
attempt_identity,
|
||||||
retry_index: 0,
|
|
||||||
model: &input.requested_model,
|
model: &input.requested_model,
|
||||||
provider_name: &resolved.transport.provider.name,
|
provider_name: &resolved.transport.provider.name,
|
||||||
provider_id: &candidate.provider_id,
|
provider_id: &candidate.provider_id,
|
||||||
@@ -144,6 +146,7 @@ pub(crate) async fn maybe_build_local_openai_cli_decision_payload_for_candidate(
|
|||||||
provider_api_format: &resolved.provider_api_format,
|
provider_api_format: &resolved.provider_api_format,
|
||||||
client_api_format: spec_metadata.api_format,
|
client_api_format: spec_metadata.api_format,
|
||||||
mapped_model: Some(&resolved.mapped_model),
|
mapped_model: Some(&resolved.mapped_model),
|
||||||
|
candidate_group_id: candidate_group_id.as_deref(),
|
||||||
upstream_url: Some(&resolved.upstream_url),
|
upstream_url: Some(&resolved.upstream_url),
|
||||||
provider_request_method: Some(serde_json::Value::Null),
|
provider_request_method: Some(serde_json::Value::Null),
|
||||||
provider_request_headers: Some(&resolved.provider_request_headers),
|
provider_request_headers: Some(&resolved.provider_request_headers),
|
||||||
|
|||||||
@@ -230,6 +230,7 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
|
|||||||
Some(&input.requested_model),
|
Some(&input.requested_model),
|
||||||
&candidates,
|
&candidates,
|
||||||
);
|
);
|
||||||
|
let available_candidate_count = candidates.len() as u32;
|
||||||
let attempts = persist_available_local_execution_candidates_with_context(
|
let attempts = persist_available_local_execution_candidates_with_context(
|
||||||
planner_state,
|
planner_state,
|
||||||
trace_id,
|
trace_id,
|
||||||
@@ -269,7 +270,7 @@ pub(crate) async fn materialize_local_openai_cli_candidate_attempts(
|
|||||||
state,
|
state,
|
||||||
trace_id,
|
trace_id,
|
||||||
persistence_policy.skipped,
|
persistence_policy.skipped,
|
||||||
attempts.len() as u32,
|
available_candidate_count,
|
||||||
skipped_candidates,
|
skipped_candidates,
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ impl<'a> PlannerAppState<'a> {
|
|||||||
api_key_id: &str,
|
api_key_id: &str,
|
||||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||||
candidate_index: u32,
|
candidate_index: u32,
|
||||||
|
retry_index: u32,
|
||||||
candidate_id: &str,
|
candidate_id: &str,
|
||||||
required_capabilities: Option<&Value>,
|
required_capabilities: Option<&Value>,
|
||||||
extra_data: Option<Value>,
|
extra_data: Option<Value>,
|
||||||
@@ -42,6 +43,7 @@ impl<'a> PlannerAppState<'a> {
|
|||||||
api_key_id,
|
api_key_id,
|
||||||
candidate,
|
candidate,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
|
retry_index,
|
||||||
candidate_id,
|
candidate_id,
|
||||||
required_capabilities,
|
required_capabilities,
|
||||||
extra_data,
|
extra_data,
|
||||||
@@ -59,6 +61,7 @@ impl<'a> PlannerAppState<'a> {
|
|||||||
api_key_id: &str,
|
api_key_id: &str,
|
||||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||||
candidate_index: u32,
|
candidate_index: u32,
|
||||||
|
retry_index: u32,
|
||||||
candidate_id: &str,
|
candidate_id: &str,
|
||||||
required_capabilities: Option<&Value>,
|
required_capabilities: Option<&Value>,
|
||||||
skip_reason: &str,
|
skip_reason: &str,
|
||||||
@@ -73,6 +76,7 @@ impl<'a> PlannerAppState<'a> {
|
|||||||
api_key_id,
|
api_key_id,
|
||||||
candidate,
|
candidate,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
|
retry_index,
|
||||||
candidate_id,
|
candidate_id,
|
||||||
required_capabilities,
|
required_capabilities,
|
||||||
skip_reason,
|
skip_reason,
|
||||||
|
|||||||
@@ -1,24 +1,11 @@
|
|||||||
use std::collections::BTreeSet;
|
|
||||||
|
|
||||||
use aether_contracts::{ExecutionPlan, ExecutionResult};
|
use aether_contracts::{ExecutionPlan, ExecutionResult};
|
||||||
use regex::Regex;
|
use serde_json::Value;
|
||||||
use serde_json::{json, Value};
|
|
||||||
use tracing::debug;
|
|
||||||
|
|
||||||
use crate::provider_transport::GatewayProviderTransportSnapshot;
|
use crate::orchestration::{
|
||||||
|
resolve_local_failover_analysis_for_attempt, LocalFailoverAnalysis, LocalFailoverDecision,
|
||||||
|
};
|
||||||
use crate::AppState;
|
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 {
|
fn sync_plan_kind_disables_local_candidate_failover(plan_kind: &str) -> bool {
|
||||||
matches!(
|
matches!(
|
||||||
plan_kind,
|
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(
|
pub(crate) async fn should_retry_next_local_candidate_sync(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
plan: &ExecutionPlan,
|
plan: &ExecutionPlan,
|
||||||
@@ -66,22 +21,43 @@ pub(crate) async fn should_retry_next_local_candidate_sync(
|
|||||||
result: &ExecutionResult,
|
result: &ExecutionResult,
|
||||||
response_text: Option<&str>,
|
response_text: Option<&str>,
|
||||||
) -> bool {
|
) -> bool {
|
||||||
if sync_plan_kind_disables_local_candidate_failover(plan_kind) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
matches!(
|
matches!(
|
||||||
resolve_local_failover_decision(
|
analyze_local_candidate_failover_sync(
|
||||||
state,
|
state,
|
||||||
plan,
|
plan,
|
||||||
|
plan_kind,
|
||||||
report_context,
|
report_context,
|
||||||
result.status_code,
|
result,
|
||||||
response_text,
|
response_text,
|
||||||
)
|
)
|
||||||
.await,
|
.await
|
||||||
|
.decision,
|
||||||
LocalFailoverDecision::RetryNextCandidate
|
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(
|
pub(crate) async fn should_stop_local_candidate_failover_sync(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
plan: &ExecutionPlan,
|
plan: &ExecutionPlan,
|
||||||
@@ -90,19 +66,20 @@ pub(crate) async fn should_stop_local_candidate_failover_sync(
|
|||||||
result: &ExecutionResult,
|
result: &ExecutionResult,
|
||||||
response_text: Option<&str>,
|
response_text: Option<&str>,
|
||||||
) -> bool {
|
) -> bool {
|
||||||
if sync_plan_kind_disables_local_candidate_failover(plan_kind) {
|
|
||||||
return false;
|
|
||||||
}
|
|
||||||
matches!(
|
matches!(
|
||||||
resolve_local_failover_decision(
|
analyze_local_candidate_failover_sync(
|
||||||
state,
|
state,
|
||||||
plan,
|
plan,
|
||||||
|
plan_kind,
|
||||||
report_context,
|
report_context,
|
||||||
result.status_code,
|
result,
|
||||||
response_text,
|
response_text,
|
||||||
)
|
)
|
||||||
.await,
|
.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>,
|
response_text: Option<&str>,
|
||||||
) -> bool {
|
) -> bool {
|
||||||
matches!(
|
matches!(
|
||||||
resolve_local_candidate_failover_decision_stream(
|
resolve_local_candidate_failover_analysis_stream(
|
||||||
state,
|
state,
|
||||||
plan,
|
plan,
|
||||||
report_context,
|
report_context,
|
||||||
@@ -203,7 +180,10 @@ pub(crate) async fn should_retry_next_local_candidate_stream(
|
|||||||
response_text,
|
response_text,
|
||||||
)
|
)
|
||||||
.await,
|
.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>,
|
response_text: Option<&str>,
|
||||||
) -> bool {
|
) -> bool {
|
||||||
matches!(
|
matches!(
|
||||||
resolve_local_candidate_failover_decision_stream(
|
resolve_local_candidate_failover_analysis_stream(
|
||||||
state,
|
state,
|
||||||
plan,
|
plan,
|
||||||
report_context,
|
report_context,
|
||||||
@@ -224,10 +204,30 @@ pub(crate) async fn should_stop_local_candidate_failover_stream(
|
|||||||
response_text,
|
response_text,
|
||||||
)
|
)
|
||||||
.await,
|
.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(
|
pub(crate) async fn resolve_local_candidate_failover_decision_stream(
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
plan: &ExecutionPlan,
|
plan: &ExecutionPlan,
|
||||||
@@ -235,7 +235,15 @@ pub(crate) async fn resolve_local_candidate_failover_decision_stream(
|
|||||||
status_code: u16,
|
status_code: u16,
|
||||||
response_text: Option<&str>,
|
response_text: Option<&str>,
|
||||||
) -> LocalFailoverDecision {
|
) -> 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(
|
pub(crate) fn local_failover_response_text(
|
||||||
@@ -255,304 +263,6 @@ pub(crate) fn local_failover_response_text(
|
|||||||
.map(ToOwned::to_owned)
|
.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(
|
pub(crate) fn should_fallback_to_control_stream(
|
||||||
plan_kind: &str,
|
plan_kind: &str,
|
||||||
status_code: u16,
|
status_code: u16,
|
||||||
@@ -623,13 +333,15 @@ mod tests {
|
|||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
resolve_core_stream_error_finalize_report_kind,
|
resolve_core_stream_error_finalize_report_kind,
|
||||||
resolve_core_sync_error_finalize_report_kind, resolve_local_failover_policy,
|
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_stream,
|
||||||
should_fallback_to_control_stream, should_fallback_to_control_sync,
|
should_fallback_to_control_sync, should_retry_next_local_candidate_stream,
|
||||||
should_retry_next_local_candidate_stream, should_retry_next_local_candidate_sync,
|
should_retry_next_local_candidate_sync, should_stop_local_candidate_failover_stream,
|
||||||
should_stop_local_candidate_failover_stream, should_stop_local_candidate_failover_sync,
|
should_stop_local_candidate_failover_sync,
|
||||||
LocalFailoverPolicy, LocalFailoverRegexRule,
|
|
||||||
};
|
};
|
||||||
use crate::data::GatewayDataState;
|
use crate::data::GatewayDataState;
|
||||||
|
use crate::orchestration::{
|
||||||
|
resolve_local_failover_policy, LocalFailoverPolicy, LocalFailoverRegexRule,
|
||||||
|
};
|
||||||
use crate::AppState;
|
use crate::AppState;
|
||||||
|
|
||||||
fn sample_plan() -> aether_contracts::ExecutionPlan {
|
fn sample_plan() -> aether_contracts::ExecutionPlan {
|
||||||
@@ -1122,7 +834,7 @@ mod tests {
|
|||||||
.await
|
.await
|
||||||
);
|
);
|
||||||
assert!(
|
assert!(
|
||||||
!should_retry_next_local_candidate_stream(
|
should_retry_next_local_candidate_stream(
|
||||||
&state,
|
&state,
|
||||||
&plan,
|
&plan,
|
||||||
"openai_chat_stream",
|
"openai_chat_stream",
|
||||||
|
|||||||
@@ -6,7 +6,6 @@ use serde_json::{Map, Value};
|
|||||||
mod constants;
|
mod constants;
|
||||||
mod fallback;
|
mod fallback;
|
||||||
pub(crate) mod ndjson;
|
pub(crate) mod ndjson;
|
||||||
mod pool_feedback;
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
pub(crate) mod remote_compat;
|
pub(crate) mod remote_compat;
|
||||||
mod server;
|
mod server;
|
||||||
@@ -20,18 +19,17 @@ pub(crate) use self::constants::{
|
|||||||
MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES,
|
MAX_ERROR_BODY_BYTES, MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES,
|
||||||
};
|
};
|
||||||
pub(crate) use self::fallback::{
|
pub(crate) use self::fallback::{
|
||||||
append_local_failover_policy_to_value, local_failover_response_text,
|
analyze_local_candidate_failover_sync, local_failover_response_text,
|
||||||
resolve_core_stream_direct_finalize_report_kind,
|
resolve_core_stream_direct_finalize_report_kind,
|
||||||
resolve_core_stream_error_finalize_report_kind, resolve_core_sync_error_finalize_report_kind,
|
resolve_core_stream_error_finalize_report_kind, resolve_core_sync_error_finalize_report_kind,
|
||||||
|
resolve_local_candidate_failover_analysis_stream,
|
||||||
resolve_local_candidate_failover_decision_stream, should_fallback_to_control_stream,
|
resolve_local_candidate_failover_decision_stream, should_fallback_to_control_stream,
|
||||||
should_fallback_to_control_sync, should_finalize_sync_response,
|
should_fallback_to_control_sync, should_finalize_sync_response,
|
||||||
should_retry_next_local_candidate_stream, should_retry_next_local_candidate_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,
|
should_stop_local_candidate_failover_stream, should_stop_local_candidate_failover_sync,
|
||||||
LocalFailoverDecision,
|
|
||||||
};
|
};
|
||||||
pub(crate) use pool_feedback::{
|
pub(crate) use crate::orchestration::{
|
||||||
record_pool_error_feedback, record_pool_stream_timeout_feedback,
|
append_local_failover_policy_to_value, LocalFailoverAnalysis, LocalFailoverDecision,
|
||||||
record_stream_pool_success_feedback, record_sync_pool_success_feedback,
|
|
||||||
};
|
};
|
||||||
pub use server::{
|
pub use server::{
|
||||||
build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit,
|
build_execution_runtime_router, build_execution_runtime_router_with_request_concurrency_limit,
|
||||||
|
|||||||
@@ -1,188 +0,0 @@
|
|||||||
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
|
|
||||||
use aether_usage_runtime::{
|
|
||||||
build_stream_terminal_usage_outcome, build_sync_terminal_usage_outcome, TerminalUsageOutcome,
|
|
||||||
};
|
|
||||||
use serde_json::Value;
|
|
||||||
use std::collections::BTreeMap;
|
|
||||||
use tracing::warn;
|
|
||||||
|
|
||||||
use crate::ai_pipeline::extract_pool_sticky_session_token;
|
|
||||||
use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_config_value;
|
|
||||||
use crate::handlers::shared::provider_pool::{
|
|
||||||
record_admin_provider_pool_error, record_admin_provider_pool_stream_timeout,
|
|
||||||
record_admin_provider_pool_success, AdminProviderPoolConfig,
|
|
||||||
};
|
|
||||||
use crate::usage::{GatewayStreamReportRequest, GatewaySyncReportRequest};
|
|
||||||
use crate::AppState;
|
|
||||||
|
|
||||||
struct PoolFeedbackContext {
|
|
||||||
runner: aether_data::redis::RedisKvRunner,
|
|
||||||
pool_config: AdminProviderPoolConfig,
|
|
||||||
sticky_session_token: Option<String>,
|
|
||||||
}
|
|
||||||
|
|
||||||
fn pool_feedback_request_body<'a>(
|
|
||||||
plan: &'a ExecutionPlan,
|
|
||||||
report_context: Option<&'a Value>,
|
|
||||||
) -> Option<&'a Value> {
|
|
||||||
report_context
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.and_then(|object| object.get("original_request_body"))
|
|
||||||
.filter(|value| !value.is_null())
|
|
||||||
.or(plan.body.json_body.as_ref())
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn resolve_pool_feedback_context(
|
|
||||||
state: &AppState,
|
|
||||||
plan: &ExecutionPlan,
|
|
||||||
report_context: Option<&Value>,
|
|
||||||
) -> Option<PoolFeedbackContext> {
|
|
||||||
let Some(runner) = state.redis_kv_runner() else {
|
|
||||||
return None;
|
|
||||||
};
|
|
||||||
|
|
||||||
let transport = match state
|
|
||||||
.read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
Ok(Some(transport)) => transport,
|
|
||||||
Ok(None) => return None,
|
|
||||||
Err(err) => {
|
|
||||||
warn!(
|
|
||||||
"gateway execution runtime pool feedback: failed to read transport snapshot for provider {} endpoint {} key {}: {:?}",
|
|
||||||
plan.provider_id, plan.endpoint_id, plan.key_id, err
|
|
||||||
);
|
|
||||||
return None;
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
let Some(pool_config) =
|
|
||||||
admin_provider_pool_config_from_config_value(transport.provider.config.as_ref())
|
|
||||||
else {
|
|
||||||
return None;
|
|
||||||
};
|
|
||||||
|
|
||||||
let sticky_session_token = pool_feedback_request_body(plan, report_context)
|
|
||||||
.and_then(extract_pool_sticky_session_token);
|
|
||||||
|
|
||||||
Some(PoolFeedbackContext {
|
|
||||||
runner,
|
|
||||||
pool_config,
|
|
||||||
sticky_session_token,
|
|
||||||
})
|
|
||||||
}
|
|
||||||
|
|
||||||
fn total_tokens_used(outcome: &TerminalUsageOutcome) -> u64 {
|
|
||||||
outcome
|
|
||||||
.standardized_usage
|
|
||||||
.as_ref()
|
|
||||||
.map(|usage| {
|
|
||||||
usage
|
|
||||||
.input_tokens
|
|
||||||
.saturating_add(usage.output_tokens)
|
|
||||||
.max(0) as u64
|
|
||||||
})
|
|
||||||
.unwrap_or(0)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn resolve_ttfb_ms(telemetry: Option<&ExecutionTelemetry>) -> Option<u64> {
|
|
||||||
telemetry.and_then(|telemetry| telemetry.ttfb_ms.or(telemetry.elapsed_ms))
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(crate) async fn record_sync_pool_success_feedback(
|
|
||||||
state: &AppState,
|
|
||||||
plan: &ExecutionPlan,
|
|
||||||
report_context: Option<&Value>,
|
|
||||||
payload: &GatewaySyncReportRequest,
|
|
||||||
) {
|
|
||||||
let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
|
|
||||||
let usage_outcome = build_sync_terminal_usage_outcome(plan, report_context, payload);
|
|
||||||
record_admin_provider_pool_success(
|
|
||||||
&context.runner,
|
|
||||||
&plan.provider_id,
|
|
||||||
&plan.key_id,
|
|
||||||
&context.pool_config,
|
|
||||||
context.sticky_session_token.as_deref(),
|
|
||||||
total_tokens_used(&usage_outcome),
|
|
||||||
resolve_ttfb_ms(payload.telemetry.as_ref()),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(crate) async fn record_stream_pool_success_feedback(
|
|
||||||
state: &AppState,
|
|
||||||
plan: &ExecutionPlan,
|
|
||||||
report_context: Option<&Value>,
|
|
||||||
payload: &GatewayStreamReportRequest,
|
|
||||||
) {
|
|
||||||
let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
|
|
||||||
let usage_outcome = build_stream_terminal_usage_outcome(plan, report_context, payload);
|
|
||||||
record_admin_provider_pool_success(
|
|
||||||
&context.runner,
|
|
||||||
&plan.provider_id,
|
|
||||||
&plan.key_id,
|
|
||||||
&context.pool_config,
|
|
||||||
context.sticky_session_token.as_deref(),
|
|
||||||
total_tokens_used(&usage_outcome),
|
|
||||||
resolve_ttfb_ms(payload.telemetry.as_ref()),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(crate) async fn record_pool_error_feedback(
|
|
||||||
state: &AppState,
|
|
||||||
plan: &ExecutionPlan,
|
|
||||||
report_context: Option<&Value>,
|
|
||||||
status_code: u16,
|
|
||||||
headers: &BTreeMap<String, String>,
|
|
||||||
error_body: Option<&str>,
|
|
||||||
) {
|
|
||||||
if status_code < 400 {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
|
|
||||||
if status_code == 401 {
|
|
||||||
let _ = state
|
|
||||||
.invalidate_local_oauth_refresh_entry(&plan.key_id)
|
|
||||||
.await;
|
|
||||||
}
|
|
||||||
|
|
||||||
record_admin_provider_pool_error(
|
|
||||||
&context.runner,
|
|
||||||
&plan.provider_id,
|
|
||||||
&plan.key_id,
|
|
||||||
&context.pool_config,
|
|
||||||
status_code,
|
|
||||||
error_body,
|
|
||||||
Some(headers),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(crate) async fn record_pool_stream_timeout_feedback(
|
|
||||||
state: &AppState,
|
|
||||||
plan: &ExecutionPlan,
|
|
||||||
report_context: Option<&Value>,
|
|
||||||
) {
|
|
||||||
let Some(context) = resolve_pool_feedback_context(state, plan, report_context).await else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
|
|
||||||
record_admin_provider_pool_stream_timeout(
|
|
||||||
&context.runner,
|
|
||||||
&plan.provider_id,
|
|
||||||
&plan.key_id,
|
|
||||||
&context.pool_config,
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
}
|
|
||||||
@@ -53,14 +53,18 @@ use crate::execution_runtime::transport::{
|
|||||||
DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
|
DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
|
||||||
};
|
};
|
||||||
use crate::execution_runtime::{
|
use crate::execution_runtime::{
|
||||||
local_failover_response_text, record_pool_error_feedback, record_stream_pool_success_feedback,
|
local_failover_response_text, resolve_core_stream_direct_finalize_report_kind,
|
||||||
resolve_core_stream_direct_finalize_report_kind,
|
|
||||||
resolve_core_stream_error_finalize_report_kind,
|
resolve_core_stream_error_finalize_report_kind,
|
||||||
resolve_local_candidate_failover_decision_stream, should_fallback_to_control_stream,
|
resolve_local_candidate_failover_analysis_stream, should_fallback_to_control_stream,
|
||||||
should_retry_next_local_candidate_stream, LocalFailoverDecision,
|
should_retry_next_local_candidate_stream, LocalFailoverDecision,
|
||||||
};
|
};
|
||||||
use crate::execution_runtime::{MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES};
|
use crate::execution_runtime::{MAX_STREAM_PREFETCH_BYTES, MAX_STREAM_PREFETCH_FRAMES};
|
||||||
use crate::log_ids::short_request_id;
|
use crate::log_ids::short_request_id;
|
||||||
|
use crate::orchestration::{
|
||||||
|
apply_local_execution_effect, LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect,
|
||||||
|
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
|
||||||
|
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
|
||||||
|
};
|
||||||
use crate::request_candidate_runtime::{
|
use crate::request_candidate_runtime::{
|
||||||
ensure_execution_request_candidate_slot, record_local_request_candidate_status,
|
ensure_execution_request_candidate_slot, record_local_request_candidate_status,
|
||||||
};
|
};
|
||||||
@@ -566,16 +570,7 @@ async fn execute_stream_from_frame_stream(
|
|||||||
let (body_json, body_base64) = decode_stream_error_body(&headers, &error_body);
|
let (body_json, body_base64) = decode_stream_error_body(&headers, &error_body);
|
||||||
let error_response_text =
|
let error_response_text =
|
||||||
local_failover_response_text(body_json.as_ref(), &error_body, None);
|
local_failover_response_text(body_json.as_ref(), &error_body, None);
|
||||||
record_pool_error_feedback(
|
let failover_analysis = resolve_local_candidate_failover_analysis_stream(
|
||||||
state,
|
|
||||||
&plan,
|
|
||||||
report_context.as_ref(),
|
|
||||||
status_code,
|
|
||||||
&headers,
|
|
||||||
error_response_text.as_deref(),
|
|
||||||
)
|
|
||||||
.await;
|
|
||||||
let failover_decision = resolve_local_candidate_failover_decision_stream(
|
|
||||||
state,
|
state,
|
||||||
&plan,
|
&plan,
|
||||||
report_context.as_ref(),
|
report_context.as_ref(),
|
||||||
@@ -583,6 +578,70 @@ async fn execute_stream_from_frame_stream(
|
|||||||
error_response_text.as_deref(),
|
error_response_text.as_deref(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan: &plan,
|
||||||
|
report_context: report_context.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
|
||||||
|
status_code,
|
||||||
|
classification: failover_analysis.classification,
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan: &plan,
|
||||||
|
report_context: report_context.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect {
|
||||||
|
status_code,
|
||||||
|
classification: failover_analysis.classification,
|
||||||
|
headers: Some(&headers),
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan: &plan,
|
||||||
|
report_context: report_context.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
|
||||||
|
status_code,
|
||||||
|
classification: failover_analysis.classification,
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan: &plan,
|
||||||
|
report_context: report_context.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
|
||||||
|
status_code,
|
||||||
|
response_text: error_response_text.as_deref(),
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan: &plan,
|
||||||
|
report_context: report_context.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::PoolError(LocalPoolErrorEffect {
|
||||||
|
status_code,
|
||||||
|
classification: failover_analysis.classification,
|
||||||
|
headers: &headers,
|
||||||
|
error_body: error_response_text.as_deref(),
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
let failover_decision = failover_analysis.decision;
|
||||||
debug!(
|
debug!(
|
||||||
event_name = "execution_runtime_stream_failover_decided",
|
event_name = "execution_runtime_stream_failover_decided",
|
||||||
log_type = "debug",
|
log_type = "debug",
|
||||||
@@ -1502,11 +1561,24 @@ async fn execute_stream_from_frame_stream(
|
|||||||
}),
|
}),
|
||||||
telemetry: telemetry.clone(),
|
telemetry: telemetry.clone(),
|
||||||
};
|
};
|
||||||
record_stream_pool_success_feedback(
|
apply_local_execution_effect(
|
||||||
&state_for_report,
|
&state_for_report,
|
||||||
&plan_for_report,
|
LocalExecutionEffectContext {
|
||||||
report_context_owned.as_ref(),
|
plan: &plan_for_report,
|
||||||
&usage_payload,
|
report_context: report_context_owned.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
&state_for_report,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan: &plan_for_report,
|
||||||
|
report_context: report_context_owned.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::PoolSuccessStream {
|
||||||
|
payload: &usage_payload,
|
||||||
|
},
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
record_stream_terminal_usage(
|
record_stream_terminal_usage(
|
||||||
|
|||||||
@@ -16,8 +16,13 @@ use crate::control::GatewayControlDecision;
|
|||||||
use crate::execution_runtime::submission::{
|
use crate::execution_runtime::submission::{
|
||||||
resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize,
|
resolve_core_error_background_report_kind, submit_local_core_error_or_sync_finalize,
|
||||||
};
|
};
|
||||||
use crate::execution_runtime::{record_pool_error_feedback, record_pool_stream_timeout_feedback};
|
|
||||||
use crate::log_ids::short_request_id;
|
use crate::log_ids::short_request_id;
|
||||||
|
use crate::orchestration::{
|
||||||
|
apply_local_execution_effect, resolve_local_failover_analysis_for_attempt,
|
||||||
|
LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect, LocalExecutionEffect,
|
||||||
|
LocalExecutionEffectContext, LocalHealthFailureEffect, LocalOAuthInvalidationEffect,
|
||||||
|
LocalPoolErrorEffect,
|
||||||
|
};
|
||||||
use crate::request_candidate_runtime::record_report_request_candidate_status;
|
use crate::request_candidate_runtime::record_report_request_candidate_status;
|
||||||
use crate::usage::submit_sync_report;
|
use crate::usage::submit_sync_report;
|
||||||
use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
|
use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
|
||||||
@@ -123,22 +128,92 @@ async fn record_stream_sync_failure(
|
|||||||
failure: &StreamFailureReport,
|
failure: &StreamFailureReport,
|
||||||
started_at_unix_ms: Option<u64>,
|
started_at_unix_ms: Option<u64>,
|
||||||
) {
|
) {
|
||||||
if matches!(
|
|
||||||
failure.error_type.as_str(),
|
|
||||||
"first_byte_timeout" | "read_timeout"
|
|
||||||
) {
|
|
||||||
record_pool_stream_timeout_feedback(state, plan, report_context).await;
|
|
||||||
}
|
|
||||||
let error_body = serde_json::to_string(&failure.body_json).ok();
|
let error_body = serde_json::to_string(&failure.body_json).ok();
|
||||||
record_pool_error_feedback(
|
let failure_analysis = resolve_local_failover_analysis_for_attempt(
|
||||||
state,
|
state,
|
||||||
plan,
|
plan,
|
||||||
report_context,
|
report_context,
|
||||||
failure.status_code,
|
failure.status_code,
|
||||||
&payload.headers,
|
|
||||||
error_body.as_deref(),
|
error_body.as_deref(),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
|
if matches!(
|
||||||
|
failure.error_type.as_str(),
|
||||||
|
"first_byte_timeout" | "read_timeout"
|
||||||
|
) {
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan,
|
||||||
|
report_context,
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::PoolStreamTimeout,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
}
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan,
|
||||||
|
report_context,
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
|
||||||
|
status_code: failure.status_code,
|
||||||
|
classification: failure_analysis.classification,
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan,
|
||||||
|
report_context,
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect {
|
||||||
|
status_code: failure.status_code,
|
||||||
|
classification: failure_analysis.classification,
|
||||||
|
headers: Some(&payload.headers),
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan,
|
||||||
|
report_context,
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
|
||||||
|
status_code: failure.status_code,
|
||||||
|
classification: failure_analysis.classification,
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan,
|
||||||
|
report_context,
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
|
||||||
|
status_code: failure.status_code,
|
||||||
|
response_text: error_body.as_deref(),
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan,
|
||||||
|
report_context,
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::PoolError(LocalPoolErrorEffect {
|
||||||
|
status_code: failure.status_code,
|
||||||
|
classification: failure_analysis.classification,
|
||||||
|
headers: &payload.headers,
|
||||||
|
error_body: error_body.as_deref(),
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
let context_seed = build_terminal_usage_context_seed(plan, report_context);
|
let context_seed = build_terminal_usage_context_seed(plan, report_context);
|
||||||
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
|
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
|
||||||
state
|
state
|
||||||
|
|||||||
@@ -30,12 +30,16 @@ use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_
|
|||||||
use crate::execution_runtime::submission::submit_local_core_error_or_sync_finalize;
|
use crate::execution_runtime::submission::submit_local_core_error_or_sync_finalize;
|
||||||
use crate::execution_runtime::transport::DirectSyncExecutionRuntime;
|
use crate::execution_runtime::transport::DirectSyncExecutionRuntime;
|
||||||
use crate::execution_runtime::{
|
use crate::execution_runtime::{
|
||||||
local_failover_response_text, record_pool_error_feedback, record_sync_pool_success_feedback,
|
analyze_local_candidate_failover_sync, local_failover_response_text,
|
||||||
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_sync,
|
resolve_core_sync_error_finalize_report_kind, should_fallback_to_control_sync,
|
||||||
should_finalize_sync_response, should_retry_next_local_candidate_sync,
|
should_finalize_sync_response, LocalFailoverDecision,
|
||||||
should_stop_local_candidate_failover_sync,
|
|
||||||
};
|
};
|
||||||
use crate::log_ids::short_request_id;
|
use crate::log_ids::short_request_id;
|
||||||
|
use crate::orchestration::{
|
||||||
|
apply_local_execution_effect, LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect,
|
||||||
|
LocalExecutionEffect, LocalExecutionEffectContext, LocalHealthFailureEffect,
|
||||||
|
LocalHealthSuccessEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
|
||||||
|
};
|
||||||
use crate::request_candidate_runtime::{
|
use crate::request_candidate_runtime::{
|
||||||
ensure_execution_request_candidate_slot, record_local_request_candidate_status,
|
ensure_execution_request_candidate_slot, record_local_request_candidate_status,
|
||||||
};
|
};
|
||||||
@@ -247,7 +251,7 @@ pub(crate) async fn execute_execution_runtime_sync(
|
|||||||
&body_bytes,
|
&body_bytes,
|
||||||
result.error.as_ref().map(|error| error.message.as_str()),
|
result.error.as_ref().map(|error| error.message.as_str()),
|
||||||
);
|
);
|
||||||
let stop_local_failover = should_stop_local_candidate_failover_sync(
|
let local_failover_analysis = analyze_local_candidate_failover_sync(
|
||||||
state,
|
state,
|
||||||
&plan,
|
&plan,
|
||||||
plan_kind,
|
plan_kind,
|
||||||
@@ -257,27 +261,74 @@ pub(crate) async fn execute_execution_runtime_sync(
|
|||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
if result.status_code >= 400 {
|
if result.status_code >= 400 {
|
||||||
record_pool_error_feedback(
|
apply_local_execution_effect(
|
||||||
state,
|
state,
|
||||||
&plan,
|
LocalExecutionEffectContext {
|
||||||
report_context.as_ref(),
|
plan: &plan,
|
||||||
result.status_code,
|
report_context: report_context.as_ref(),
|
||||||
&headers,
|
},
|
||||||
local_failover_response_text.as_deref(),
|
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
|
||||||
|
status_code: result.status_code,
|
||||||
|
classification: local_failover_analysis.classification,
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan: &plan,
|
||||||
|
report_context: report_context.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect {
|
||||||
|
status_code: result.status_code,
|
||||||
|
classification: local_failover_analysis.classification,
|
||||||
|
headers: Some(&headers),
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan: &plan,
|
||||||
|
report_context: report_context.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
|
||||||
|
status_code: result.status_code,
|
||||||
|
classification: local_failover_analysis.classification,
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan: &plan,
|
||||||
|
report_context: report_context.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
|
||||||
|
status_code: result.status_code,
|
||||||
|
response_text: local_failover_response_text.as_deref(),
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan: &plan,
|
||||||
|
report_context: report_context.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::PoolError(LocalPoolErrorEffect {
|
||||||
|
status_code: result.status_code,
|
||||||
|
classification: local_failover_analysis.classification,
|
||||||
|
headers: &headers,
|
||||||
|
error_body: local_failover_response_text.as_deref(),
|
||||||
|
}),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
if should_retry_next_local_candidate_sync(
|
if matches!(
|
||||||
state,
|
local_failover_analysis.decision,
|
||||||
&plan,
|
LocalFailoverDecision::RetryNextCandidate
|
||||||
plan_kind,
|
) {
|
||||||
report_context.as_ref(),
|
|
||||||
&result,
|
|
||||||
local_failover_response_text.as_deref(),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
&& !stop_local_failover
|
|
||||||
{
|
|
||||||
let terminal_unix_secs = current_request_candidate_unix_ms();
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
||||||
record_local_request_candidate_status(
|
record_local_request_candidate_status(
|
||||||
state,
|
state,
|
||||||
@@ -341,16 +392,17 @@ pub(crate) async fn execute_execution_runtime_sync(
|
|||||||
mapped_error_finalize_kind.clone()
|
mapped_error_finalize_kind.clone()
|
||||||
};
|
};
|
||||||
|
|
||||||
if !stop_local_failover
|
if !matches!(
|
||||||
&& should_fallback_to_control_sync(
|
local_failover_analysis.decision,
|
||||||
plan_kind,
|
LocalFailoverDecision::StopLocalFailover
|
||||||
&result,
|
) && should_fallback_to_control_sync(
|
||||||
body_json.as_ref(),
|
plan_kind,
|
||||||
has_body_bytes,
|
&result,
|
||||||
explicit_finalize || implicit_finalize.is_some(),
|
body_json.as_ref(),
|
||||||
mapped_error_finalize_kind.is_some(),
|
has_body_bytes,
|
||||||
)
|
explicit_finalize || implicit_finalize.is_some(),
|
||||||
{
|
mapped_error_finalize_kind.is_some(),
|
||||||
|
) {
|
||||||
let terminal_unix_secs = current_request_candidate_unix_ms();
|
let terminal_unix_secs = current_request_candidate_unix_ms();
|
||||||
record_local_request_candidate_status(
|
record_local_request_candidate_status(
|
||||||
state,
|
state,
|
||||||
@@ -406,11 +458,24 @@ pub(crate) async fn execute_execution_runtime_sync(
|
|||||||
telemetry: result.telemetry.clone(),
|
telemetry: result.telemetry.clone(),
|
||||||
};
|
};
|
||||||
if result.status_code < 400 {
|
if result.status_code < 400 {
|
||||||
record_sync_pool_success_feedback(
|
apply_local_execution_effect(
|
||||||
state,
|
state,
|
||||||
&plan,
|
LocalExecutionEffectContext {
|
||||||
report_context.as_ref(),
|
plan: &plan,
|
||||||
&base_usage_payload,
|
report_context: report_context.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
apply_local_execution_effect(
|
||||||
|
state,
|
||||||
|
LocalExecutionEffectContext {
|
||||||
|
plan: &plan,
|
||||||
|
report_context: report_context.as_ref(),
|
||||||
|
},
|
||||||
|
LocalExecutionEffect::PoolSuccessSync {
|
||||||
|
payload: &base_usage_payload,
|
||||||
|
},
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -128,7 +128,7 @@ impl<'a> AdminAppState<'a> {
|
|||||||
display_name: Option<&str>,
|
display_name: Option<&str>,
|
||||||
mime_type: Option<&str>,
|
mime_type: Option<&str>,
|
||||||
) -> Result<(), GatewayError> {
|
) -> Result<(), GatewayError> {
|
||||||
crate::usage::reporting::store_local_gemini_file_mapping(
|
crate::orchestration::store_local_gemini_file_mapping(
|
||||||
self.app,
|
self.app,
|
||||||
file_name,
|
file_name,
|
||||||
key_id,
|
key_id,
|
||||||
|
|||||||
@@ -48,6 +48,7 @@ mod log_ids;
|
|||||||
mod maintenance;
|
mod maintenance;
|
||||||
pub(crate) mod middleware;
|
pub(crate) mod middleware;
|
||||||
mod model_fetch;
|
mod model_fetch;
|
||||||
|
mod orchestration;
|
||||||
mod provider_key_auth;
|
mod provider_key_auth;
|
||||||
pub(crate) use aether_provider_transport as provider_transport;
|
pub(crate) use aether_provider_transport as provider_transport;
|
||||||
mod query;
|
mod query;
|
||||||
|
|||||||
335
apps/aether-gateway/src/orchestration/adaptive.rs
Normal file
335
apps/aether-gateway/src/orchestration/adaptive.rs
Normal 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
234
apps/aether-gateway/src/orchestration/attempt.rs
Normal 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
510
apps/aether-gateway/src/orchestration/classifier.rs
Normal 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
1215
apps/aether-gateway/src/orchestration/effects.rs
Normal file
File diff suppressed because it is too large
Load Diff
188
apps/aether-gateway/src/orchestration/health.rs
Normal file
188
apps/aether-gateway/src/orchestration/health.rs
Normal 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
78
apps/aether-gateway/src/orchestration/mod.rs
Normal 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
358
apps/aether-gateway/src/orchestration/policy.rs
Normal 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
150
apps/aether-gateway/src/orchestration/recovery.rs
Normal 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
424
apps/aether-gateway/src/orchestration/report_effects.rs
Normal 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(¤t_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();
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -363,6 +363,7 @@ pub(crate) async fn persist_available_local_candidate(
|
|||||||
api_key_id: &str,
|
api_key_id: &str,
|
||||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||||
candidate_index: u32,
|
candidate_index: u32,
|
||||||
|
retry_index: u32,
|
||||||
candidate_id: &str,
|
candidate_id: &str,
|
||||||
required_capabilities: Option<&Value>,
|
required_capabilities: Option<&Value>,
|
||||||
extra_data: Option<serde_json::Value>,
|
extra_data: Option<serde_json::Value>,
|
||||||
@@ -378,7 +379,7 @@ pub(crate) async fn persist_available_local_candidate(
|
|||||||
username: None,
|
username: None,
|
||||||
api_key_name: None,
|
api_key_name: None,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
retry_index: 0,
|
retry_index,
|
||||||
provider_id: Some(candidate.provider_id.clone()),
|
provider_id: Some(candidate.provider_id.clone()),
|
||||||
endpoint_id: Some(candidate.endpoint_id.clone()),
|
endpoint_id: Some(candidate.endpoint_id.clone()),
|
||||||
key_id: Some(candidate.key_id.clone()),
|
key_id: Some(candidate.key_id.clone()),
|
||||||
@@ -405,7 +406,7 @@ pub(crate) async fn persist_available_local_candidate(
|
|||||||
request_id = %short_request_id(trace_id),
|
request_id = %short_request_id(trace_id),
|
||||||
candidate_id = %stored.id,
|
candidate_id = %stored.id,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
retry_index = 0,
|
retry_index,
|
||||||
status = "available",
|
status = "available",
|
||||||
source = "planner_available",
|
source = "planner_available",
|
||||||
provider_id = %candidate.provider_id,
|
provider_id = %candidate.provider_id,
|
||||||
@@ -423,7 +424,7 @@ pub(crate) async fn persist_available_local_candidate(
|
|||||||
request_id = %short_request_id(trace_id),
|
request_id = %short_request_id(trace_id),
|
||||||
candidate_id = %candidate_id,
|
candidate_id = %candidate_id,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
retry_index = 0,
|
retry_index,
|
||||||
status = "available",
|
status = "available",
|
||||||
source = "planner_available",
|
source = "planner_available",
|
||||||
provider_id = %candidate.provider_id,
|
provider_id = %candidate.provider_id,
|
||||||
@@ -452,6 +453,7 @@ pub(crate) async fn persist_skipped_local_candidate(
|
|||||||
api_key_id: &str,
|
api_key_id: &str,
|
||||||
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
candidate: &SchedulerMinimalCandidateSelectionCandidate,
|
||||||
candidate_index: u32,
|
candidate_index: u32,
|
||||||
|
retry_index: u32,
|
||||||
candidate_id: &str,
|
candidate_id: &str,
|
||||||
required_capabilities: Option<&Value>,
|
required_capabilities: Option<&Value>,
|
||||||
skip_reason: &str,
|
skip_reason: &str,
|
||||||
@@ -468,7 +470,7 @@ pub(crate) async fn persist_skipped_local_candidate(
|
|||||||
username: None,
|
username: None,
|
||||||
api_key_name: None,
|
api_key_name: None,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
retry_index: 0,
|
retry_index,
|
||||||
provider_id: Some(candidate.provider_id.clone()),
|
provider_id: Some(candidate.provider_id.clone()),
|
||||||
endpoint_id: Some(candidate.endpoint_id.clone()),
|
endpoint_id: Some(candidate.endpoint_id.clone()),
|
||||||
key_id: Some(candidate.key_id.clone()),
|
key_id: Some(candidate.key_id.clone()),
|
||||||
@@ -495,7 +497,7 @@ pub(crate) async fn persist_skipped_local_candidate(
|
|||||||
request_id = %short_request_id(trace_id),
|
request_id = %short_request_id(trace_id),
|
||||||
candidate_id = %stored.id,
|
candidate_id = %stored.id,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
retry_index = 0,
|
retry_index,
|
||||||
status = "skipped",
|
status = "skipped",
|
||||||
skip_reason,
|
skip_reason,
|
||||||
source = "planner_skipped",
|
source = "planner_skipped",
|
||||||
@@ -513,7 +515,7 @@ pub(crate) async fn persist_skipped_local_candidate(
|
|||||||
request_id = %short_request_id(trace_id),
|
request_id = %short_request_id(trace_id),
|
||||||
candidate_id = %candidate_id,
|
candidate_id = %candidate_id,
|
||||||
candidate_index,
|
candidate_index,
|
||||||
retry_index = 0,
|
retry_index,
|
||||||
status = "skipped",
|
status = "skipped",
|
||||||
skip_reason,
|
skip_reason,
|
||||||
source = "planner_skipped",
|
source = "planner_skipped",
|
||||||
@@ -892,6 +894,7 @@ mod tests {
|
|||||||
"api-key-1",
|
"api-key-1",
|
||||||
&sample_minimal_candidate(),
|
&sample_minimal_candidate(),
|
||||||
0,
|
0,
|
||||||
|
0,
|
||||||
"cand-runtime-cap-123",
|
"cand-runtime-cap-123",
|
||||||
Some(&required_capabilities),
|
Some(&required_capabilities),
|
||||||
None,
|
None,
|
||||||
|
|||||||
@@ -688,6 +688,39 @@ impl AppState {
|
|||||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn update_provider_catalog_key_format_health(
|
||||||
|
&self,
|
||||||
|
key_id: &str,
|
||||||
|
api_format: &str,
|
||||||
|
health_by_format: &serde_json::Value,
|
||||||
|
) -> Result<bool, GatewayError> {
|
||||||
|
let api_format = api_format.trim();
|
||||||
|
if api_format.is_empty() {
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
|
||||||
|
let Some(current_key) = self
|
||||||
|
.read_provider_catalog_keys_by_ids(&[key_id.to_string()])
|
||||||
|
.await?
|
||||||
|
.into_iter()
|
||||||
|
.next()
|
||||||
|
else {
|
||||||
|
return Ok(false);
|
||||||
|
};
|
||||||
|
|
||||||
|
if current_key.health_by_format.as_ref() == Some(health_by_format) {
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
|
||||||
|
self.update_provider_catalog_key_health_state(
|
||||||
|
key_id,
|
||||||
|
current_key.is_active,
|
||||||
|
Some(health_by_format),
|
||||||
|
current_key.circuit_breaker_by_format.as_ref(),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) async fn update_provider_catalog_key_health_state(
|
pub(crate) async fn update_provider_catalog_key_health_state(
|
||||||
&self,
|
&self,
|
||||||
key_id: &str,
|
key_id: &str,
|
||||||
|
|||||||
@@ -7,6 +7,8 @@ use crate::handlers::shared::default_provider_key_status_snapshot;
|
|||||||
use crate::provider_transport::LocalOAuthHttpExecutor;
|
use crate::provider_transport::LocalOAuthHttpExecutor;
|
||||||
|
|
||||||
use super::super::provider_transport;
|
use super::super::provider_transport;
|
||||||
|
use crate::provider_key_auth::provider_key_is_oauth_managed;
|
||||||
|
use aether_admin::provider::quota as admin_provider_quota_pure;
|
||||||
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
|
use aether_contracts::{ExecutionPlan, ExecutionTimeouts, RequestBody};
|
||||||
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
|
||||||
use base64::{engine::general_purpose::STANDARD, Engine as _};
|
use base64::{engine::general_purpose::STANDARD, Engine as _};
|
||||||
@@ -676,6 +678,65 @@ impl AppState {
|
|||||||
self.oauth_refresh.invalidate_cached_entry(key_id).await
|
self.oauth_refresh.invalidate_cached_entry(key_id).await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn mark_provider_catalog_key_oauth_invalid(
|
||||||
|
&self,
|
||||||
|
key_id: &str,
|
||||||
|
provider_type: &str,
|
||||||
|
invalid_reason: &str,
|
||||||
|
) -> Result<bool, GatewayError> {
|
||||||
|
let invalid_reason = invalid_reason.trim();
|
||||||
|
if invalid_reason.is_empty() {
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
|
||||||
|
let Some(mut latest_key) = self
|
||||||
|
.data
|
||||||
|
.list_provider_catalog_keys_by_ids(&[key_id.to_string()])
|
||||||
|
.await
|
||||||
|
.map_err(|err| GatewayError::Internal(err.to_string()))?
|
||||||
|
.into_iter()
|
||||||
|
.next()
|
||||||
|
else {
|
||||||
|
return Ok(false);
|
||||||
|
};
|
||||||
|
|
||||||
|
if !provider_key_is_oauth_managed(&latest_key, provider_type) {
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
|
||||||
|
let now_unix_secs = SystemTime::now()
|
||||||
|
.duration_since(UNIX_EPOCH)
|
||||||
|
.ok()
|
||||||
|
.map(|duration| duration.as_secs())
|
||||||
|
.unwrap_or(0);
|
||||||
|
let (oauth_invalid_at_unix_secs, oauth_invalid_reason) = merge_runtime_oauth_invalid_state(
|
||||||
|
provider_type,
|
||||||
|
&latest_key,
|
||||||
|
invalid_reason,
|
||||||
|
now_unix_secs,
|
||||||
|
);
|
||||||
|
if oauth_invalid_at_unix_secs == latest_key.oauth_invalid_at_unix_secs
|
||||||
|
&& oauth_invalid_reason == latest_key.oauth_invalid_reason
|
||||||
|
{
|
||||||
|
return Ok(false);
|
||||||
|
}
|
||||||
|
|
||||||
|
latest_key.oauth_invalid_at_unix_secs = oauth_invalid_at_unix_secs;
|
||||||
|
latest_key.oauth_invalid_reason = oauth_invalid_reason;
|
||||||
|
latest_key.updated_at_unix_secs = Some(now_unix_secs);
|
||||||
|
let current_status_snapshot = latest_key.status_snapshot.take();
|
||||||
|
latest_key.status_snapshot =
|
||||||
|
sync_provider_key_oauth_status_snapshot(current_status_snapshot, &latest_key);
|
||||||
|
let updated = self
|
||||||
|
.update_provider_catalog_key(&latest_key)
|
||||||
|
.await?
|
||||||
|
.is_some();
|
||||||
|
if updated {
|
||||||
|
let _ = self.invalidate_local_oauth_refresh_entry(key_id).await;
|
||||||
|
}
|
||||||
|
Ok(updated)
|
||||||
|
}
|
||||||
|
|
||||||
async fn persist_local_oauth_refresh_entry(
|
async fn persist_local_oauth_refresh_entry(
|
||||||
&self,
|
&self,
|
||||||
transport: &provider_transport::GatewayProviderTransportSnapshot,
|
transport: &provider_transport::GatewayProviderTransportSnapshot,
|
||||||
@@ -858,6 +919,43 @@ impl AppState {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn merge_runtime_oauth_invalid_state(
|
||||||
|
provider_type: &str,
|
||||||
|
key: &StoredProviderCatalogKey,
|
||||||
|
invalid_reason: &str,
|
||||||
|
now_unix_secs: u64,
|
||||||
|
) -> (Option<u64>, Option<String>) {
|
||||||
|
let candidate_reason = invalid_reason.trim();
|
||||||
|
if candidate_reason.is_empty() {
|
||||||
|
return (
|
||||||
|
key.oauth_invalid_at_unix_secs,
|
||||||
|
key.oauth_invalid_reason.clone(),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
if provider_type.trim().eq_ignore_ascii_case("codex") {
|
||||||
|
return admin_provider_quota_pure::codex_build_invalid_state(
|
||||||
|
key,
|
||||||
|
candidate_reason.to_string(),
|
||||||
|
now_unix_secs,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
let current_reason = key
|
||||||
|
.oauth_invalid_reason
|
||||||
|
.as_deref()
|
||||||
|
.map(str::trim)
|
||||||
|
.unwrap_or_default();
|
||||||
|
if current_reason == candidate_reason {
|
||||||
|
return (
|
||||||
|
key.oauth_invalid_at_unix_secs,
|
||||||
|
(!current_reason.is_empty()).then_some(current_reason.to_string()),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
(Some(now_unix_secs), Some(candidate_reason.to_string()))
|
||||||
|
}
|
||||||
|
|
||||||
fn local_oauth_execution_body_text(result: &aether_contracts::ExecutionResult) -> String {
|
fn local_oauth_execution_body_text(result: &aether_contracts::ExecutionResult) -> String {
|
||||||
result
|
result
|
||||||
.body
|
.body
|
||||||
|
|||||||
@@ -177,13 +177,11 @@ fn usage_runtime_paths_depend_on_shared_crates_not_app_runtime_shims() {
|
|||||||
"is_local_ai_sync_report_kind",
|
"is_local_ai_sync_report_kind",
|
||||||
"is_local_ai_stream_report_kind",
|
"is_local_ai_stream_report_kind",
|
||||||
"sync_report_represents_failure",
|
"sync_report_represents_failure",
|
||||||
"extract_gemini_file_mapping_entries",
|
|
||||||
"gemini_file_mapping_cache_key",
|
|
||||||
"normalize_gemini_file_name",
|
|
||||||
"report_request_id",
|
"report_request_id",
|
||||||
"should_handle_local_sync_report",
|
"should_handle_local_sync_report",
|
||||||
"should_handle_local_stream_report",
|
"should_handle_local_stream_report",
|
||||||
"GEMINI_FILE_MAPPING_TTL_SECONDS",
|
"apply_local_report_effect",
|
||||||
|
"LocalReportEffect",
|
||||||
] {
|
] {
|
||||||
assert!(
|
assert!(
|
||||||
usage_reporting_mod.contains(pattern),
|
usage_reporting_mod.contains(pattern),
|
||||||
@@ -206,12 +204,39 @@ fn usage_runtime_paths_depend_on_shared_crates_not_app_runtime_shims() {
|
|||||||
"fn should_handle_local_sync_report(",
|
"fn should_handle_local_sync_report(",
|
||||||
"fn should_handle_local_stream_report(",
|
"fn should_handle_local_stream_report(",
|
||||||
"\"openai_video_delete_sync_success\" && payload.status_code == 404",
|
"\"openai_video_delete_sync_success\" && payload.status_code == 404",
|
||||||
|
"sync_codex_quota_from_response_headers(",
|
||||||
|
"apply_local_gemini_file_mapping_report_effect(",
|
||||||
|
"pub(crate) async fn store_local_gemini_file_mapping(",
|
||||||
] {
|
] {
|
||||||
assert!(
|
assert!(
|
||||||
!usage_reporting_mod.contains(pattern),
|
!usage_reporting_mod.contains(pattern),
|
||||||
"usage/reporting/mod.rs should not own local report classification logic {pattern}"
|
"usage/reporting/mod.rs should not own local report classification logic {pattern}"
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let report_effects =
|
||||||
|
read_workspace_file("apps/aether-gateway/src/orchestration/report_effects.rs");
|
||||||
|
assert!(
|
||||||
|
report_effects.contains("aether_usage_runtime"),
|
||||||
|
"orchestration/report_effects.rs should depend on aether_usage_runtime"
|
||||||
|
);
|
||||||
|
for pattern in [
|
||||||
|
"extract_gemini_file_mapping_entries",
|
||||||
|
"gemini_file_mapping_cache_key",
|
||||||
|
"normalize_gemini_file_name",
|
||||||
|
"report_request_id",
|
||||||
|
"GEMINI_FILE_MAPPING_TTL_SECONDS",
|
||||||
|
"sync_codex_quota_from_response_headers",
|
||||||
|
"store_local_gemini_file_mapping",
|
||||||
|
"delete_local_gemini_file_mapping",
|
||||||
|
"GatewaySyncReportRequest",
|
||||||
|
"GatewayStreamReportRequest",
|
||||||
|
] {
|
||||||
|
assert!(
|
||||||
|
report_effects.contains(pattern),
|
||||||
|
"orchestration/report_effects.rs should own local report effect detail {pattern}"
|
||||||
|
);
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
|||||||
@@ -2768,6 +2768,49 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
|
|||||||
}),
|
}),
|
||||||
);
|
);
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
struct SeenExecutionRuntimeRequest {
|
||||||
|
url: String,
|
||||||
|
authorization: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
let seen_execution_runtime = Arc::new(Mutex::new(None::<SeenExecutionRuntimeRequest>));
|
||||||
|
let seen_execution_runtime_clone = Arc::clone(&seen_execution_runtime);
|
||||||
|
let execution_runtime = Router::new().route(
|
||||||
|
"/v1/execute/sync",
|
||||||
|
any(move |request: Request| {
|
||||||
|
let seen_execution_runtime_inner = Arc::clone(&seen_execution_runtime_clone);
|
||||||
|
async move {
|
||||||
|
let plan: aether_contracts::ExecutionPlan = serde_json::from_slice(
|
||||||
|
&to_bytes(request.into_body(), usize::MAX)
|
||||||
|
.await
|
||||||
|
.expect("body should read"),
|
||||||
|
)
|
||||||
|
.expect("plan should parse");
|
||||||
|
*seen_execution_runtime_inner
|
||||||
|
.lock()
|
||||||
|
.expect("mutex should lock") = Some(SeenExecutionRuntimeRequest {
|
||||||
|
url: plan.url.clone(),
|
||||||
|
authorization: plan
|
||||||
|
.headers
|
||||||
|
.get("authorization")
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or_default(),
|
||||||
|
});
|
||||||
|
let result = aether_contracts::ExecutionResult {
|
||||||
|
request_id: plan.request_id,
|
||||||
|
candidate_id: None,
|
||||||
|
status_code: 401,
|
||||||
|
headers: std::collections::BTreeMap::new(),
|
||||||
|
body: None,
|
||||||
|
telemetry: None,
|
||||||
|
error: None,
|
||||||
|
};
|
||||||
|
(StatusCode::OK, Json(result))
|
||||||
|
}
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
|
||||||
let mut provider = sample_provider("provider-codex", "codex", 10);
|
let mut provider = sample_provider("provider-codex", "codex", 10);
|
||||||
provider.provider_type = "codex".to_string();
|
provider.provider_type = "codex".to_string();
|
||||||
let endpoint = sample_endpoint(
|
let endpoint = sample_endpoint(
|
||||||
@@ -2832,6 +2875,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
|
|||||||
|
|
||||||
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
let (upstream_url, upstream_handle) = start_server(upstream).await;
|
||||||
let (token_url, token_handle) = start_server(token_server).await;
|
let (token_url, token_handle) = start_server(token_server).await;
|
||||||
|
let (execution_runtime_url, execution_runtime_handle) = start_server(execution_runtime).await;
|
||||||
let oauth_refresh =
|
let oauth_refresh =
|
||||||
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
|
crate::provider_transport::LocalOAuthRefreshCoordinator::with_adapters_for_tests(vec![
|
||||||
Arc::new(
|
Arc::new(
|
||||||
@@ -2840,8 +2884,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
|
|||||||
),
|
),
|
||||||
]);
|
]);
|
||||||
let gateway = build_router_with_state(
|
let gateway = build_router_with_state(
|
||||||
AppState::new()
|
build_state_with_execution_runtime_override(execution_runtime_url)
|
||||||
.expect("gateway should build")
|
|
||||||
.with_data_state_for_tests(
|
.with_data_state_for_tests(
|
||||||
GatewayDataState::with_provider_catalog_repository_for_tests(
|
GatewayDataState::with_provider_catalog_repository_for_tests(
|
||||||
provider_catalog_repository.clone(),
|
provider_catalog_repository.clone(),
|
||||||
@@ -2889,6 +2932,19 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
|
|||||||
}
|
}
|
||||||
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);
|
||||||
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
|
assert_eq!(*token_hits.lock().expect("mutex should lock"), 1);
|
||||||
|
let seen_execution_runtime_request = seen_execution_runtime
|
||||||
|
.lock()
|
||||||
|
.expect("mutex should lock")
|
||||||
|
.clone()
|
||||||
|
.expect("execution runtime request should be captured");
|
||||||
|
assert_eq!(
|
||||||
|
seen_execution_runtime_request.url,
|
||||||
|
"https://chatgpt.com/backend-api/wham/usage"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
seen_execution_runtime_request.authorization,
|
||||||
|
"Bearer refreshed-codex-access-token"
|
||||||
|
);
|
||||||
|
|
||||||
let stored_key = provider_catalog_repository
|
let stored_key = provider_catalog_repository
|
||||||
.list_keys_by_ids(&["key-codex-oauth-refresh".to_string()])
|
.list_keys_by_ids(&["key-codex-oauth-refresh".to_string()])
|
||||||
@@ -2974,6 +3030,7 @@ async fn gateway_refreshes_admin_provider_oauth_key_locally_with_trusted_admin_p
|
|||||||
);
|
);
|
||||||
|
|
||||||
gateway_handle.abort();
|
gateway_handle.abort();
|
||||||
|
execution_runtime_handle.abort();
|
||||||
token_handle.abort();
|
token_handle.abort();
|
||||||
upstream_handle.abort();
|
upstream_handle.abort();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,625 +0,0 @@
|
|||||||
use std::collections::{BTreeMap, HashMap};
|
|
||||||
use std::sync::{Mutex, OnceLock};
|
|
||||||
use std::time::{Duration, Instant};
|
|
||||||
|
|
||||||
use crate::clock::current_unix_secs;
|
|
||||||
use crate::handlers::shared::sync_provider_key_quota_status_snapshot;
|
|
||||||
use crate::{AppState, GatewayError};
|
|
||||||
use aether_admin::provider::quota as admin_provider_quota_pure;
|
|
||||||
use serde_json::Value;
|
|
||||||
|
|
||||||
const CACHE_TTL_SECONDS: u64 = 30;
|
|
||||||
const CACHE_MAX_ENTRIES: usize = 4096;
|
|
||||||
|
|
||||||
type HeaderFingerprintCache = Mutex<HashMap<String, (String, Instant)>>;
|
|
||||||
|
|
||||||
static HEADER_FINGERPRINT_CACHE: OnceLock<HeaderFingerprintCache> = OnceLock::new();
|
|
||||||
|
|
||||||
fn header_fingerprint_cache() -> &'static HeaderFingerprintCache {
|
|
||||||
HEADER_FINGERPRINT_CACHE.get_or_init(|| Mutex::new(HashMap::new()))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn report_context_key_id(report_context: Option<&Value>) -> Option<String> {
|
|
||||||
report_context
|
|
||||||
.and_then(|context| context.get("key_id"))
|
|
||||||
.and_then(Value::as_str)
|
|
||||||
.map(str::trim)
|
|
||||||
.filter(|value| !value.is_empty())
|
|
||||||
.map(ToOwned::to_owned)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn is_volatile_compare_field(key: &str) -> bool {
|
|
||||||
key == "updated_at" || key.ends_with("_reset_seconds") || key.ends_with("_reset_after_seconds")
|
|
||||||
}
|
|
||||||
|
|
||||||
fn canonicalize_value(value: &Value) -> Value {
|
|
||||||
match value {
|
|
||||||
Value::Array(items) => Value::Array(items.iter().map(canonicalize_value).collect()),
|
|
||||||
Value::Object(object) => {
|
|
||||||
let mut entries = object.iter().collect::<Vec<_>>();
|
|
||||||
entries.sort_by(|left, right| left.0.cmp(right.0));
|
|
||||||
let mut normalized = serde_json::Map::new();
|
|
||||||
for (key, value) in entries {
|
|
||||||
normalized.insert(key.clone(), canonicalize_value(value));
|
|
||||||
}
|
|
||||||
Value::Object(normalized)
|
|
||||||
}
|
|
||||||
_ => value.clone(),
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn fingerprint_codex_payload(value: &Value) -> Option<String> {
|
|
||||||
let object = value.as_object()?;
|
|
||||||
let mut entries = object
|
|
||||||
.iter()
|
|
||||||
.filter(|(key, _)| !is_volatile_compare_field(key))
|
|
||||||
.collect::<Vec<_>>();
|
|
||||||
entries.sort_by(|left, right| left.0.cmp(right.0));
|
|
||||||
|
|
||||||
let mut normalized = serde_json::Map::new();
|
|
||||||
for (key, value) in entries {
|
|
||||||
normalized.insert(key.clone(), canonicalize_value(value));
|
|
||||||
}
|
|
||||||
serde_json::to_string(&Value::Object(normalized)).ok()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn get_cached_fingerprint(key_id: &str, now: Instant) -> Option<String> {
|
|
||||||
let mut cache = header_fingerprint_cache()
|
|
||||||
.lock()
|
|
||||||
.expect("codex realtime quota cache should lock");
|
|
||||||
match cache.get(key_id) {
|
|
||||||
Some((fingerprint, expires_at)) if *expires_at > now => Some(fingerprint.clone()),
|
|
||||||
Some(_) => {
|
|
||||||
cache.remove(key_id);
|
|
||||||
None
|
|
||||||
}
|
|
||||||
None => None,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn set_cached_fingerprint(key_id: &str, fingerprint: String, now: Instant) {
|
|
||||||
let mut cache = header_fingerprint_cache()
|
|
||||||
.lock()
|
|
||||||
.expect("codex realtime quota cache should lock");
|
|
||||||
cache.insert(
|
|
||||||
key_id.to_string(),
|
|
||||||
(
|
|
||||||
fingerprint,
|
|
||||||
now.checked_add(Duration::from_secs(CACHE_TTL_SECONDS))
|
|
||||||
.unwrap_or(now),
|
|
||||||
),
|
|
||||||
);
|
|
||||||
|
|
||||||
cache.retain(|_, (_, expires_at)| *expires_at > now);
|
|
||||||
if cache.len() <= CACHE_MAX_ENTRIES {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
let mut entries = cache
|
|
||||||
.iter()
|
|
||||||
.map(|(key, (_, expires_at))| (key.clone(), *expires_at))
|
|
||||||
.collect::<Vec<_>>();
|
|
||||||
entries.sort_by_key(|entry| entry.1);
|
|
||||||
for (key, _) in entries.into_iter().take(cache.len() - CACHE_MAX_ENTRIES) {
|
|
||||||
cache.remove(&key);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
fn merge_metadata_object(
|
|
||||||
current: Option<&Value>,
|
|
||||||
section_key: &str,
|
|
||||||
section_value: Value,
|
|
||||||
) -> Option<Value> {
|
|
||||||
let mut merged = current
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.cloned()
|
|
||||||
.unwrap_or_default();
|
|
||||||
merged.insert(section_key.to_string(), section_value);
|
|
||||||
Some(Value::Object(merged))
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(super) async fn sync_codex_quota_from_response_headers(
|
|
||||||
state: &AppState,
|
|
||||||
report_context: Option<&Value>,
|
|
||||||
headers: &BTreeMap<String, String>,
|
|
||||||
) -> Result<bool, GatewayError> {
|
|
||||||
let key_id = match report_context_key_id(report_context) {
|
|
||||||
Some(value) => value,
|
|
||||||
None => return Ok(false),
|
|
||||||
};
|
|
||||||
|
|
||||||
let now_unix_secs = current_unix_secs();
|
|
||||||
let Some(parsed) = admin_provider_quota_pure::parse_codex_usage_headers(headers, now_unix_secs)
|
|
||||||
else {
|
|
||||||
return Ok(false);
|
|
||||||
};
|
|
||||||
let Some(incoming_fingerprint) = fingerprint_codex_payload(&parsed) else {
|
|
||||||
return Ok(false);
|
|
||||||
};
|
|
||||||
|
|
||||||
let now = Instant::now();
|
|
||||||
if get_cached_fingerprint(&key_id, now).as_deref() == Some(incoming_fingerprint.as_str()) {
|
|
||||||
return Ok(false);
|
|
||||||
}
|
|
||||||
|
|
||||||
let Some(key) = state
|
|
||||||
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&key_id))
|
|
||||||
.await?
|
|
||||||
.into_iter()
|
|
||||||
.next()
|
|
||||||
else {
|
|
||||||
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
|
|
||||||
return Ok(false);
|
|
||||||
};
|
|
||||||
|
|
||||||
let Some(provider) = state
|
|
||||||
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&key.provider_id))
|
|
||||||
.await?
|
|
||||||
.into_iter()
|
|
||||||
.next()
|
|
||||||
else {
|
|
||||||
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
|
|
||||||
return Ok(false);
|
|
||||||
};
|
|
||||||
if !provider.provider_type.trim().eq_ignore_ascii_case("codex") {
|
|
||||||
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
|
|
||||||
return Ok(false);
|
|
||||||
}
|
|
||||||
|
|
||||||
let current_codex = key
|
|
||||||
.upstream_metadata
|
|
||||||
.as_ref()
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.and_then(|metadata| metadata.get("codex"))
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.cloned()
|
|
||||||
.unwrap_or_else(serde_json::Map::new);
|
|
||||||
let current_codex = Value::Object(current_codex);
|
|
||||||
let Some(current_fingerprint) = fingerprint_codex_payload(¤t_codex) else {
|
|
||||||
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
|
|
||||||
return Ok(false);
|
|
||||||
};
|
|
||||||
if current_fingerprint == incoming_fingerprint {
|
|
||||||
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
|
|
||||||
return Ok(false);
|
|
||||||
}
|
|
||||||
|
|
||||||
let updated_upstream_metadata =
|
|
||||||
merge_metadata_object(key.upstream_metadata.as_ref(), "codex", parsed);
|
|
||||||
let updated_status_snapshot = sync_provider_key_quota_status_snapshot(
|
|
||||||
key.status_snapshot.as_ref(),
|
|
||||||
provider.provider_type.as_str(),
|
|
||||||
updated_upstream_metadata.as_ref(),
|
|
||||||
"response_headers",
|
|
||||||
);
|
|
||||||
let mut updated_key = key;
|
|
||||||
updated_key.upstream_metadata = updated_upstream_metadata;
|
|
||||||
updated_key.status_snapshot = updated_status_snapshot;
|
|
||||||
updated_key.updated_at_unix_secs = Some(now_unix_secs);
|
|
||||||
|
|
||||||
let updated = state
|
|
||||||
.update_provider_catalog_key(&updated_key)
|
|
||||||
.await?
|
|
||||||
.is_some();
|
|
||||||
if updated {
|
|
||||||
set_cached_fingerprint(&key_id, incoming_fingerprint, now);
|
|
||||||
}
|
|
||||||
Ok(updated)
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
pub(super) fn clear_codex_quota_fingerprint_cache() {
|
|
||||||
if let Some(cache) = HEADER_FINGERPRINT_CACHE.get() {
|
|
||||||
cache
|
|
||||||
.lock()
|
|
||||||
.expect("codex realtime quota cache should lock")
|
|
||||||
.clear();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
#[cfg(test)]
|
|
||||||
mod tests {
|
|
||||||
use super::{clear_codex_quota_fingerprint_cache, sync_codex_quota_from_response_headers};
|
|
||||||
use crate::data::GatewayDataState;
|
|
||||||
use crate::AppState;
|
|
||||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
|
||||||
use aether_data_contracts::repository::provider_catalog::{
|
|
||||||
ProviderCatalogReadRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
|
|
||||||
};
|
|
||||||
use serde_json::{json, Value};
|
|
||||||
use std::collections::BTreeMap;
|
|
||||||
use std::sync::Arc;
|
|
||||||
|
|
||||||
fn build_state(repository: Arc<InMemoryProviderCatalogReadRepository>) -> AppState {
|
|
||||||
AppState::new()
|
|
||||||
.expect("gateway state should build")
|
|
||||||
.with_data_state_for_tests(
|
|
||||||
GatewayDataState::with_provider_catalog_repository_for_tests(repository),
|
|
||||||
)
|
|
||||||
}
|
|
||||||
|
|
||||||
fn sample_provider(provider_id: &str, provider_type: &str) -> StoredProviderCatalogProvider {
|
|
||||||
StoredProviderCatalogProvider::new(
|
|
||||||
provider_id.to_string(),
|
|
||||||
provider_type.to_string(),
|
|
||||||
None,
|
|
||||||
provider_type.to_string(),
|
|
||||||
)
|
|
||||||
.expect("provider should build")
|
|
||||||
}
|
|
||||||
|
|
||||||
fn sample_key(
|
|
||||||
key_id: &str,
|
|
||||||
provider_id: &str,
|
|
||||||
upstream_metadata: Option<Value>,
|
|
||||||
) -> StoredProviderCatalogKey {
|
|
||||||
let mut key = StoredProviderCatalogKey::new(
|
|
||||||
key_id.to_string(),
|
|
||||||
provider_id.to_string(),
|
|
||||||
"default".to_string(),
|
|
||||||
"bearer".to_string(),
|
|
||||||
None,
|
|
||||||
true,
|
|
||||||
)
|
|
||||||
.expect("key should build")
|
|
||||||
.with_transport_fields(
|
|
||||||
Some(json!(["openai:cli"])),
|
|
||||||
"sk-codex-test".to_string(),
|
|
||||||
None,
|
|
||||||
None,
|
|
||||||
None,
|
|
||||||
None,
|
|
||||||
None,
|
|
||||||
None,
|
|
||||||
None,
|
|
||||||
)
|
|
||||||
.expect("key transport should build");
|
|
||||||
key.upstream_metadata = upstream_metadata;
|
|
||||||
key
|
|
||||||
}
|
|
||||||
|
|
||||||
fn quota_snapshot<'a>(key: &'a StoredProviderCatalogKey) -> &'a serde_json::Map<String, Value> {
|
|
||||||
key.status_snapshot
|
|
||||||
.as_ref()
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.and_then(|snapshot| snapshot.get("quota"))
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.expect("quota snapshot should exist")
|
|
||||||
}
|
|
||||||
|
|
||||||
fn paid_headers(
|
|
||||||
primary_used_percent: &str,
|
|
||||||
secondary_used_percent: &str,
|
|
||||||
primary_reset_after_seconds: &str,
|
|
||||||
secondary_reset_after_seconds: &str,
|
|
||||||
) -> BTreeMap<String, String> {
|
|
||||||
BTreeMap::from([
|
|
||||||
("x-codex-plan-type".to_string(), "team".to_string()),
|
|
||||||
(
|
|
||||||
"x-codex-primary-used-percent".to_string(),
|
|
||||||
primary_used_percent.to_string(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"x-codex-secondary-used-percent".to_string(),
|
|
||||||
secondary_used_percent.to_string(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"x-codex-primary-window-minutes".to_string(),
|
|
||||||
"300".to_string(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"x-codex-secondary-window-minutes".to_string(),
|
|
||||||
"10080".to_string(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"x-codex-primary-reset-after-seconds".to_string(),
|
|
||||||
primary_reset_after_seconds.to_string(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"x-codex-secondary-reset-after-seconds".to_string(),
|
|
||||||
secondary_reset_after_seconds.to_string(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"x-codex-primary-reset-at".to_string(),
|
|
||||||
"1776148929".to_string(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"x-codex-secondary-reset-at".to_string(),
|
|
||||||
"1776657828".to_string(),
|
|
||||||
),
|
|
||||||
(
|
|
||||||
"x-codex-credits-has-credits".to_string(),
|
|
||||||
"False".to_string(),
|
|
||||||
),
|
|
||||||
("x-codex-credits-balance".to_string(), "".to_string()),
|
|
||||||
("x-codex-credits-unlimited".to_string(), "False".to_string()),
|
|
||||||
])
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn sync_codex_quota_replaces_existing_codex_fields_and_preserves_other_sections() {
|
|
||||||
clear_codex_quota_fingerprint_cache();
|
|
||||||
|
|
||||||
let mut key = sample_key(
|
|
||||||
"key-codex-1",
|
|
||||||
"provider-codex",
|
|
||||||
Some(json!({
|
|
||||||
"codex": {
|
|
||||||
"legacy_marker": "drop-me",
|
|
||||||
"secondary_used_percent": 2.0,
|
|
||||||
"credits_balance": 42.0,
|
|
||||||
"account_disabled": true,
|
|
||||||
"reason": "deactivated_workspace"
|
|
||||||
},
|
|
||||||
"other": {
|
|
||||||
"value": true
|
|
||||||
}
|
|
||||||
})),
|
|
||||||
);
|
|
||||||
key.status_snapshot = Some(json!({
|
|
||||||
"oauth": {
|
|
||||||
"code": "valid",
|
|
||||||
"label": "有效",
|
|
||||||
"requires_reauth": false,
|
|
||||||
"expiring_soon": false
|
|
||||||
},
|
|
||||||
"account": {
|
|
||||||
"code": "ok",
|
|
||||||
"blocked": false,
|
|
||||||
"recoverable": false
|
|
||||||
},
|
|
||||||
"quota": {
|
|
||||||
"code": "unknown",
|
|
||||||
"exhausted": false
|
|
||||||
}
|
|
||||||
}));
|
|
||||||
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
|
||||||
vec![sample_provider("provider-codex", "codex")],
|
|
||||||
Vec::new(),
|
|
||||||
vec![key],
|
|
||||||
));
|
|
||||||
let state = build_state(Arc::clone(&repository));
|
|
||||||
|
|
||||||
let updated = sync_codex_quota_from_response_headers(
|
|
||||||
&state,
|
|
||||||
Some(&json!({
|
|
||||||
"request_id": "req-codex-realtime-1",
|
|
||||||
"key_id": "key-codex-1"
|
|
||||||
})),
|
|
||||||
&paid_headers("100", "31", "15160", "524059"),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("codex realtime sync should succeed");
|
|
||||||
|
|
||||||
assert!(updated);
|
|
||||||
let reloaded = repository
|
|
||||||
.list_keys_by_ids(&["key-codex-1".to_string()])
|
|
||||||
.await
|
|
||||||
.expect("keys should list");
|
|
||||||
let codex = reloaded[0]
|
|
||||||
.upstream_metadata
|
|
||||||
.as_ref()
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.and_then(|metadata| metadata.get("codex"))
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.expect("codex metadata should exist");
|
|
||||||
assert_eq!(codex.get("plan_type"), Some(&json!("team")));
|
|
||||||
assert_eq!(codex.get("primary_used_percent"), Some(&json!(31.0)));
|
|
||||||
assert_eq!(codex.get("secondary_used_percent"), Some(&json!(100.0)));
|
|
||||||
assert_eq!(codex.get("has_credits"), Some(&json!(false)));
|
|
||||||
assert_eq!(codex.get("credits_unlimited"), Some(&json!(false)));
|
|
||||||
assert!(codex.get("legacy_marker").is_none());
|
|
||||||
assert!(codex.get("credits_balance").is_none());
|
|
||||||
assert!(codex.get("account_disabled").is_none());
|
|
||||||
assert!(codex.get("reason").is_none());
|
|
||||||
assert!(codex.get("updated_at").and_then(Value::as_u64).is_some());
|
|
||||||
assert_eq!(
|
|
||||||
reloaded[0]
|
|
||||||
.upstream_metadata
|
|
||||||
.as_ref()
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.and_then(|metadata| metadata.get("other")),
|
|
||||||
Some(&json!({"value": true}))
|
|
||||||
);
|
|
||||||
let quota = quota_snapshot(&reloaded[0]);
|
|
||||||
assert_eq!(quota.get("version"), Some(&json!(2)));
|
|
||||||
assert_eq!(quota.get("provider_type"), Some(&json!("codex")));
|
|
||||||
assert_eq!(quota.get("source"), Some(&json!("response_headers")));
|
|
||||||
assert_eq!(quota.get("code"), Some(&json!("exhausted")));
|
|
||||||
assert_eq!(quota.get("exhausted"), Some(&json!(true)));
|
|
||||||
assert_eq!(quota.get("plan_type"), Some(&json!("team")));
|
|
||||||
assert_eq!(quota.get("usage_ratio"), Some(&json!(1.0)));
|
|
||||||
assert_eq!(quota.get("updated_at"), quota.get("observed_at"));
|
|
||||||
let windows = quota
|
|
||||||
.get("windows")
|
|
||||||
.and_then(Value::as_array)
|
|
||||||
.expect("windows should be array");
|
|
||||||
assert_eq!(windows.len(), 2);
|
|
||||||
assert_eq!(windows[0].get("code"), Some(&json!("weekly")));
|
|
||||||
assert_eq!(windows[1].get("code"), Some(&json!("5h")));
|
|
||||||
let oauth = reloaded[0]
|
|
||||||
.status_snapshot
|
|
||||||
.as_ref()
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.and_then(|snapshot| snapshot.get("oauth"))
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.expect("oauth snapshot should exist");
|
|
||||||
assert_eq!(oauth.get("code"), Some(&json!("valid")));
|
|
||||||
let account = reloaded[0]
|
|
||||||
.status_snapshot
|
|
||||||
.as_ref()
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.and_then(|snapshot| snapshot.get("account"))
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.expect("account snapshot should exist");
|
|
||||||
assert_eq!(account.get("code"), Some(&json!("ok")));
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn sync_codex_quota_skips_non_codex_provider() {
|
|
||||||
clear_codex_quota_fingerprint_cache();
|
|
||||||
|
|
||||||
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
|
||||||
vec![sample_provider("provider-openai", "openai")],
|
|
||||||
Vec::new(),
|
|
||||||
vec![sample_key("key-openai-1", "provider-openai", None)],
|
|
||||||
));
|
|
||||||
let state = build_state(Arc::clone(&repository));
|
|
||||||
|
|
||||||
let updated = sync_codex_quota_from_response_headers(
|
|
||||||
&state,
|
|
||||||
Some(&json!({
|
|
||||||
"request_id": "req-openai-realtime-1",
|
|
||||||
"key_id": "key-openai-1"
|
|
||||||
})),
|
|
||||||
&paid_headers("100", "31", "15160", "524059"),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("non-codex realtime sync should not fail");
|
|
||||||
|
|
||||||
assert!(!updated);
|
|
||||||
let reloaded = repository
|
|
||||||
.list_keys_by_ids(&["key-openai-1".to_string()])
|
|
||||||
.await
|
|
||||||
.expect("keys should list");
|
|
||||||
assert_eq!(reloaded[0].upstream_metadata, None);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn sync_codex_quota_skips_when_headers_do_not_contain_codex_metadata() {
|
|
||||||
clear_codex_quota_fingerprint_cache();
|
|
||||||
|
|
||||||
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
|
||||||
vec![sample_provider("provider-codex", "codex")],
|
|
||||||
Vec::new(),
|
|
||||||
vec![sample_key("key-codex-2", "provider-codex", None)],
|
|
||||||
));
|
|
||||||
let state = build_state(Arc::clone(&repository));
|
|
||||||
|
|
||||||
let updated = sync_codex_quota_from_response_headers(
|
|
||||||
&state,
|
|
||||||
Some(&json!({
|
|
||||||
"request_id": "req-codex-realtime-2",
|
|
||||||
"key_id": "key-codex-2"
|
|
||||||
})),
|
|
||||||
&BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("empty codex headers should not fail");
|
|
||||||
|
|
||||||
assert!(!updated);
|
|
||||||
let reloaded = repository
|
|
||||||
.list_keys_by_ids(&["key-codex-2".to_string()])
|
|
||||||
.await
|
|
||||||
.expect("keys should list");
|
|
||||||
assert_eq!(reloaded[0].upstream_metadata, None);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn sync_codex_quota_cache_hit_skips_when_only_reset_countdown_changes() {
|
|
||||||
clear_codex_quota_fingerprint_cache();
|
|
||||||
|
|
||||||
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
|
||||||
vec![sample_provider("provider-codex", "codex")],
|
|
||||||
Vec::new(),
|
|
||||||
vec![sample_key("key-codex-3", "provider-codex", None)],
|
|
||||||
));
|
|
||||||
let state = build_state(Arc::clone(&repository));
|
|
||||||
let key_ids = ["key-codex-3".to_string()];
|
|
||||||
|
|
||||||
let first_updated = sync_codex_quota_from_response_headers(
|
|
||||||
&state,
|
|
||||||
Some(&json!({
|
|
||||||
"request_id": "req-codex-realtime-3",
|
|
||||||
"key_id": "key-codex-3"
|
|
||||||
})),
|
|
||||||
&paid_headers("100", "31", "15160", "524059"),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("first realtime sync should succeed");
|
|
||||||
assert!(first_updated);
|
|
||||||
let first_snapshot = repository
|
|
||||||
.list_keys_by_ids(&key_ids)
|
|
||||||
.await
|
|
||||||
.expect("keys should list")
|
|
||||||
.into_iter()
|
|
||||||
.next()
|
|
||||||
.expect("key should exist");
|
|
||||||
|
|
||||||
let second_updated = sync_codex_quota_from_response_headers(
|
|
||||||
&state,
|
|
||||||
Some(&json!({
|
|
||||||
"request_id": "req-codex-realtime-3",
|
|
||||||
"key_id": "key-codex-3"
|
|
||||||
})),
|
|
||||||
&paid_headers("100", "31", "42", "84"),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("second realtime sync should succeed");
|
|
||||||
assert!(!second_updated);
|
|
||||||
let second_snapshot = repository
|
|
||||||
.list_keys_by_ids(&key_ids)
|
|
||||||
.await
|
|
||||||
.expect("keys should list")
|
|
||||||
.into_iter()
|
|
||||||
.next()
|
|
||||||
.expect("key should exist");
|
|
||||||
assert_eq!(
|
|
||||||
first_snapshot.upstream_metadata,
|
|
||||||
second_snapshot.upstream_metadata
|
|
||||||
);
|
|
||||||
assert_eq!(
|
|
||||||
first_snapshot.updated_at_unix_secs,
|
|
||||||
second_snapshot.updated_at_unix_secs
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
|
||||||
async fn sync_codex_quota_updates_when_usage_percent_changes() {
|
|
||||||
clear_codex_quota_fingerprint_cache();
|
|
||||||
|
|
||||||
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
|
||||||
vec![sample_provider("provider-codex", "codex")],
|
|
||||||
Vec::new(),
|
|
||||||
vec![sample_key("key-codex-4", "provider-codex", None)],
|
|
||||||
));
|
|
||||||
let state = build_state(Arc::clone(&repository));
|
|
||||||
|
|
||||||
let first_updated = sync_codex_quota_from_response_headers(
|
|
||||||
&state,
|
|
||||||
Some(&json!({
|
|
||||||
"request_id": "req-codex-realtime-4",
|
|
||||||
"key_id": "key-codex-4"
|
|
||||||
})),
|
|
||||||
&paid_headers("98", "31", "15160", "524059"),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("first realtime sync should succeed");
|
|
||||||
assert!(first_updated);
|
|
||||||
|
|
||||||
let second_updated = sync_codex_quota_from_response_headers(
|
|
||||||
&state,
|
|
||||||
Some(&json!({
|
|
||||||
"request_id": "req-codex-realtime-4",
|
|
||||||
"key_id": "key-codex-4"
|
|
||||||
})),
|
|
||||||
&paid_headers("100", "31", "15160", "524059"),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.expect("second realtime sync should succeed");
|
|
||||||
assert!(second_updated);
|
|
||||||
|
|
||||||
let reloaded = repository
|
|
||||||
.list_keys_by_ids(&["key-codex-4".to_string()])
|
|
||||||
.await
|
|
||||||
.expect("keys should list");
|
|
||||||
let codex = reloaded[0]
|
|
||||||
.upstream_metadata
|
|
||||||
.as_ref()
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.and_then(|metadata| metadata.get("codex"))
|
|
||||||
.and_then(Value::as_object)
|
|
||||||
.expect("codex metadata should exist");
|
|
||||||
assert_eq!(codex.get("secondary_used_percent"), Some(&json!(100.0)));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
@@ -4,22 +4,20 @@ use aether_contracts::ExecutionError;
|
|||||||
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
|
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
|
||||||
use aether_scheduler_core::{execution_error_details, SchedulerRequestCandidateStatusUpdate};
|
use aether_scheduler_core::{execution_error_details, SchedulerRequestCandidateStatusUpdate};
|
||||||
use tracing::{debug, warn};
|
use tracing::{debug, warn};
|
||||||
use uuid::Uuid;
|
|
||||||
|
|
||||||
use crate::clock::{current_unix_ms, current_unix_secs};
|
use crate::clock::current_unix_ms;
|
||||||
use crate::log_ids::short_request_id;
|
use crate::log_ids::short_request_id;
|
||||||
|
use crate::orchestration::{apply_local_report_effect, LocalReportEffect};
|
||||||
use crate::request_candidate_runtime::record_report_request_candidate_status;
|
use crate::request_candidate_runtime::record_report_request_candidate_status;
|
||||||
use crate::{AppState, GatewayError};
|
use crate::{AppState, GatewayError};
|
||||||
|
|
||||||
mod codex_realtime_quota;
|
|
||||||
mod context;
|
mod context;
|
||||||
use context::{report_context_is_locally_actionable, resolve_locally_actionable_report_context};
|
use context::{report_context_is_locally_actionable, resolve_locally_actionable_report_context};
|
||||||
|
|
||||||
use aether_usage_runtime::{
|
use aether_usage_runtime::{
|
||||||
extract_gemini_file_mapping_entries, gemini_file_mapping_cache_key,
|
is_local_ai_stream_report_kind, is_local_ai_sync_report_kind, report_request_id,
|
||||||
is_local_ai_stream_report_kind, is_local_ai_sync_report_kind, normalize_gemini_file_name,
|
should_handle_local_stream_report, should_handle_local_sync_report,
|
||||||
report_request_id, should_handle_local_stream_report, should_handle_local_sync_report,
|
sync_report_represents_failure,
|
||||||
sync_report_represents_failure, GEMINI_FILE_MAPPING_TTL_SECONDS,
|
|
||||||
};
|
};
|
||||||
pub(crate) use aether_usage_runtime::{GatewayStreamReportRequest, GatewaySyncReportRequest};
|
pub(crate) use aether_usage_runtime::{GatewayStreamReportRequest, GatewaySyncReportRequest};
|
||||||
|
|
||||||
@@ -188,7 +186,6 @@ pub(crate) async fn submit_stream_report(
|
|||||||
}
|
}
|
||||||
|
|
||||||
async fn handle_local_sync_report(state: &AppState, payload: &GatewaySyncReportRequest) {
|
async fn handle_local_sync_report(state: &AppState, payload: &GatewaySyncReportRequest) {
|
||||||
apply_local_gemini_file_mapping_side_effect(state, payload).await;
|
|
||||||
let terminal_unix_ms = current_unix_ms();
|
let terminal_unix_ms = current_unix_ms();
|
||||||
let (error_type, error_message) =
|
let (error_type, error_message) =
|
||||||
execution_error_details(None::<&ExecutionError>, payload.body_json.as_ref());
|
execution_error_details(None::<&ExecutionError>, payload.body_json.as_ref());
|
||||||
@@ -215,22 +212,7 @@ async fn handle_local_sync_report(state: &AppState, payload: &GatewaySyncReportR
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
if let Err(err) = codex_realtime_quota::sync_codex_quota_from_response_headers(
|
apply_local_report_effect(state, LocalReportEffect::Sync { payload }).await;
|
||||||
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 handle_local_stream_report(state: &AppState, payload: &GatewayStreamReportRequest) {
|
async fn handle_local_stream_report(state: &AppState, payload: &GatewayStreamReportRequest) {
|
||||||
@@ -253,154 +235,7 @@ async fn handle_local_stream_report(state: &AppState, payload: &GatewayStreamRep
|
|||||||
},
|
},
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
if let Err(err) = codex_realtime_quota::sync_codex_quota_from_response_headers(
|
apply_local_report_effect(state, LocalReportEffect::Stream { payload }).await;
|
||||||
state,
|
|
||||||
payload.report_context.as_ref(),
|
|
||||||
&payload.headers,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
warn!(
|
|
||||||
event_name = "codex_realtime_quota_sync_failed",
|
|
||||||
log_type = "ops",
|
|
||||||
report_kind = %payload.report_kind,
|
|
||||||
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
|
|
||||||
error = ?err,
|
|
||||||
"gateway failed to persist codex realtime quota from stream response headers"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn apply_local_gemini_file_mapping_side_effect(
|
|
||||||
state: &AppState,
|
|
||||||
payload: &GatewaySyncReportRequest,
|
|
||||||
) {
|
|
||||||
match payload.report_kind.as_str() {
|
|
||||||
"gemini_files_store_mapping" => {
|
|
||||||
if payload.status_code >= 300 {
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
let key_id = payload
|
|
||||||
.report_context
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|context| context.get("file_key_id"))
|
|
||||||
.and_then(serde_json::Value::as_str)
|
|
||||||
.map(str::trim)
|
|
||||||
.filter(|value| !value.is_empty());
|
|
||||||
let user_id = payload
|
|
||||||
.report_context
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|context| context.get("user_id"))
|
|
||||||
.and_then(serde_json::Value::as_str)
|
|
||||||
.map(str::trim)
|
|
||||||
.filter(|value| !value.is_empty());
|
|
||||||
let Some(key_id) = key_id else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
|
|
||||||
for entry in extract_gemini_file_mapping_entries(payload) {
|
|
||||||
if let Err(err) = store_local_gemini_file_mapping(
|
|
||||||
state,
|
|
||||||
entry.file_name.as_str(),
|
|
||||||
key_id,
|
|
||||||
user_id,
|
|
||||||
entry.display_name.as_deref(),
|
|
||||||
entry.mime_type.as_deref(),
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
{
|
|
||||||
warn!(
|
|
||||||
event_name = "gemini_file_mapping_store_failed",
|
|
||||||
log_type = "ops",
|
|
||||||
report_kind = %payload.report_kind,
|
|
||||||
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
|
|
||||||
file_name = %entry.file_name,
|
|
||||||
error = ?err,
|
|
||||||
"gateway failed to persist gemini file mapping locally"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
"gemini_files_delete_mapping" if payload.status_code < 300 => {
|
|
||||||
let file_name = payload
|
|
||||||
.report_context
|
|
||||||
.as_ref()
|
|
||||||
.and_then(|context| context.get("file_name"))
|
|
||||||
.and_then(serde_json::Value::as_str)
|
|
||||||
.and_then(normalize_gemini_file_name);
|
|
||||||
let Some(file_name) = file_name else {
|
|
||||||
return;
|
|
||||||
};
|
|
||||||
|
|
||||||
if let Err(err) = delete_local_gemini_file_mapping(state, file_name.as_str()).await {
|
|
||||||
warn!(
|
|
||||||
event_name = "gemini_file_mapping_delete_failed",
|
|
||||||
log_type = "ops",
|
|
||||||
report_kind = %payload.report_kind,
|
|
||||||
report_request_id = %short_request_id(report_request_id(payload.report_context.as_ref())),
|
|
||||||
file_name = %file_name,
|
|
||||||
error = ?err,
|
|
||||||
"gateway failed to delete gemini file mapping locally"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
_ => {}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(crate) async fn store_local_gemini_file_mapping(
|
|
||||||
state: &AppState,
|
|
||||||
file_name: &str,
|
|
||||||
key_id: &str,
|
|
||||||
user_id: Option<&str>,
|
|
||||||
display_name: Option<&str>,
|
|
||||||
mime_type: Option<&str>,
|
|
||||||
) -> Result<(), GatewayError> {
|
|
||||||
let Some(file_name) = normalize_gemini_file_name(file_name) else {
|
|
||||||
return Ok(());
|
|
||||||
};
|
|
||||||
let expires_at_unix_secs = current_unix_secs().saturating_add(GEMINI_FILE_MAPPING_TTL_SECONDS);
|
|
||||||
|
|
||||||
let _stored = state
|
|
||||||
.upsert_gemini_file_mapping(
|
|
||||||
aether_data::repository::gemini_file_mappings::UpsertGeminiFileMappingRecord {
|
|
||||||
id: Uuid::new_v4().to_string(),
|
|
||||||
file_name: file_name.clone(),
|
|
||||||
key_id: key_id.to_string(),
|
|
||||||
user_id: user_id.map(ToOwned::to_owned),
|
|
||||||
display_name: display_name.map(ToOwned::to_owned),
|
|
||||||
mime_type: mime_type.map(ToOwned::to_owned),
|
|
||||||
source_hash: None,
|
|
||||||
expires_at_unix_secs,
|
|
||||||
},
|
|
||||||
)
|
|
||||||
.await?;
|
|
||||||
state
|
|
||||||
.cache_set_string_with_ttl(
|
|
||||||
gemini_file_mapping_cache_key(file_name.as_str()).as_str(),
|
|
||||||
key_id,
|
|
||||||
GEMINI_FILE_MAPPING_TTL_SECONDS,
|
|
||||||
)
|
|
||||||
.await?;
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn delete_local_gemini_file_mapping(
|
|
||||||
state: &AppState,
|
|
||||||
file_name: &str,
|
|
||||||
) -> Result<(), GatewayError> {
|
|
||||||
let Some(file_name) = normalize_gemini_file_name(file_name) else {
|
|
||||||
return Ok(());
|
|
||||||
};
|
|
||||||
|
|
||||||
let _deleted = state
|
|
||||||
.delete_gemini_file_mapping_by_file_name(file_name.as_str())
|
|
||||||
.await?;
|
|
||||||
state
|
|
||||||
.cache_delete_key(gemini_file_mapping_cache_key(file_name.as_str()).as_str())
|
|
||||||
.await?;
|
|
||||||
Ok(())
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
@@ -824,7 +659,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn submit_sync_report_updates_codex_quota_from_response_headers() {
|
async fn submit_sync_report_updates_codex_quota_from_response_headers() {
|
||||||
super::codex_realtime_quota::clear_codex_quota_fingerprint_cache();
|
crate::orchestration::clear_local_report_effect_caches_for_tests();
|
||||||
|
|
||||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||||
vec![sample_provider_catalog_provider(
|
vec![sample_provider_catalog_provider(
|
||||||
@@ -888,7 +723,7 @@ mod tests {
|
|||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn submit_stream_report_updates_codex_quota_from_response_headers() {
|
async fn submit_stream_report_updates_codex_quota_from_response_headers() {
|
||||||
super::codex_realtime_quota::clear_codex_quota_fingerprint_cache();
|
crate::orchestration::clear_local_report_effect_caches_for_tests();
|
||||||
|
|
||||||
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||||
vec![sample_provider_catalog_provider(
|
vec![sample_provider_catalog_provider(
|
||||||
|
|||||||
@@ -473,6 +473,24 @@ pub fn codex_structured_invalid_reason(status_code: u16, upstream_message: Optio
|
|||||||
message.to_string()
|
message.to_string()
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn codex_runtime_invalid_reason(
|
||||||
|
status_code: u16,
|
||||||
|
upstream_message: Option<&str>,
|
||||||
|
) -> Option<String> {
|
||||||
|
match status_code {
|
||||||
|
401 => Some(codex_structured_invalid_reason(401, upstream_message)),
|
||||||
|
402 if codex_looks_like_workspace_deactivated(upstream_message) => {
|
||||||
|
Some(codex_structured_invalid_reason(402, upstream_message))
|
||||||
|
}
|
||||||
|
403 if codex_looks_like_token_invalidated(upstream_message)
|
||||||
|
|| codex_looks_like_account_deactivated(upstream_message) =>
|
||||||
|
{
|
||||||
|
Some(codex_structured_invalid_reason(403, upstream_message))
|
||||||
|
}
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub fn codex_soft_request_failure_reason(
|
pub fn codex_soft_request_failure_reason(
|
||||||
status_code: u16,
|
status_code: u16,
|
||||||
upstream_message: Option<&str>,
|
upstream_message: Option<&str>,
|
||||||
@@ -634,3 +652,31 @@ pub fn parse_kiro_usage_response(
|
|||||||
|
|
||||||
Some(serde_json::Value::Object(result))
|
Some(serde_json::Value::Object(result))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::{codex_runtime_invalid_reason, OAUTH_ACCOUNT_BLOCK_PREFIX, OAUTH_EXPIRED_PREFIX};
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn codex_runtime_invalid_reason_marks_401_as_expired() {
|
||||||
|
assert_eq!(
|
||||||
|
codex_runtime_invalid_reason(401, Some("session expired")),
|
||||||
|
Some(format!("{OAUTH_EXPIRED_PREFIX}session expired"))
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn codex_runtime_invalid_reason_marks_account_deactivated_403() {
|
||||||
|
assert_eq!(
|
||||||
|
codex_runtime_invalid_reason(403, Some("account has been deactivated")),
|
||||||
|
Some(format!(
|
||||||
|
"{OAUTH_ACCOUNT_BLOCK_PREFIX}account has been deactivated"
|
||||||
|
))
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn codex_runtime_invalid_reason_ignores_generic_403() {
|
||||||
|
assert_eq!(codex_runtime_invalid_reason(403, Some("forbidden")), None);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user