i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-08 18:37:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(gateway): harden provider request execution

Browse Source 
Preserve exact request payloads and model client surface and API operation explicitly.

Add Anthropic compatibility profiles, bounded stream commitment, and scoped OAuth retry behavior across provider transports.
This commit is contained in:
elky
2026-07-27 09:36:31 +08:00
parent 79b70f7b5c
commit 531cf11025
152 changed files with 13984 additions and 2075 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/ai_serving/planner/candidate_materialization.rs
+250 -8
View File
@@ -68,6 +68,8 @@ pub(crate) struct LocalExecutionCandidateAttempt {
pub(crate) struct LocalExecutionCandidateAttemptSource<'a> {
items: VecDeque<LocalExecutionCandidateAttemptSourceItem<'a>>,
skipped_provider_ids: BTreeSet<String>,
skipped_endpoint_ids: BTreeSet<String>,
skipped_credential_ids: BTreeSet<String>,
}
type DecorateSkippedCandidateFn<'a> = Arc<
@@ -80,6 +82,10 @@ pub(crate) trait LocalExecutionAttemptSource<T>: Send {
async fn drain_execution_attempts(&mut self) -> Result<Vec<T>, GatewayError>;
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError>;
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError>;
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError>;
}
@@ -111,6 +117,8 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
Self {
items,
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
}
}
@@ -123,7 +131,12 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
};
match front {
LocalExecutionCandidateAttemptSourceItem::Static { attempts } => {
if dispatch_sequence_provider_is_skipped(attempts, &self.skipped_provider_ids) {
if dispatch_sequence_candidate_is_skipped(
attempts,
&self.skipped_provider_ids,
&self.skipped_endpoint_ids,
&self.skipped_credential_ids,
) {
self.items.pop_front();
continue;
}
@@ -141,10 +154,20 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
pending_attempts,
pool_exhaustion_persistence,
} => {
if self.skipped_provider_ids.contains(cursor.provider_id()) {
if self.skipped_provider_ids.contains(cursor.provider_id())
|| self.skipped_endpoint_ids.contains(cursor.endpoint_id())
{
self.items.pop_front();
continue;
}
if dispatch_sequence_candidate_is_skipped(
pending_attempts,
&self.skipped_provider_ids,
&self.skipped_endpoint_ids,
&self.skipped_credential_ids,
) {
*pending_attempts = DispatchSequence::new(Vec::new());
}
if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) {
return Ok(Some(attempt));
}
@@ -162,6 +185,14 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
self.items.pop_front();
continue;
};
if candidate_is_skipped(
&candidate,
&self.skipped_provider_ids,
&self.skipped_endpoint_ids,
&self.skipped_credential_ids,
) {
continue;
}
*pending_attempts = dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
candidate,
@@ -174,6 +205,12 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
for provider_id in &self.skipped_provider_ids {
cursor.skip_provider(provider_id);
}
for endpoint_id in &self.skipped_endpoint_ids {
cursor.skip_endpoint(endpoint_id);
}
for key_id in &self.skipped_credential_ids {
cursor.skip_credential(key_id);
}
let Some(attempt) = cursor.next_attempt().await? else {
self.items.pop_front();
continue;
@@ -201,6 +238,32 @@ impl<'a> LocalExecutionCandidateAttemptSource<'a> {
}
}
}
pub(crate) fn skip_endpoint(&mut self, endpoint_id: &str) {
let endpoint_id = endpoint_id.trim();
if endpoint_id.is_empty() {
return;
}
self.skipped_endpoint_ids.insert(endpoint_id.to_string());
for item in &mut self.items {
if let LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { cursor } = item {
cursor.skip_endpoint(endpoint_id);
}
}
}
pub(crate) fn skip_credential(&mut self, key_id: &str) {
let key_id = key_id.trim();
if key_id.is_empty() {
return;
}
self.skipped_credential_ids.insert(key_id.to_string());
for item in &mut self.items {
if let LocalExecutionCandidateAttemptSourceItem::RequestedModelPage { cursor } = item {
cursor.skip_credential(key_id);
}
}
}
}
impl LocalExecutionCandidateAttempt {
@@ -668,6 +731,8 @@ where
LocalExecutionCandidateAttemptSource {
items,
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
},
candidate_count,
)
@@ -802,6 +867,8 @@ where
page_cursor,
pending_items: VecDeque::new(),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
candidate_count: 0,
next_candidate_index: 0,
remembered_affinity: false,
@@ -825,6 +892,8 @@ where
LocalExecutionCandidateAttemptSource {
items,
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
},
candidate_count,
)
@@ -851,6 +920,8 @@ struct RequestedModelAttemptPageCursor<'a> {
page_cursor: LocalCandidatePreselectionPageCursor<'a>,
pending_items: VecDeque<LocalExecutionCandidateAttemptSourceItem<'a>>,
skipped_provider_ids: BTreeSet<String>,
skipped_endpoint_ids: BTreeSet<String>,
skipped_credential_ids: BTreeSet<String>,
candidate_count: usize,
next_candidate_index: u32,
remembered_affinity: bool,
@@ -864,6 +935,14 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
self.skipped_provider_ids.insert(provider_id.to_string());
}
fn skip_endpoint(&mut self, endpoint_id: &str) {
self.skipped_endpoint_ids.insert(endpoint_id.to_string());
}
fn skip_credential(&mut self, key_id: &str) {
self.skipped_credential_ids.insert(key_id.to_string());
}
async fn next_attempt(
&mut self,
) -> Result<Option<LocalExecutionCandidateAttempt>, GatewayError> {
@@ -871,8 +950,13 @@ impl<'a> RequestedModelAttemptPageCursor<'a> {
return Err(error);
}
loop {
if let Some(attempt) =
pop_attempt_from_items(&mut self.pending_items, &self.skipped_provider_ids).await
if let Some(attempt) = pop_attempt_from_items(
&mut self.pending_items,
&self.skipped_provider_ids,
&self.skipped_endpoint_ids,
&self.skipped_credential_ids,
)
.await
{
return Ok(Some(attempt));
}
@@ -1075,12 +1159,19 @@ fn page_is_exact_auth_api_key_concurrency_limited(
async fn pop_attempt_from_items(
items: &mut VecDeque<LocalExecutionCandidateAttemptSourceItem<'_>>,
skipped_provider_ids: &BTreeSet<String>,
skipped_endpoint_ids: &BTreeSet<String>,
skipped_credential_ids: &BTreeSet<String>,
) -> Option<LocalExecutionCandidateAttempt> {
loop {
let front = items.front_mut()?;
match front {
LocalExecutionCandidateAttemptSourceItem::Static { attempts } => {
if dispatch_sequence_provider_is_skipped(attempts, skipped_provider_ids) {
if dispatch_sequence_candidate_is_skipped(
attempts,
skipped_provider_ids,
skipped_endpoint_ids,
skipped_credential_ids,
) {
items.pop_front();
continue;
}
@@ -1098,10 +1189,20 @@ async fn pop_attempt_from_items(
pending_attempts,
pool_exhaustion_persistence,
} => {
if skipped_provider_ids.contains(cursor.provider_id()) {
if skipped_provider_ids.contains(cursor.provider_id())
|| skipped_endpoint_ids.contains(cursor.endpoint_id())
{
items.pop_front();
continue;
}
if dispatch_sequence_candidate_is_skipped(
pending_attempts,
skipped_provider_ids,
skipped_endpoint_ids,
skipped_credential_ids,
) {
*pending_attempts = DispatchSequence::new(Vec::new());
}
if let Some(attempt) = next_attempt_from_dispatch_sequence(pending_attempts) {
return Some(attempt);
}
@@ -1119,6 +1220,14 @@ async fn pop_attempt_from_items(
items.pop_front();
continue;
};
if candidate_is_skipped(
&candidate,
skipped_provider_ids,
skipped_endpoint_ids,
skipped_credential_ids,
) {
continue;
}
*pending_attempts = dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
candidate,
@@ -1780,15 +1889,33 @@ fn next_attempt_from_dispatch_sequence(
Some(attempt)
}
fn dispatch_sequence_provider_is_skipped(
fn dispatch_sequence_candidate_is_skipped(
sequence: &DispatchSequence<LocalExecutionCandidateAttempt>,
skipped_provider_ids: &BTreeSet<String>,
skipped_endpoint_ids: &BTreeSet<String>,
skipped_credential_ids: &BTreeSet<String>,
) -> bool {
sequence.peek_current().is_some_and(|item| {
skipped_provider_ids.contains(&item.candidate.eligible.candidate.provider_id)
candidate_is_skipped(
&item.candidate.eligible,
skipped_provider_ids,
skipped_endpoint_ids,
skipped_credential_ids,
)
})
}
fn candidate_is_skipped(
candidate: &EligibleLocalExecutionCandidate,
skipped_provider_ids: &BTreeSet<String>,
skipped_endpoint_ids: &BTreeSet<String>,
skipped_credential_ids: &BTreeSet<String>,
) -> bool {
skipped_provider_ids.contains(&candidate.candidate.provider_id)
|| skipped_endpoint_ids.contains(&candidate.candidate.endpoint_id)
|| skipped_credential_ids.contains(&candidate.candidate.key_id)
}
fn dispatch_sequence_exhausted(
sequence: &mut DispatchSequence<LocalExecutionCandidateAttempt>,
) -> bool {
@@ -2341,6 +2468,8 @@ mod tests {
page_cursor,
pending_items: VecDeque::new(),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
candidate_count: 0,
next_candidate_index: 0,
remembered_affinity: false,
@@ -2436,6 +2565,8 @@ mod tests {
page_cursor,
pending_items: VecDeque::new(),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
candidate_count: 0,
next_candidate_index: 0,
remembered_affinity: false,
@@ -2634,6 +2765,8 @@ mod tests {
),
}]),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
};
let first = source
@@ -2652,6 +2785,111 @@ mod tests {
.is_none());
}
#[tokio::test]
async fn dynamic_attempt_source_skips_credentials_and_endpoints_across_static_candidates() {
let key_a = sample_eligible("key-a", None);
let key_b = sample_eligible("key-b", None);
let mut key_c = sample_eligible("key-c", None);
key_c.candidate.endpoint_id = "endpoint-2".to_string();
Arc::make_mut(&mut key_c.transport).endpoint.id = "endpoint-2".to_string();
let static_item =
|candidate, candidate_index| LocalExecutionCandidateAttemptSourceItem::Static {
attempts: dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
candidate,
candidate_index,
)
.into(),
),
};
let mut source = LocalExecutionCandidateAttemptSource {
items: VecDeque::from([
static_item(key_a, 0),
static_item(key_b, 1),
static_item(key_c, 2),
]),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
};
source.skip_credential("key-a");
let key_b_attempt = source
.next_attempt()
.await
.expect("candidate source should succeed")
.expect("a different credential should remain");
assert_eq!(key_b_attempt.eligible.candidate.key_id, "key-b");
source.skip_endpoint("endpoint-1");
let endpoint_2_attempt = source
.next_attempt()
.await
.expect("candidate source should succeed")
.expect("a different endpoint should remain");
assert_eq!(endpoint_2_attempt.eligible.candidate.key_id, "key-c");
assert_eq!(
endpoint_2_attempt.eligible.candidate.endpoint_id,
"endpoint-2"
);
}
#[tokio::test]
async fn dynamic_attempt_source_filters_skipped_pool_pending_credential() {
let app = AppState::new().expect("state should build");
let mut pool_group = sample_eligible("pool-group", None);
pool_group.kind = LocalExecutionCandidateKind::PoolGroup;
pool_group.transport = sample_transport("pool-group", Some(json!({ "pool_advanced": {} })));
let pool_cursor = PoolKeyCursor::new(
PlannerAppState::new(&app),
pool_group,
None,
Some("gpt-5"),
None,
);
let pool_key_attempts = dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(
sample_eligible("pool-key-a", None),
0,
)
.into(),
);
let mut fallback = sample_eligible("fallback-key", None);
fallback.candidate.provider_id = "provider-b".to_string();
Arc::make_mut(&mut fallback.transport).provider.id = "provider-b".to_string();
Arc::make_mut(&mut fallback.transport).key.provider_id = "provider-b".to_string();
let fallback_attempts = dispatch_sequence_from_attempts(
build_unpersisted_local_execution_candidate_attempts(fallback, 1).into(),
);
let mut source = LocalExecutionCandidateAttemptSource {
items: VecDeque::from([
LocalExecutionCandidateAttemptSourceItem::Pool {
cursor: pool_cursor,
candidate_index: 0,
pending_attempts: pool_key_attempts,
pool_exhaustion_persistence: None,
},
LocalExecutionCandidateAttemptSourceItem::Static {
attempts: fallback_attempts,
},
]),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
};
source.skip_credential("pool-key-a");
let attempt = source
.next_attempt()
.await
.expect("candidate source should succeed")
.expect("fallback credential should remain");
assert_eq!(attempt.eligible.candidate.provider_id, "provider-b");
assert_eq!(attempt.eligible.candidate.key_id, "fallback-key");
}
#[tokio::test]
async fn skipped_provider_discards_pool_cursor_and_continues_with_next_provider() {
let app = AppState::new().expect("state should build");
@@ -2686,6 +2924,8 @@ mod tests {
},
]),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
};
source.skip_provider("provider-1");
@@ -2752,6 +2992,8 @@ mod tests {
pool_exhaustion_persistence: Some(pool_exhaustion_persistence),
}]),
skipped_provider_ids: BTreeSet::new(),
skipped_endpoint_ids: BTreeSet::new(),
skipped_credential_ids: BTreeSet::new(),
};
assert!(source