i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-08 02:17:46 +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/tests/ai_execute/sync/chat/pii_redaction.rs
+21 -73
View File
@@ -782,82 +782,30 @@ async fn ai_execute_pii_redaction_restores_executed_candidate_session_after_late
}
large_stack_async_test!(
pii_redaction_performance_limits_do_not_forward_unredacted_body_upstream,
pii_redaction_performance_limits_do_not_forward_unredacted_body_upstream_impl
pii_redaction_forwards_text_above_previous_scan_cap_only_after_masking,
pii_redaction_forwards_text_above_previous_scan_cap_only_after_masking_impl
);
async fn pii_redaction_performance_limits_do_not_forward_unredacted_body_upstream_impl() {
let provider_hits = Arc::new(AtomicUsize::new(0));
let provider_hits_clone = Arc::clone(&provider_hits);
let provider_app = Router::new().route(
"/v1/chat/completions",
any(move |_request: Request| {
let provider_hits_inner = Arc::clone(&provider_hits_clone);
async move {
provider_hits_inner.fetch_add(1, Ordering::SeqCst);
Json(json!({
"id": "unexpected",
"choices": [{"message": {"role": "assistant", "content": "unexpected"}}]
}))
}
}),
async fn pii_redaction_forwards_text_above_previous_scan_cap_only_after_masking_impl() {
let mut request = rich_pii_request();
request["messages"]
.as_array_mut()
.expect("messages should be an array")
.push(json!({
"role": "user",
"content": "x".repeat(2 * 1024 * 1024 + 1),
}));
let (response_json, seen) =
run_sync_redaction_case("pii-redaction-large-request", true, true, "known", request).await;
let provider_body_text = serde_json::to_string(&seen.body).expect("body should serialize");
assert!(!provider_body_text.contains("[email protected]"));
assert!(provider_body_text.contains("<AETHER:EMAIL:"));
assert_eq!(
response_json["choices"][0]["message"]["content"],
"restored [email protected] access_token=accessValueABCDEF1234567890abcdef secret_key=secretValueABCDEF1234567890abcdef"
);
let (provider_url, provider_handle) = start_server(provider_app).await;
let auth_repository =
auth_repository_with_redaction_feature_settings("pii-redaction-limit", true);
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
candidate_row("pii-redaction-limit"),
]));
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider("pii-redaction-limit", true)],
vec![endpoint("pii-redaction-limit", provider_url)],
vec![key("pii-redaction-limit")],
));
let data_state = crate::data::GatewayDataState::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
DEVELOPMENT_ENCRYPTION_KEY,
)
.with_system_config_values_for_tests(redaction_config(true));
let gateway_state = AppState::new()
.expect("gateway state should build")
.with_data_state_for_tests(data_state);
let gateway = build_router_with_state(gateway_state);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let original = format!("[email protected] {}", "x".repeat(2 * 1024 * 1024));
let response = reqwest::Client::new()
.post(format!("{gateway_url}/v1/chat/completions"))
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-pii-redaction-limit",
)
.header(TRACE_ID_HEADER, "trace-pii-redaction-limit")
.body(
json!({
"model": "gpt-5",
"messages": [{"role": "user", "content": original}]
})
.to_string(),
)
.send()
.await
.expect("request should complete");
let status = response.status();
let response_text = response.text().await.expect("body should read");
assert_eq!(status, StatusCode::PAYLOAD_TOO_LARGE, "{response_text}");
assert!(response_text.contains("scanned text limit exceeded"));
assert!(!response_text.contains("[email protected]"));
assert_eq!(provider_hits.load(Ordering::SeqCst), 0);
gateway_handle.abort();
provider_handle.abort();
}
large_stack_async_test!(
apps/aether-gateway/src/tests/control/internal.rs
+40 -20
View File
@@ -444,15 +444,18 @@ async fn gateway_streams_tunnel_relay_body_to_attachment_owner() {
);
let (owner_url, owner_handle) = start_server(owner).await;
let data_state = GatewayDataState::disabled().with_system_config_values_for_tests(vec![(
"tunnel.attachments.node-123".to_string(),
json!({
"gateway_instance_id": "gateway-b",
"relay_base_url": owner_url,
"conn_count": 1,
"observed_at_unix_secs": 4_102_444_800u64,
}),
)]);
let data_state = GatewayDataState::disabled().with_system_config_values_for_tests(vec![
(
"tunnel.attachments.node-123".to_string(),
json!({
"gateway_instance_id": "gateway-b",
"relay_base_url": owner_url,
"conn_count": 1,
"observed_at_unix_secs": 4_102_444_800u64,
}),
),
("max_request_body_size".to_string(), json!(8)),
]);
let mut state = AppState::new()
.expect("gateway should build")
.with_data_state_for_tests(data_state)
@@ -539,16 +542,21 @@ async fn gateway_does_not_forward_tunnel_relay_twice() {
}
#[tokio::test]
async fn gateway_rejects_owner_relay_body_above_configured_limit() {
let owner_hits = Arc::new(Mutex::new(0usize));
let owner_hits_clone = Arc::clone(&owner_hits);
async fn gateway_forwards_owner_relay_body_above_recording_limit() {
const RECORDING_LIMIT_BYTES: usize = 8 * 1024 * 1024;
let captured_body = Arc::new(Mutex::new(None::<Bytes>));
let captured_body_clone = Arc::clone(&captured_body);
let owner = Router::new().route(
"/api/internal/tunnel/relay/node-123",
post(move |_request: Request| {
let owner_hits_inner = Arc::clone(&owner_hits_clone);
post(move |body: Body| {
let captured_body_inner = Arc::clone(&captured_body_clone);
async move {
*owner_hits_inner.lock().expect("mutex should lock") += 1;
(StatusCode::OK, Body::from("unexpected owner hit"))
let body = axum::body::to_bytes(body, usize::MAX)
.await
.expect("owner body should read");
*captured_body_inner.lock().expect("mutex should lock") = Some(body);
StatusCode::OK
}
}),
);
@@ -564,7 +572,10 @@ async fn gateway_rejects_owner_relay_body_above_configured_limit() {
"observed_at_unix_secs": 4_102_444_800u64,
}),
),
("max_request_body_size".to_string(), json!(8)),
(
"max_request_body_size".to_string(),
json!(RECORDING_LIMIT_BYTES),
),
]);
let gateway = build_router_with_state(
AppState::new()
@@ -574,15 +585,24 @@ async fn gateway_rejects_owner_relay_body_above_configured_limit() {
);
let (gateway_url, gateway_handle) = start_server(gateway).await;
let request_payload = vec![b'x'; RECORDING_LIMIT_BYTES + 1];
let envelope = relay_envelope(
&relay_request_meta(false, Some(60_000), None),
&request_payload,
);
assert!(envelope.len() > RECORDING_LIMIT_BYTES);
let response = reqwest::Client::new()
.post(format!("{gateway_url}/api/internal/tunnel/relay/node-123"))
.body("relay-envelope")
.body(envelope.clone())
.send()
.await
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE);
assert_eq!(*owner_hits.lock().expect("mutex should lock"), 0);
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
captured_body.lock().expect("mutex should lock").as_ref(),
Some(&Bytes::from(envelope))
);
gateway_handle.abort();
owner_handle.abort();
apps/aether-gateway/src/tests/usage/local.rs
+165 -14
View File
@@ -171,6 +171,7 @@ async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when
}
},
"telemetry": {
"ttfb_ms": 10,
"elapsed_ms": 25
}
}))
@@ -241,6 +242,19 @@ async fn gateway_handles_local_openai_chat_sync_report_with_local_reporting_when
assert_eq!(stored_usage.status, "completed");
assert_eq!(stored_usage.total_tokens, 5);
assert_eq!(stored_usage.response_time_ms, Some(25));
let end_to_end_time_ms = stored_usage
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("end_to_end_time_ms"))
.and_then(serde_json::Value::as_u64)
.expect("end-to-end latency should be persisted separately");
let end_to_end_first_byte_time_ms = stored_usage
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("end_to_end_first_byte_time_ms"))
.and_then(serde_json::Value::as_u64)
.expect("end-to-end first-byte latency should be persisted separately");
assert!(end_to_end_first_byte_time_ms <= end_to_end_time_ms);
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-local-report-sync-123")
@@ -422,14 +436,14 @@ async fn gateway_truncates_deep_request_echo_for_local_openai_chat_sync_usage_im
}
#[test]
fn gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_usage() {
fn gateway_ignores_legacy_max_request_body_size_for_local_openai_chat_sync_usage() {
run_async_test_on_large_stack(
"gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_usage",
gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_usage_impl(),
"gateway_ignores_legacy_max_request_body_size_for_local_openai_chat_sync_usage",
gateway_ignores_legacy_max_request_body_size_for_local_openai_chat_sync_usage_impl(),
);
}
async fn gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_usage_impl() {
async fn gateway_ignores_legacy_max_request_body_size_for_local_openai_chat_sync_usage_impl() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
@@ -546,13 +560,13 @@ async fn gateway_applies_system_max_request_body_size_to_local_openai_chat_sync_
assert_eq!(stored_usage.total_tokens, 5);
assert_eq!(
stored_usage.request_body_state,
Some(UsageBodyCaptureState::Truncated)
Some(UsageBodyCaptureState::Inline)
);
assert_eq!(
stored_usage.provider_request_body_state,
Some(UsageBodyCaptureState::Truncated)
Some(UsageBodyCaptureState::Inline)
);
assert_eq!(
assert_ne!(
stored_usage
.request_body
.as_ref()
@@ -905,6 +919,7 @@ async fn gateway_records_failed_usage_when_sync_runtime_transport_is_unavailable
.expect("gateway should build")
.with_execution_runtime_sync_override_for_tests(move |_plan| {
*execution_hits_clone.lock().expect("mutex should lock") += 1;
std::thread::sleep(std::time::Duration::from_millis(5));
Err(crate::GatewayError::Internal(
"simulated transport unavailable".to_string(),
))
@@ -968,12 +983,148 @@ async fn gateway_records_failed_usage_when_sync_runtime_transport_is_unavailable
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
assert!(stored_candidates[0]
.latency_ms
.is_some_and(|value| value >= 5));
assert_eq!(
stored_candidates[0].error_type.as_deref(),
Some("execution_runtime_unavailable")
);
}
#[test]
fn sync_transport_error_policy_stops_or_retries_candidates_end_to_end() {
run_async_test_on_large_stack(
"sync_transport_error_policy_stops_or_retries_candidates_end_to_end",
sync_transport_error_policy_stops_or_retries_candidates_end_to_end_impl(),
);
}
async fn sync_transport_error_policy_stops_or_retries_candidates_end_to_end_impl() {
async fn run_case(stop_on_transport_errors: bool) -> (StatusCode, usize, Vec<Option<u16>>) {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let execution_hits = Arc::new(std::sync::atomic::AtomicUsize::new(0));
let execution_hits_for_override = Arc::clone(&execution_hits);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-transport-policy")),
sample_local_openai_auth_snapshot("api-key-transport-policy", "user-transport-policy"),
)]));
let mut second_candidate = sample_local_openai_candidate_row();
second_candidate.key_id = "key-openai-usage-local-2".to_string();
second_candidate.key_name = "secondary".to_string();
second_candidate.key_internal_priority = second_candidate.key_internal_priority - 1;
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(),
second_candidate,
]));
let mut provider = sample_local_openai_provider();
provider.config = stop_on_transport_errors.then(|| {
json!({
"failover_rules": {
"stop_on_transport_errors": true,
}
})
});
let mut second_key = sample_local_openai_key();
second_key.id = "key-openai-usage-local-2".to_string();
second_key.name = "secondary".to_string();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key(), second_key],
));
let gateway_state = crate::AppState::new()
.expect("gateway should build")
.with_execution_runtime_sync_override_for_tests(move |plan| {
let hit = execution_hits_for_override
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
if hit == 0 {
return Err(crate::GatewayError::Internal(
"simulated transport unavailable".to_string(),
));
}
Ok(aether_contracts::ExecutionResult {
request_id: plan.request_id.clone(),
candidate_id: plan.candidate_id.clone(),
status_code: 200,
headers: std::collections::BTreeMap::from([(
"content-type".to_string(),
"application/json".to_string(),
)]),
body: Some(aether_contracts::ResponseBody {
json_body: Some(json!({
"id": "chatcmpl-transport-policy",
"choices": [{"message": {"role": "assistant", "content": "ok"}}]
})),
body_bytes_b64: None,
}),
telemetry: Some(aether_contracts::ExecutionTelemetry {
ttfb_ms: Some(1),
elapsed_ms: Some(1),
upstream_bytes: None,
}),
error: None,
})
})
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
usage_repository,
DEVELOPMENT_ENCRYPTION_KEY,
),
);
let gateway = build_router_with_state(gateway_state);
let trace_id = if stop_on_transport_errors {
"trace-transport-policy-stop"
} else {
"trace-transport-policy-retry"
};
let request = Request::builder()
.method(http::Method::POST)
.uri("/v1/chat/completions")
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-transport-policy",
)
.header(TRACE_ID_HEADER, trace_id)
.body(Body::from("{\"model\":\"gpt-5\",\"messages\":[]}"))
.expect("request should build");
let response = send_request(gateway, request).await;
let candidates = request_candidate_repository
.list_by_request_id(trace_id)
.await
.expect("request candidates should read");
(
response.status(),
execution_hits.load(std::sync::atomic::Ordering::SeqCst),
candidates
.into_iter()
.filter(|candidate| candidate.status == RequestCandidateStatus::Failed)
.map(|candidate| candidate.status_code)
.collect(),
)
}
let (retry_status, retry_hits, retry_failure_statuses) = run_case(false).await;
assert_eq!(retry_status, StatusCode::OK);
assert!(retry_hits >= 2);
assert_eq!(retry_failure_statuses.first(), Some(&None));
let (stop_status, stop_hits, stop_failure_statuses) = run_case(true).await;
assert_eq!(stop_status, StatusCode::BAD_GATEWAY);
assert_eq!(stop_hits, 1);
assert_eq!(stop_failure_statuses, vec![None]);
}
#[test]
fn gateway_records_failed_usage_for_claude_runtime_miss_without_execution_exhaustion() {
run_async_test_on_large_stack(
@@ -1365,14 +1516,14 @@ async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_wh
}
#[test]
fn gateway_preserves_stream_usage_when_max_response_body_size_truncates_capture() {
fn gateway_ignores_legacy_max_response_body_size_for_stream_usage() {
run_async_test_on_large_stack(
"gateway_preserves_stream_usage_when_max_response_body_size_truncates_capture",
gateway_preserves_stream_usage_when_max_response_body_size_truncates_capture_impl(),
"gateway_ignores_legacy_max_response_body_size_for_stream_usage",
gateway_ignores_legacy_max_response_body_size_for_stream_usage_impl(),
);
}
async fn gateway_preserves_stream_usage_when_max_response_body_size_truncates_capture_impl() {
async fn gateway_ignores_legacy_max_response_body_size_for_stream_usage_impl() {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
@@ -1524,13 +1675,13 @@ async fn gateway_preserves_stream_usage_when_max_response_body_size_truncates_ca
assert_eq!(stored_usage.total_tokens, 6);
assert_eq!(
stored_usage.response_body_state,
Some(UsageBodyCaptureState::Truncated)
Some(UsageBodyCaptureState::Inline)
);
assert_eq!(
stored_usage.client_response_body_state,
Some(UsageBodyCaptureState::Truncated)
Some(UsageBodyCaptureState::Inline)
);
assert_eq!(
assert_ne!(
stored_usage
.response_body
.as_ref()