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

feat(gateway): harden failover and payload handling

Browse Source 
Retry pre-response transport failures across candidates with an explicit stop policy, and propagate end-to-end timing into usage records and UI diagnostics.

Remove legacy body, import, cookie, PII, and tunnel replay caps while preserving optional operator-configured gateway limits.
This commit is contained in:
elky
2026-07-30 01:03:27 +08:00
parent a97acc07fc
commit a04673a90d
80 changed files with 3640 additions and 1236 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/executor/candidate_loop.rs
+220 -40
View File
@@ -20,8 +20,9 @@ use crate::ai_serving::LocalExecutionAttemptSource;
use crate::clock::current_unix_ms;
use crate::control::GatewayControlDecision;
use crate::execution_runtime::{
execute_execution_runtime_stream_with_retry_scope,
build_transport_error_stop_response, execute_execution_runtime_stream_with_retry_scope,
execute_execution_runtime_sync_with_retry_scope,
mark_stream_candidate_watchdog_terminal_started, StreamCandidateWatchdogProgress,
};
use crate::executor::{
build_local_execution_exhaustion, mark_deferred_upstream_response, LocalExecutionRequestOutcome,
@@ -30,7 +31,9 @@ use crate::handlers::shared::provider_pool::release_admin_provider_pool_key_leas
use crate::log_ids::short_request_id;
use crate::orchestration::{
local_execution_candidate_metadata_from_report_context,
local_failover_policy_from_report_context, resolve_local_failover_policy, LocalFailoverPolicy,
local_failover_policy_from_report_context, resolve_local_failover_policy,
resolve_local_transport_failover_analysis_for_attempt, LocalFailoverDecision,
LocalFailoverPolicy,
};
use crate::privacy::RedactionExecutionCandidateId;
use crate::request_candidate_runtime::{
@@ -1013,12 +1016,24 @@ where
let execution_decision = self.decision.clone();
let execution_report_kind = attempt.report_kind();
let execution_plan = plan.clone();
let mut execution = execute_stream_candidate_with_watchdog(
let stop_on_transport_errors = matches!(
resolve_local_transport_failover_analysis_for_attempt(
self.state,
plan,
watchdog_report_context,
)
.await
.decision,
LocalFailoverDecision::StopLocalFailover
);
let watchdog_started_at = std::time::Instant::now();
let execution = execute_stream_candidate_with_watchdog(
self.state,
self.trace_id,
self.plan_kind,
plan,
watchdog_report_context,
stop_on_transport_errors,
move || async move {
execute_execution_runtime_stream_with_retry_scope(
&execution_state,
@@ -1033,6 +1048,25 @@ where
},
)
.await?;
let mut execution = match execution {
StreamCandidateWatchdogOutcome::TransportTimeout => {
AiAttemptExecutionOutcome::Responded(
build_transport_error_stop_response(
self.state,
plan,
watchdog_report_context,
self.trace_id,
self.decision,
http::StatusCode::GATEWAY_TIMEOUT.as_u16(),
"local_stream_candidate_watchdog_timeout",
stream_candidate_watchdog_timeout_message(),
watchdog_started_at.elapsed().as_millis() as u64,
)
.await?,
)
}
StreamCandidateWatchdogOutcome::Executed(execution) => execution,
};
match &mut execution {
AiAttemptExecutionOutcome::Responded(response)
| AiAttemptExecutionOutcome::Retry {
@@ -1297,20 +1331,28 @@ fn log_stream_candidate_admission_timeout(
);
}
#[derive(Debug)]
enum StreamCandidateWatchdogOutcome {
Executed(AiAttemptExecutionOutcome<Response<Body>>),
TransportTimeout,
}
async fn execute_stream_candidate_with_watchdog<Fut>(
state: &(impl RequestCandidateRuntimeWriter + UpstreamExecutionGateProvider + ?Sized),
trace_id: &str,
plan_kind: &str,
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&serde_json::Value>,
stop_on_transport_errors: bool,
execute: impl FnOnce() -> Fut,
) -> Result<AiAttemptExecutionOutcome<Response<Body>>, GatewayError>
) -> Result<StreamCandidateWatchdogOutcome, GatewayError>
where
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_at = std::time::Instant::now();
let candidate_started_unix_ms = current_unix_ms();
let permit = match acquire_upstream_execution_gate(state, trace_id).await {
Ok(permit) => permit,
@@ -1324,17 +1366,33 @@ where
)
.await;
log_stream_candidate_admission_timeout(trace_id, plan_kind, plan, report_context, &err);
return Ok(AiAttemptExecutionOutcome::retry(
AiAttemptRetryScope::Candidate,
return Ok(StreamCandidateWatchdogOutcome::Executed(
AiAttemptExecutionOutcome::retry(AiAttemptRetryScope::Candidate),
));
}
Err(err) => return Err(err),
};
let permit_hold = permit.map(UpstreamExecutionPermitHold::new);
let watchdog_started_at = std::time::Instant::now();
let outcome = match timeout(timeout_duration, execute()).await {
Ok(result) => result,
Err(_) => {
let watchdog_progress = StreamCandidateWatchdogProgress::shared();
let execution = watchdog_progress.clone().scope(execute());
tokio::pin!(execution);
let deadline = tokio::time::sleep(timeout_duration);
tokio::pin!(deadline);
let execution_result = tokio::select! {
biased;
result = &mut execution => Some(result),
() = &mut deadline => {
if watchdog_progress.terminal_started() {
Some(execution.await)
} else {
None
}
}
};
let outcome = match execution_result {
Some(result) => result.map(StreamCandidateWatchdogOutcome::Executed),
None => {
let finished_at_unix_ms = current_unix_ms();
let request_id = short_request_id(plan.request_id.as_str());
let provider_name = plan.provider_name.as_deref().unwrap_or("-");
@@ -1350,10 +1408,10 @@ where
report_context,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(http::StatusCode::GATEWAY_TIMEOUT.as_u16()),
status_code: None,
error_type: Some("local_stream_candidate_watchdog_timeout".to_string()),
error_message: Some(stream_candidate_watchdog_timeout_message().to_string()),
latency_ms: None,
latency_ms: Some(candidate_started_at.elapsed().as_millis() as u64),
started_at_unix_ms: Some(candidate_started_unix_ms),
finished_at_unix_ms: Some(finished_at_unix_ms),
},
@@ -1374,9 +1432,13 @@ where
timeout_ms,
"gateway local stream candidate watchdog timed out"
);
Ok(AiAttemptExecutionOutcome::retry(
AiAttemptRetryScope::Candidate,
))
if stop_on_transport_errors {
Ok(StreamCandidateWatchdogOutcome::TransportTimeout)
} else {
Ok(StreamCandidateWatchdogOutcome::Executed(
AiAttemptExecutionOutcome::retry(AiAttemptRetryScope::Candidate),
))
}
}
};
observe_gateway_stage_ms(
@@ -1384,20 +1446,30 @@ where
watchdog_started_at.elapsed().as_millis() as u64,
);
match outcome {
Ok(AiAttemptExecutionOutcome::Responded(response)) => {
Ok(StreamCandidateWatchdogOutcome::Executed(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(StreamCandidateWatchdogOutcome::Executed(
AiAttemptExecutionOutcome::Responded(response),
))
}
Ok(AiAttemptExecutionOutcome::Retry {
Ok(StreamCandidateWatchdogOutcome::Executed(AiAttemptExecutionOutcome::Retry {
scope,
fallback_response,
}) => {
})) => {
drop(permit_hold);
Ok(AiAttemptExecutionOutcome::Retry {
scope,
fallback_response,
})
Ok(StreamCandidateWatchdogOutcome::Executed(
AiAttemptExecutionOutcome::Retry {
scope,
fallback_response,
},
))
}
Ok(StreamCandidateWatchdogOutcome::TransportTimeout) => {
drop(permit_hold);
Ok(StreamCandidateWatchdogOutcome::TransportTimeout)
}
Err(err) if is_candidate_level_admission_timeout(&err) => {
drop(permit_hold);
@@ -1412,8 +1484,8 @@ where
.await;
}
log_stream_candidate_admission_timeout(trace_id, plan_kind, plan, report_context, &err);
Ok(AiAttemptExecutionOutcome::retry(
AiAttemptRetryScope::Candidate,
Ok(StreamCandidateWatchdogOutcome::Executed(
AiAttemptExecutionOutcome::retry(AiAttemptRetryScope::Candidate),
))
}
Err(err) => {
@@ -2329,6 +2401,7 @@ mod tests {
"claude_cli_stream",
&plan,
Some(&report_context),
false,
|| {
std::future::pending::<
Result<AiAttemptExecutionOutcome<Response<Body>>, GatewayError>,
@@ -2342,20 +2415,19 @@ mod tests {
let result = task.await.expect("watchdog task should join");
assert!(matches!(
result,
Ok(AiAttemptExecutionOutcome::Retry {
scope: AiAttemptRetryScope::Candidate,
fallback_response: None,
})
Ok(StreamCandidateWatchdogOutcome::Executed(
AiAttemptExecutionOutcome::Retry {
scope: AiAttemptRetryScope::Candidate,
fallback_response: None,
}
))
));
let records = writer.records.lock().await;
assert_eq!(records.len(), 1);
let record = &records[0];
assert_eq!(record.status, RequestCandidateStatus::Failed);
assert_eq!(
record.status_code,
Some(http::StatusCode::GATEWAY_TIMEOUT.as_u16())
);
assert_eq!(record.status_code, None);
assert_eq!(
record.error_type.as_deref(),
Some("local_stream_candidate_watchdog_timeout")
@@ -2367,6 +2439,108 @@ mod tests {
assert_eq!(record.candidate_index, 2);
}
#[tokio::test]
async fn stream_candidate_watchdog_can_stop_on_transport_error() {
let writer = Arc::new(TestRequestCandidateWriter::default());
let plan = test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(5),
..ExecutionTimeouts::default()
}));
let report_context = test_report_context();
let result = execute_stream_candidate_with_watchdog(
writer.as_ref(),
"trace_watchdog_stop",
"claude_cli_stream",
&plan,
Some(&report_context),
true,
|| {
std::future::pending::<
Result<AiAttemptExecutionOutcome<Response<Body>>, GatewayError>,
>()
},
)
.await;
assert!(matches!(
result,
Ok(StreamCandidateWatchdogOutcome::TransportTimeout)
));
let records = writer.records.lock().await;
assert_eq!(records.len(), 1);
assert_eq!(records[0].status_code, None);
assert_eq!(
records[0].error_type.as_deref(),
Some("local_stream_candidate_watchdog_timeout")
);
}
#[tokio::test]
async fn stream_candidate_watchdog_does_not_cancel_started_terminalization() {
let writer = Arc::new(TestRequestCandidateWriter::default());
let plan = test_plan(Some(ExecutionTimeouts {
first_byte_ms: Some(5),
..ExecutionTimeouts::default()
}));
let report_context = test_report_context();
let result = execute_stream_candidate_with_watchdog(
writer.as_ref(),
"trace_terminalization",
"claude_cli_stream",
&plan,
Some(&report_context),
true,
|| async {
mark_stream_candidate_watchdog_terminal_started();
tokio::time::sleep(Duration::from_millis(20)).await;
Ok(AiAttemptExecutionOutcome::Responded(Response::new(
Body::from("terminal response"),
)))
},
)
.await;
assert!(matches!(
result,
Ok(StreamCandidateWatchdogOutcome::Executed(
AiAttemptExecutionOutcome::Responded(_)
))
));
assert!(writer.records.lock().await.is_empty());
}
#[tokio::test]
async fn stream_candidate_watchdog_does_not_relabel_execution_error_as_timeout() {
let writer = Arc::new(TestRequestCandidateWriter::default());
let plan = test_plan(None);
let report_context = test_report_context();
let result = execute_stream_candidate_with_watchdog(
writer.as_ref(),
"trace_execution_error",
"claude_cli_stream",
&plan,
Some(&report_context),
true,
|| async {
Err(GatewayError::UpstreamUnavailable {
trace_id: "trace_execution_error".to_string(),
message: "upstream connect failed".to_string(),
})
},
)
.await;
assert!(matches!(
result,
Err(GatewayError::UpstreamUnavailable { message, .. })
if message == "upstream connect failed"
));
assert!(writer.records.lock().await.is_empty());
}
#[tokio::test]
async fn stream_candidate_upstream_execution_admission_timeout_marks_failed_and_continues() {
let writer = Arc::new(TestRequestCandidateWriter::with_upstream_gate(
@@ -2388,6 +2562,7 @@ mod tests {
"claude_cli_stream",
&plan,
Some(&report_context),
false,
|| async {
panic!("execute future should not run while upstream execution gate is saturated")
},
@@ -2396,10 +2571,12 @@ mod tests {
assert!(matches!(
result,
Ok(AiAttemptExecutionOutcome::Retry {
scope: AiAttemptRetryScope::Candidate,
fallback_response: None,
})
Ok(StreamCandidateWatchdogOutcome::Executed(
AiAttemptExecutionOutcome::Retry {
scope: AiAttemptRetryScope::Candidate,
fallback_response: None,
}
))
));
let records = writer.records.lock().await;
assert_eq!(records.len(), 1);
@@ -2432,6 +2609,7 @@ mod tests {
"claude_cli_stream",
&plan,
Some(&report_context),
false,
|| async {
Err(GatewayError::AdmissionTimeout {
trace_id: "trace_target_admission".to_string(),
@@ -2444,10 +2622,12 @@ mod tests {
assert!(matches!(
result,
Ok(AiAttemptExecutionOutcome::Retry {
scope: AiAttemptRetryScope::Candidate,
fallback_response: None,
})
Ok(StreamCandidateWatchdogOutcome::Executed(
AiAttemptExecutionOutcome::Retry {
scope: AiAttemptRetryScope::Candidate,
fallback_response: None,
}
))
));
assert!(writer.records.lock().await.is_empty());
}
apps/aether-gateway/src/executor/orchestration.rs
+185 -19
View File
@@ -55,6 +55,7 @@ use crate::executor::{
LocalExecutionRequestOutcome,
};
use crate::handlers::shared::system_config_bool;
use crate::request_diagnostics::{current_request_diagnostics, scope_request_diagnostics_with};
use crate::stage_metrics::observe_gateway_stage_ms;
use crate::{AiExecutionDecision, AppState, GatewayError};
@@ -844,15 +845,19 @@ where
let decision_for_response = decision.clone();
let started_at = Instant::now();
let (tx, rx) = mpsc::channel::<Result<Bytes, IoError>>(1);
let request_diagnostics = current_request_diagnostics();
tokio::spawn(async move {
let bytes = standard_text_sync_heartbeat_final_bytes(
client_api_format.as_str(),
redaction_slot.as_ref(),
execute(state, parts, trace_id, decision, plan_kind, started_at).await,
)
scope_request_diagnostics_with(request_diagnostics, async move {
let bytes = standard_text_sync_heartbeat_final_bytes(
client_api_format.as_str(),
redaction_slot.as_ref(),
execute(state, parts, trace_id, decision, plan_kind, started_at).await,
)
.await;
let _ = tx.send(Ok(Bytes::from(bytes))).await;
})
.await;
let _ = tx.send(Ok(Bytes::from(bytes))).await;
});
let headers = BTreeMap::from([(
@@ -1115,23 +1120,27 @@ fn build_openai_image_sync_heartbeat_shell_response(
let decision_for_response = decision.clone();
let started_at = Instant::now();
let (tx, rx) = mpsc::channel::<Result<Bytes, IoError>>(1);
let request_diagnostics = current_request_diagnostics();
tokio::spawn(async move {
let bytes = openai_image_sync_heartbeat_final_bytes(
execute_openai_image_sync_heartbeat_attempts(
state,
request_path,
trace_id,
decision,
plan_kind,
attempts,
transfer_tracker,
started_at,
scope_request_diagnostics_with(request_diagnostics, async move {
let bytes = openai_image_sync_heartbeat_final_bytes(
execute_openai_image_sync_heartbeat_attempts(
state,
request_path,
trace_id,
decision,
plan_kind,
attempts,
transfer_tracker,
started_at,
)
.await,
)
.await,
)
.await;
let _ = tx.send(Ok(Bytes::from(bytes))).await;
})
.await;
let _ = tx.send(Ok(Bytes::from(bytes))).await;
});
let headers = BTreeMap::from([(
@@ -1557,9 +1566,14 @@ pub(crate) fn decision_payload_is_direct_execution(payload: &AiExecutionDecision
#[cfg(test)]
mod tests {
use super::*;
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::repository::usage::InMemoryUsageReadRepository;
use aether_data_contracts::repository::usage::UsageReadRepository;
use aether_usage_runtime::UsageRuntimeConfig;
use futures_util::StreamExt;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
const TEST_OPENAI_IMAGE_SYNC_PLAN_KIND: &str = "openai_image_sync";
const TEST_STANDARD_TEXT_SYNC_PLAN_KIND: &str = "openai_responses_compact_sync";
@@ -1684,6 +1698,61 @@ mod tests {
}
}
fn heartbeat_usage_test_state(
response_body: Value,
) -> (AppState, Arc<InMemoryUsageReadRepository>) {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
.expect("state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
request_candidate_repository,
Arc::clone(&usage_repository),
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
})
.with_execution_runtime_sync_override_for_tests(move |plan| {
let mut result = test_openai_image_execution_result(
plan,
StatusCode::OK.as_u16(),
response_body.clone(),
);
if let Some(telemetry) = result.telemetry.as_mut() {
telemetry.ttfb_ms = Some(5);
}
Ok(result)
});
(state, usage_repository)
}
async fn assert_usage_has_end_to_end_timings(
usage_repository: &InMemoryUsageReadRepository,
request_id: &str,
) {
let usage = usage_repository
.find_by_request_id(request_id)
.await
.expect("usage should read")
.expect("terminal usage should be recorded");
let request_metadata = usage
.request_metadata
.as_ref()
.expect("terminal usage should retain request diagnostics");
let end_to_end_time_ms = request_metadata
.get("end_to_end_time_ms")
.and_then(Value::as_u64)
.expect("end-to-end time should be recorded");
let end_to_end_first_byte_time_ms = request_metadata
.get("end_to_end_first_byte_time_ms")
.and_then(Value::as_u64)
.expect("end-to-end first-byte time should be recorded");
assert!(end_to_end_first_byte_time_ms <= end_to_end_time_ms);
}
fn test_standard_text_heartbeat_decision() -> GatewayControlDecision {
GatewayControlDecision::synthetic(
"/v1/responses",
@@ -1799,6 +1868,46 @@ mod tests {
assert_eq!(body["error"]["upstream_status"], json!(503));
}
#[tokio::test]
async fn openai_image_sync_heartbeat_propagates_request_diagnostics_to_terminal_usage() {
let (state, usage_repository) = heartbeat_usage_test_state(json!({
"data": [{"b64_json": "heartbeat-image"}]
}));
let response = crate::request_diagnostics::scope_request_diagnostics(async move {
crate::request_diagnostics::record_request_accepted_at(
Instant::now() - Duration::from_millis(25),
);
build_openai_image_sync_heartbeat_shell_response(
state,
"/v1/images/generations".to_string(),
"trace-image-heartbeat-retry".to_string(),
test_openai_image_heartbeat_decision(),
TEST_OPENAI_IMAGE_SYNC_PLAN_KIND.to_string(),
vec![test_openai_image_heartbeat_attempt(
0,
"endpoint-success",
"candidate-success",
)],
ProviderTransferTracker::default(),
)
})
.await
.expect("heartbeat shell should build");
let body = to_bytes(
response.into_body(),
crate::headers::max_internal_buffered_body_bytes(),
)
.await
.expect("heartbeat response body should complete");
assert!(!body.is_empty());
assert_usage_has_end_to_end_timings(
usage_repository.as_ref(),
"trace-image-heartbeat-retry",
)
.await;
}
#[tokio::test]
async fn openai_image_sync_heartbeat_attempts_retry_first_candidate_then_return_second() {
let call_count = Arc::new(AtomicUsize::new(0));
@@ -2065,6 +2174,63 @@ mod tests {
let _ = release_tx.send(());
}
#[tokio::test]
async fn standard_text_sync_heartbeat_propagates_request_diagnostics_to_terminal_usage() {
let (state, usage_repository) = heartbeat_usage_test_state(json!({
"id": "resp_heartbeat",
"output": []
}));
let (parts, _) = http::Request::builder()
.method(http::Method::POST)
.uri("/v1/responses")
.body(())
.expect("request should build")
.into_parts();
let response = crate::request_diagnostics::scope_request_diagnostics(async move {
crate::request_diagnostics::record_request_accepted_at(
Instant::now() - Duration::from_millis(25),
);
build_standard_text_sync_heartbeat_shell_response(
state,
parts,
"trace-standard-text-heartbeat-retry".to_string(),
test_standard_text_heartbeat_decision(),
TEST_STANDARD_TEXT_SYNC_PLAN_KIND.to_string(),
move |state, parts, trace_id, decision, plan_kind, _started_at| async move {
execute_sync_attempt_source::<AiSyncAttempt, _>(
&state,
&parts,
trace_id.as_str(),
&decision,
plan_kind.as_str(),
TestSyncAttemptSource::new(vec![test_standard_text_heartbeat_attempt(
0,
"endpoint-success",
"candidate-success",
"openai:responses:compact",
)]),
)
.await
},
)
})
.await
.expect("heartbeat shell should build");
let body = to_bytes(
response.into_body(),
crate::headers::max_internal_buffered_body_bytes(),
)
.await
.expect("heartbeat response body should complete");
assert!(!body.is_empty());
assert_usage_has_end_to_end_timings(
usage_repository.as_ref(),
"trace-standard-text-heartbeat-retry",
)
.await;
}
#[test]
fn standard_text_sync_heartbeat_compact_non_json_error_body_is_wrapped_in_client_format() {
let bytes = standard_text_sync_heartbeat_error_body_from_response(