i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-06 17:37:47 +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/executor/candidate_loop.rs
+219 -39
View File
@@ -1,7 +1,8 @@
use std::collections::{BTreeMap, BTreeSet};
use aether_ai_serving::{
run_ai_attempt_loop, AiAttemptLoopOutcome, AiAttemptLoopPort, AiExecutionAttempt,
run_ai_attempt_loop, AiAttemptExecutionOutcome, AiAttemptLoopOutcome, AiAttemptLoopPort,
AiAttemptRetryScope, AiExecutionAttempt,
};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_runtime::ConcurrencyPermit;
@@ -18,8 +19,13 @@ use tracing::{debug, warn, Instrument};
use crate::ai_serving::LocalExecutionAttemptSource;
use crate::clock::current_unix_ms;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::{execute_execution_runtime_stream, execute_execution_runtime_sync};
use crate::executor::{build_local_execution_exhaustion, LocalExecutionRequestOutcome};
use crate::execution_runtime::{
execute_execution_runtime_stream_with_retry_scope,
execute_execution_runtime_sync_with_retry_scope,
};
use crate::executor::{
build_local_execution_exhaustion, mark_deferred_upstream_response, LocalExecutionRequestOutcome,
};
use crate::handlers::shared::provider_pool::release_admin_provider_pool_key_lease;
use crate::log_ids::short_request_id;
use crate::orchestration::{
@@ -124,6 +130,9 @@ where
AiAttemptLoopOutcome::Responded(response) => {
Ok(LocalExecutionRequestOutcome::responded(response))
}
AiAttemptLoopOutcome::Deferred(response) => Ok(
LocalExecutionRequestOutcome::responded(mark_deferred_upstream_response(response)),
),
AiAttemptLoopOutcome::Exhausted(exhaustion) => {
Ok(LocalExecutionRequestOutcome::Exhausted(exhaustion))
}
@@ -251,7 +260,10 @@ where
Ok(())
}
async fn execute_attempt(&self, attempt: &T) -> Result<Option<Self::Response>, Self::Error> {
async fn execute_attempt(
&self,
attempt: &T,
) -> Result<AiAttemptExecutionOutcome<Self::Response>, Self::Error> {
let plan = attempt.execution_plan();
let report_context = attempt.report_context();
if let Some(response) = execution_plan_balance_capacity_response(
@@ -263,12 +275,12 @@ where
)
.await?
{
return Ok(Some(response));
return Ok(AiAttemptExecutionOutcome::Responded(response));
}
prewarm_direct_reqwest_candidate_client(plan);
let _permit = acquire_upstream_execution_gate(self.state, self.trace_id).await?;
let upstream_execution_gate_held_started_at = std::time::Instant::now();
let mut response = execute_execution_runtime_sync(
let mut execution = execute_execution_runtime_sync_with_retry_scope(
self.state,
self.parts.uri.path(),
plan.clone(),
@@ -285,10 +297,18 @@ where
.elapsed()
.as_millis() as u64,
);
if let Some(response) = response.as_mut() {
attach_redaction_execution_candidate(response, plan.candidate_id.as_deref());
match &mut execution {
AiAttemptExecutionOutcome::Responded(response)
| AiAttemptExecutionOutcome::Retry {
fallback_response: Some(response),
..
} => attach_redaction_execution_candidate(response, plan.candidate_id.as_deref()),
AiAttemptExecutionOutcome::Retry {
fallback_response: None,
..
} => {}
}
Ok(response)
Ok(execution)
}
async fn mark_unused_attempts(&self, attempts: Vec<T>) -> Result<(), Self::Error> {
@@ -389,6 +409,9 @@ where
AiAttemptLoopOutcome::Responded(response) => {
Ok(LocalExecutionRequestOutcome::responded(response))
}
AiAttemptLoopOutcome::Deferred(response) => Ok(
LocalExecutionRequestOutcome::responded(mark_deferred_upstream_response(response)),
),
AiAttemptLoopOutcome::Exhausted(exhaustion) => {
Ok(LocalExecutionRequestOutcome::Exhausted(exhaustion))
}
@@ -760,6 +783,7 @@ where
Attempt: AiExecutionAttempt + Send + Sync + 'static,
{
let mut last_attempted = None;
let mut fallback_response = None;
loop {
let next_started_at = std::time::Instant::now();
@@ -781,8 +805,8 @@ where
}
port.record_attempt_started(&attempt).await?;
let execute_started_at = std::time::Instant::now();
let response = match port.execute_attempt(&attempt).await {
Ok(response) => response,
let execution = match port.execute_attempt(&attempt).await {
Ok(execution) => execution,
Err(err) => {
let remaining = source.drain_execution_attempts().await?;
port.mark_unused_attempts(remaining).await?;
@@ -793,15 +817,26 @@ where
"stream_candidate_execute",
execute_started_at.elapsed().as_millis() as u64,
);
if let Some(response) = response {
let remaining = source.drain_execution_attempts().await?;
let unused_started_at = std::time::Instant::now();
port.mark_unused_attempts(remaining).await?;
observe_gateway_stage_ms(
"stream_candidate_unused",
unused_started_at.elapsed().as_millis() as u64,
);
return Ok(LocalExecutionRequestOutcome::responded(response));
match execution {
AiAttemptExecutionOutcome::Responded(response) => {
let remaining = source.drain_execution_attempts().await?;
let unused_started_at = std::time::Instant::now();
port.mark_unused_attempts(remaining).await?;
observe_gateway_stage_ms(
"stream_candidate_unused",
unused_started_at.elapsed().as_millis() as u64,
);
return Ok(LocalExecutionRequestOutcome::responded(response));
}
AiAttemptExecutionOutcome::Retry {
scope,
fallback_response: attempt_fallback_response,
} => {
if attempt_fallback_response.is_some() {
fallback_response = attempt_fallback_response;
}
apply_attempt_retry_scope(source, &attempt, scope).await?;
}
}
port.record_attempt_failed(&attempt).await?;
@@ -816,6 +851,12 @@ where
last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context()));
}
if let Some(response) = fallback_response {
return Ok(LocalExecutionRequestOutcome::responded(
mark_deferred_upstream_response(response),
));
}
let Some((last_plan, last_report_context)) = last_attempted else {
return Ok(LocalExecutionRequestOutcome::NoPath);
};
@@ -826,6 +867,24 @@ where
))
}
async fn apply_attempt_retry_scope<Source, Attempt>(
source: &mut Source,
attempt: &Attempt,
scope: AiAttemptRetryScope,
) -> Result<(), GatewayError>
where
Source: LocalExecutionAttemptSource<Attempt>,
Attempt: AiExecutionAttempt,
{
let plan = attempt.execution_plan();
match scope {
AiAttemptRetryScope::Candidate => Ok(()),
AiAttemptRetryScope::Credential => source.skip_credential(plan.key_id.as_str()).await,
AiAttemptRetryScope::Endpoint => source.skip_endpoint(plan.endpoint_id.as_str()).await,
AiAttemptRetryScope::Provider => source.skip_provider(plan.provider_id.as_str()).await,
}
}
async fn next_execution_attempt_with_timeout<Source, Attempt>(
source: &mut Source,
trace_id: &str,
@@ -901,7 +960,10 @@ where
Ok(())
}
async fn execute_attempt(&self, attempt: &T) -> Result<Option<Self::Response>, Self::Error> {
async fn execute_attempt(
&self,
attempt: &T,
) -> Result<AiAttemptExecutionOutcome<Self::Response>, Self::Error> {
let plan = attempt.execution_plan();
let report_context = attempt.report_context();
let candidate_index = parse_request_candidate_report_context(report_context.as_ref())
@@ -931,7 +993,7 @@ where
)
.await?
{
return Ok(Some(response));
return Ok(AiAttemptExecutionOutcome::Responded(response));
}
prewarm_direct_reqwest_candidate_client(plan);
// The attempt owns the canonical report context. Borrow it for the
@@ -951,14 +1013,14 @@ where
let execution_decision = self.decision.clone();
let execution_report_kind = attempt.report_kind();
let execution_plan = plan.clone();
let mut response = execute_stream_candidate_with_watchdog(
let mut execution = execute_stream_candidate_with_watchdog(
self.state,
self.trace_id,
self.plan_kind,
plan,
watchdog_report_context,
move || async move {
execute_execution_runtime_stream(
execute_execution_runtime_stream_with_retry_scope(
&execution_state,
execution_plan,
execution_trace_id.as_str(),
@@ -971,10 +1033,18 @@ where
},
)
.await?;
if let Some(response) = response.as_mut() {
attach_redaction_execution_candidate(response, plan.candidate_id.as_deref());
match &mut execution {
AiAttemptExecutionOutcome::Responded(response)
| AiAttemptExecutionOutcome::Retry {
fallback_response: Some(response),
..
} => attach_redaction_execution_candidate(response, plan.candidate_id.as_deref()),
AiAttemptExecutionOutcome::Retry {
fallback_response: None,
..
} => {}
}
Ok(response)
Ok(execution)
}
async fn mark_unused_attempts(&self, attempts: Vec<T>) -> Result<(), Self::Error> {
@@ -1234,9 +1304,11 @@ async fn execute_stream_candidate_with_watchdog<Fut>(
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
execute: impl FnOnce() -> Fut,
) -> Result<Option<Response<Body>>, GatewayError>
) -> Result<AiAttemptExecutionOutcome<Response<Body>>, GatewayError>
where
Fut: std::future::Future<Output = Result<Option<Response<Body>>, GatewayError>> + Send,
Fut: std::future::Future<
Output = Result<AiAttemptExecutionOutcome<Response<Body>>, GatewayError>,
> + Send,
{
let timeout_duration = resolve_stream_candidate_watchdog_timeout(plan, report_context);
let candidate_started_unix_ms = current_unix_ms();
@@ -1252,7 +1324,9 @@ where
)
.await;
log_stream_candidate_admission_timeout(trace_id, plan_kind, plan, report_context, &err);
return Ok(None);
return Ok(AiAttemptExecutionOutcome::retry(
AiAttemptRetryScope::Candidate,
));
}
Err(err) => return Err(err),
};
@@ -1300,7 +1374,9 @@ where
timeout_ms,
"gateway local stream candidate watchdog timed out"
);
Ok(None)
Ok(AiAttemptExecutionOutcome::retry(
AiAttemptRetryScope::Candidate,
))
}
};
observe_gateway_stage_ms(
@@ -1308,7 +1384,21 @@ where
watchdog_started_at.elapsed().as_millis() as u64,
);
match outcome {
Ok(response) => Ok(maybe_hold_upstream_execution_permit(response, permit_hold)),
Ok(AiAttemptExecutionOutcome::Responded(response)) => {
let response = maybe_hold_upstream_execution_permit(Some(response), permit_hold)
.expect("responded stream attempt must retain its response");
Ok(AiAttemptExecutionOutcome::Responded(response))
}
Ok(AiAttemptExecutionOutcome::Retry {
scope,
fallback_response,
}) => {
drop(permit_hold);
Ok(AiAttemptExecutionOutcome::Retry {
scope,
fallback_response,
})
}
Err(err) if is_candidate_level_admission_timeout(&err) => {
drop(permit_hold);
if should_record_candidate_admission_timeout(&err) {
@@ -1322,7 +1412,9 @@ where
.await;
}
log_stream_candidate_admission_timeout(trace_id, plan_kind, plan, report_context, &err);
Ok(None)
Ok(AiAttemptExecutionOutcome::retry(
AiAttemptRetryScope::Candidate,
))
}
Err(err) => {
drop(permit_hold);
@@ -1605,6 +1697,14 @@ mod tests {
Ok(Vec::new())
}
async fn skip_credential(&mut self, _key_id: &str) -> Result<(), GatewayError> {
Ok(())
}
async fn skip_endpoint(&mut self, _endpoint_id: &str) -> Result<(), GatewayError> {
Ok(())
}
async fn skip_provider(&mut self, _provider_id: &str) -> Result<(), GatewayError> {
Ok(())
}
@@ -1638,6 +1738,7 @@ mod tests {
struct TransferTestPort<'a> {
state: &'a AppState,
tracker: ProviderTransferTracker,
retry_scope: AiAttemptRetryScope,
executed: StdMutex<Vec<&'static str>>,
unused: StdMutex<Vec<&'static str>>,
}
@@ -1651,6 +1752,17 @@ mod tests {
Self {
state,
tracker,
retry_scope: AiAttemptRetryScope::Candidate,
executed: StdMutex::new(Vec::new()),
unused: StdMutex::new(Vec::new()),
}
}
fn with_retry_scope(state: &'a AppState, retry_scope: AiAttemptRetryScope) -> Self {
Self {
state,
tracker: ProviderTransferTracker::default(),
retry_scope,
executed: StdMutex::new(Vec::new()),
unused: StdMutex::new(Vec::new()),
}
@@ -1702,9 +1814,13 @@ mod tests {
async fn execute_attempt(
&self,
attempt: &TransferTestAttempt,
) -> Result<Option<Self::Response>, Self::Error> {
) -> Result<AiAttemptExecutionOutcome<Self::Response>, Self::Error> {
self.executed.lock().unwrap().push(attempt.label);
Ok((attempt.plan.provider_id == "provider-b").then(|| Response::new(Body::from("ok"))))
Ok(if attempt.plan.provider_id == "provider-b" {
AiAttemptExecutionOutcome::Responded(Response::new(Body::from("ok")))
} else {
AiAttemptExecutionOutcome::retry(self.retry_scope)
})
}
async fn mark_unused_attempts(
@@ -1751,6 +1867,18 @@ mod tests {
Ok(self.attempts.drain(..).collect())
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.attempts
.retain(|attempt| attempt.plan.key_id != key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.attempts
.retain(|attempt| attempt.plan.endpoint_id != endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.skipped_providers.push(provider_id.to_string());
self.attempts
@@ -1872,6 +2000,36 @@ mod tests {
assert_eq!(source.skipped_providers, ["provider-a"]);
}
#[tokio::test]
async fn dynamic_loop_applies_provider_scoped_retry_to_candidate_source() {
let state = AppState::new().expect("state should build");
let port = TransferTestPort::with_retry_scope(&state, AiAttemptRetryScope::Provider);
let mut source = TransferTestAttemptSource {
attempts: transfer_test_attempts().into(),
skipped_providers: Vec::new(),
};
let outcome = run_dynamic_attempt_loop(
&port,
&mut source,
"trace-provider-scope-test",
"provider_scope_test",
Duration::from_secs(1),
)
.await
.expect("dynamic attempt loop should succeed");
assert!(matches!(
outcome,
LocalExecutionRequestOutcome::Responded(_)
));
assert_eq!(
port.executed.lock().unwrap().as_slice(),
["a-key1-retry0", "b-key1-retry0"]
);
assert_eq!(source.skipped_providers, ["provider-a"]);
}
#[test]
fn transfer_timeout_is_checked_at_candidate_boundary_and_zero_disables_limits() {
let started_at = Instant::now();
@@ -2171,14 +2329,24 @@ mod tests {
"claude_cli_stream",
&plan,
Some(&report_context),
|| std::future::pending::<Result<Option<Response<Body>>, GatewayError>>(),
|| {
std::future::pending::<
Result<AiAttemptExecutionOutcome<Response<Body>>, GatewayError>,
>()
},
)
.await
});
tokio::time::sleep(Duration::from_millis(40)).await;
let result = task.await.expect("watchdog task should join");
assert!(matches!(result, Ok(None)));
assert!(matches!(
result,
Ok(AiAttemptExecutionOutcome::Retry {
scope: AiAttemptRetryScope::Candidate,
fallback_response: None,
})
));
let records = writer.records.lock().await;
assert_eq!(records.len(), 1);
@@ -2226,7 +2394,13 @@ mod tests {
)
.await;
assert!(matches!(result, Ok(None)));
assert!(matches!(
result,
Ok(AiAttemptExecutionOutcome::Retry {
scope: AiAttemptRetryScope::Candidate,
fallback_response: None,
})
));
let records = writer.records.lock().await;
assert_eq!(records.len(), 1);
let record = &records[0];
@@ -2268,7 +2442,13 @@ mod tests {
)
.await;
assert!(matches!(result, Ok(None)));
assert!(matches!(
result,
Ok(AiAttemptExecutionOutcome::Retry {
scope: AiAttemptRetryScope::Candidate,
fallback_response: None,
})
));
assert!(writer.records.lock().await.is_empty());
}
}
apps/aether-gateway/src/executor/mod.rs
+2 -1
View File
@@ -19,7 +19,8 @@ pub(crate) use orchestration::*;
pub(crate) use outcome::{
beautify_local_execution_client_error_message, build_fast_local_execution_exhaustion,
build_fast_local_execution_runtime_miss_context, build_local_execution_exhaustion,
build_local_execution_runtime_miss_context, record_failed_usage_for_exhausted_request,
build_local_execution_runtime_miss_context, is_deferred_upstream_response,
mark_deferred_upstream_response, record_failed_usage_for_exhausted_request,
record_failed_usage_for_runtime_miss_request, LocalExecutionExhaustion,
LocalExecutionRequestOutcome, LocalExecutionRuntimeMissContext,
};
apps/aether-gateway/src/executor/orchestration.rs
+13 -3
View File
@@ -1091,7 +1091,7 @@ fn standard_text_sync_heartbeat_error_kind(status_code: u16) -> LocalCoreSyncErr
401 => LocalCoreSyncErrorKind::Authentication,
403 => LocalCoreSyncErrorKind::PermissionDenied,
404 => LocalCoreSyncErrorKind::NotFound,
413 => LocalCoreSyncErrorKind::ContextLengthExceeded,
413 => LocalCoreSyncErrorKind::RequestTooLarge,
429 => LocalCoreSyncErrorKind::RateLimit,
503 => LocalCoreSyncErrorKind::Overloaded,
_ => LocalCoreSyncErrorKind::ServerError,
@@ -1586,6 +1586,18 @@ mod tests {
Ok(self.attempts.drain(..).collect())
}
async fn skip_credential(&mut self, key_id: &str) -> Result<(), GatewayError> {
self.attempts
.retain(|attempt| attempt.plan.key_id != key_id);
Ok(())
}
async fn skip_endpoint(&mut self, endpoint_id: &str) -> Result<(), GatewayError> {
self.attempts
.retain(|attempt| attempt.plan.endpoint_id != endpoint_id);
Ok(())
}
async fn skip_provider(&mut self, provider_id: &str) -> Result<(), GatewayError> {
self.attempts
.retain(|attempt| attempt.plan.provider_id != provider_id);
@@ -1813,7 +1825,6 @@ mod tests {
test_openai_image_heartbeat_attempt(0, "endpoint-retry", "candidate-retry"),
test_openai_image_heartbeat_attempt(1, "endpoint-success", "candidate-success"),
];
let outcome = execute_openai_image_sync_heartbeat_attempts(
state,
"/v1/images/generations".to_string(),
@@ -2104,7 +2115,6 @@ mod tests {
"openai:responses:compact",
),
];
let (parts, _) = http::Request::builder()
.method(http::Method::POST)
.uri("/v1/responses")
apps/aether-gateway/src/executor/outcome.rs
+59 -13
View File
@@ -18,6 +18,7 @@ use base64::Engine as _;
use serde_json::{json, Map, Value};
use tracing::warn;
use crate::ai_serving::{build_core_error_body_for_client_format, LocalCoreSyncErrorKind};
use crate::constants::{
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER,
};
@@ -32,6 +33,9 @@ pub(crate) enum LocalExecutionRequestOutcome {
NoPath,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct DeferredUpstreamResponse;
#[derive(Debug, Clone)]
pub(crate) struct LocalExecutionExhaustion {
request_id: String,
@@ -70,6 +74,18 @@ impl LocalExecutionRequestOutcome {
}
}
pub(crate) fn mark_deferred_upstream_response(mut response: Response<Body>) -> Response<Body> {
response.extensions_mut().insert(DeferredUpstreamResponse);
response
}
pub(crate) fn is_deferred_upstream_response(response: &Response<Body>) -> bool {
response
.extensions()
.get::<DeferredUpstreamResponse>()
.is_some()
}
impl LocalExecutionRuntimeMissContext {
pub(crate) fn persisted_candidate_count(&self) -> usize {
self.candidate_contexts.len()
@@ -316,12 +332,12 @@ pub(crate) async fn record_failed_usage_for_exhausted_request(
);
}
data.client_response_headers = Some(Value::Object(client_headers));
data.client_response_body = Some(json!({
"error": {
"type": "http_error",
"message": beautify_local_execution_client_error_message(local_execution_runtime_miss_detail),
}
}));
let client_message =
beautify_local_execution_client_error_message(local_execution_runtime_miss_detail);
data.client_response_body = Some(runtime_miss_client_error_body(
data.api_format.as_deref(),
&client_message,
));
let mut request_metadata = match data.request_metadata.take() {
Some(Value::Object(object)) => object,
@@ -392,12 +408,7 @@ pub(crate) async fn record_failed_usage_for_runtime_miss_request(
let status_code = http::StatusCode::SERVICE_UNAVAILABLE.as_u16();
let client_message =
beautify_local_execution_client_error_message(local_execution_runtime_miss_detail);
let client_body = json!({
"error": {
"type": "http_error",
"message": client_message,
}
});
let client_body = runtime_miss_client_error_body(api_format.as_deref(), &client_message);
let mut client_headers = Map::from_iter([(
"content-type".to_string(),
Value::String("application/json".to_string()),
@@ -654,6 +665,30 @@ fn json_header_map() -> Value {
)]))
}
fn runtime_miss_client_error_body(api_format: Option<&str>, message: &str) -> Value {
let fallback = json!({
"error": {
"type": "http_error",
"message": message,
}
});
let is_claude = api_format.is_some_and(|format| {
crate::ai_serving::normalize_api_format_alias(format)
.eq_ignore_ascii_case("claude:messages")
});
if !is_claude {
return fallback;
}
build_core_error_body_for_client_format(
"claude:messages",
message,
None,
LocalCoreSyncErrorKind::Overloaded,
)
.unwrap_or(fallback)
}
fn runtime_miss_original_headers_json(headers: &HeaderMap) -> Value {
let mut headers = crate::headers::collect_control_headers(headers);
for (name, value) in headers.iter_mut() {
@@ -1135,7 +1170,7 @@ fn trimmed_non_empty(value: Option<&str>) -> Option<String> {
mod tests {
use super::{
apply_runtime_miss_usage_routing, beautify_local_execution_client_error_message,
request_candidate_represents_provider_execution,
request_candidate_represents_provider_execution, runtime_miss_client_error_body,
select_last_runtime_miss_executed_candidate, LocalExecutionRuntimeMissContext,
RuntimeMissCandidateContext,
};
@@ -1169,6 +1204,17 @@ mod tests {
);
}
#[test]
fn runtime_miss_usage_body_matches_claude_client_envelope() {
let claude = runtime_miss_client_error_body(Some("claude:messages"), "busy");
assert_eq!(claude["type"], "error");
assert_eq!(claude["error"]["type"], "overloaded_error");
let openai = runtime_miss_client_error_body(Some("openai:chat"), "busy");
assert_eq!(openai["error"]["type"], "http_error");
assert!(openai.get("type").is_none());
}
#[test]
fn runtime_miss_routing_moves_to_typed_usage_fields_and_keeps_metadata_lightweight() {
let mut data = UsageEventData::default();
apps/aether-gateway/src/executor/stream_path.rs
+24 -2
View File
@@ -1,6 +1,6 @@
use aether_ai_serving::{
run_ai_stream_execution_path, AiPlanFallbackReason, AiServingExecutionOutcome,
AiStreamExecutionPathPort, AiStreamExecutionStep,
AiStreamExecutionPathPort, AiStreamExecutionStep, OriginalRequestPayload,
};
use async_trait::async_trait;
use axum::body::{Body, Bytes};
@@ -59,6 +59,21 @@ pub(crate) async fn maybe_execute_via_stream_decision_path(
);
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let mut planning_parts = parts.clone();
if crate::ai_serving::is_json_request(&planning_parts.headers) {
if let Ok(decoded_body) = crate::ai_serving::decoded_request_body_bytes(
&planning_parts.headers,
body_bytes.as_ref(),
) {
planning_parts
.extensions
.insert(OriginalRequestPayload::from_parsed_json(
body_json.clone(),
decoded_body.as_ref(),
));
}
}
let parts = &planning_parts;
observe_gateway_stage_ms(
"frontdoor_stream_parse",
parse_started_at.elapsed().as_millis() as u64,
@@ -422,7 +437,11 @@ fn to_ai_serving_outcome(
) -> AiServingExecutionOutcome<Response<Body>, super::LocalExecutionExhaustion> {
match outcome {
LocalExecutionRequestOutcome::Responded(response) => {
AiServingExecutionOutcome::Responded(response)
if super::is_deferred_upstream_response(&response) {
AiServingExecutionOutcome::Deferred(response)
} else {
AiServingExecutionOutcome::Responded(response)
}
}
LocalExecutionRequestOutcome::Exhausted(outcome) => {
AiServingExecutionOutcome::Exhausted(outcome)
@@ -438,6 +457,9 @@ fn from_ai_serving_outcome(
AiServingExecutionOutcome::Responded(response) => {
LocalExecutionRequestOutcome::Responded(response)
}
AiServingExecutionOutcome::Deferred(response) => {
LocalExecutionRequestOutcome::Responded(response)
}
AiServingExecutionOutcome::Exhausted(outcome) => {
LocalExecutionRequestOutcome::Exhausted(outcome)
}
apps/aether-gateway/src/executor/sync_path.rs
+25 -2
View File
@@ -1,6 +1,6 @@
use aether_ai_serving::{
run_ai_sync_execution_path, AiPlanFallbackReason, AiServingExecutionOutcome,
AiSyncExecutionPathPort, AiSyncExecutionStep,
AiSyncExecutionPathPort, AiSyncExecutionStep, OriginalRequestPayload,
};
use async_trait::async_trait;
use axum::body::{Body, Bytes};
@@ -51,6 +51,22 @@ pub(crate) async fn maybe_execute_via_sync_decision_path(
return Ok(LocalExecutionRequestOutcome::NoPath);
};
let mut planning_parts = parts.clone();
if crate::ai_serving::is_json_request(&planning_parts.headers) {
if let Ok(decoded_body) = crate::ai_serving::decoded_request_body_bytes(
&planning_parts.headers,
body_bytes.as_ref(),
) {
planning_parts
.extensions
.insert(OriginalRequestPayload::from_parsed_json(
body_json.clone(),
decoded_body.as_ref(),
));
}
}
let parts = &planning_parts;
if let Some(stream_plan_kind) = resolve_execution_runtime_stream_plan_kind(parts, decision) {
if is_matching_stream_request(stream_plan_kind, parts, &body_json, body_base64.as_deref()) {
return Ok(LocalExecutionRequestOutcome::NoPath);
@@ -256,7 +272,11 @@ fn to_ai_serving_outcome(
) -> AiServingExecutionOutcome<Response<Body>, super::LocalExecutionExhaustion> {
match outcome {
LocalExecutionRequestOutcome::Responded(response) => {
AiServingExecutionOutcome::Responded(response)
if super::is_deferred_upstream_response(&response) {
AiServingExecutionOutcome::Deferred(response)
} else {
AiServingExecutionOutcome::Responded(response)
}
}
LocalExecutionRequestOutcome::Exhausted(outcome) => {
AiServingExecutionOutcome::Exhausted(outcome)
@@ -272,6 +292,9 @@ fn from_ai_serving_outcome(
AiServingExecutionOutcome::Responded(response) => {
LocalExecutionRequestOutcome::Responded(response)
}
AiServingExecutionOutcome::Deferred(response) => {
LocalExecutionRequestOutcome::Responded(response)
}
AiServingExecutionOutcome::Exhausted(outcome) => {
LocalExecutionRequestOutcome::Exhausted(outcome)
}