feat(gateway): harden failover and payload handling

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
@@ -115,7 +115,16 @@ pub(crate) async fn maybe_execute_chatgpt_web_image_sync(
let result = match execute_chatgpt_web_image(state, plan, report_context, started_at).await
{
Ok(result) => result,
Err(err) => chatgpt_web_transport_error_execution_result(plan, started_at, &err),
Err(ExecutionRuntimeTransportError::UpstreamHttpStatus {
status_code,
message,
}) => chatgpt_web_http_error_execution_result(
plan,
started_at,
status_code,
message.as_str(),
),
Err(error) => return Err(error),
};
Ok(Some(result))
})
@@ -133,7 +142,13 @@ pub(crate) async fn maybe_execute_chatgpt_web_image_stream(
let started_at = Instant::now();
let result = match execute_chatgpt_web_image(state, plan, report_context, started_at).await {
Ok(result) => result,
Err(err) => chatgpt_web_transport_error_execution_result(plan, started_at, &err),
Err(ExecutionRuntimeTransportError::UpstreamHttpStatus {
status_code,
message,
}) => {
chatgpt_web_http_error_execution_result(plan, started_at, status_code, message.as_str())
}
Err(error) => return Err(error),
};
Ok(Some(ChatGptWebImageStream {
frame_stream: execution_result_frame_stream(plan, &result, report_context),
@@ -2567,19 +2582,20 @@ fn json_execution_result(
}
}
fn chatgpt_web_transport_error_execution_result(
fn chatgpt_web_http_error_execution_result(
plan: &ExecutionPlan,
started_at: Instant,
error: &ExecutionRuntimeTransportError,
status_code: u16,
message: &str,
) -> ExecutionResult {
json_execution_result(
plan,
503,
status_code,
json!({
"error": {
"type": "upstream_error",
"code": "chatgpt_web_image_execution_unavailable",
"message": error.to_string()
"message": message
}
}),
started_at,
@@ -2798,11 +2814,14 @@ fn ensure_success(
return Ok(());
}
let body = String::from_utf8_lossy(&execution_result_body_bytes_lossy(result)).to_string();
Err(ExecutionRuntimeTransportError::UpstreamRequest(format!(
"{stage} returned {}: {}",
result.status_code,
body.chars().take(320).collect::<String>()
)))
Err(ExecutionRuntimeTransportError::UpstreamHttpStatus {
status_code: result.status_code,
message: format!(
"{stage} returned {}: {}",
result.status_code,
body.chars().take(320).collect::<String>()
),
})
}
fn chatgpt_web_base_url_from_plan(plan: &ExecutionPlan) -> String {
@@ -3992,10 +4011,14 @@ data: [DONE]
Some(&json!({"chatgpt_web_image": true})),
)
.await
.expect("executor should run")
.expect("executor should preserve the upstream HTTP response")
.expect("plan should be intercepted");
assert_ne!(result.status_code, 200);
assert_eq!(result.status_code, 500);
assert_eq!(
execution_result_json(&result).expect("error response should be json")["error"]["code"],
json!("chatgpt_web_image_execution_unavailable")
);
let metadata = reloaded_chatgpt_web_metadata(repository.as_ref()).await;
assert_eq!(metadata["image_quota_remaining"], json!(25.0));
assert_eq!(metadata["image_quota_used"], json!(0.0));
@@ -4005,6 +4028,71 @@ data: [DONE]
handle.abort();
}
#[tokio::test]
async fn chatgpt_web_image_sync_propagates_network_failure_without_synthetic_503() {
let listener = crate::test_support::bind_loopback_listener()
.await
.expect("listener should bind");
let base_url = format!(
"http://{}",
listener.local_addr().expect("local addr should resolve")
);
drop(listener);
let state = crate::AppState::new().expect("state should build");
let plan = sample_plan(
base_url.as_str(),
json!({"prompt": "draw a small test image"}),
false,
);
let error = maybe_execute_chatgpt_web_image_sync(
&state,
&plan,
Some(&json!({"chatgpt_web_image": true})),
)
.await
.expect_err("connection failure should propagate to the candidate loop");
assert!(matches!(
error,
ExecutionRuntimeTransportError::UpstreamRequest(_)
));
}
#[tokio::test]
async fn chatgpt_web_image_stream_propagates_network_failure_without_synthetic_503() {
let listener = crate::test_support::bind_loopback_listener()
.await
.expect("listener should bind");
let base_url = format!(
"http://{}",
listener.local_addr().expect("local addr should resolve")
);
drop(listener);
let state = crate::AppState::new().expect("state should build");
let plan = sample_plan(
base_url.as_str(),
json!({"prompt": "draw a small test image"}),
true,
);
let error = match maybe_execute_chatgpt_web_image_stream(
&state,
&plan,
Some(&json!({"chatgpt_web_image": true})),
)
.await
{
Err(error) => error,
Ok(_) => panic!("connection failure should propagate to the candidate loop"),
};
assert!(matches!(
error,
ExecutionRuntimeTransportError::UpstreamRequest(_)
));
}
#[tokio::test]
async fn chatgpt_web_image_stream_path_wraps_success_sse_as_ndjson_frames() {
let (base_url, handle) = start_mock_chatgpt_web().await;
@@ -940,6 +940,7 @@ mod tests {
max_transfer_timeout_seconds: 0,
stop_status_codes: [503].into_iter().collect(),
continue_status_codes: [409, 429].into_iter().collect(),
stop_on_transport_errors: false,
success_failover_patterns: Vec::new(),
error_stop_patterns: Vec::new(),
stop_cyber_policy_errors: true,
@@ -43,7 +43,6 @@ const GROK_MEDIA_POST_PATH: &str = "/rest/media/post/create";
const GROK_IMAGINE_WS_URL: &str = "wss://grok.com/ws/imagine/listen";
const GROK_STANDARD_PROVIDER_API_FORMAT: &str = "openai:responses";
const GROK_PROMPT_OVERHEAD_TOKENS: u64 = 4;
const GROK_MAX_ATTACHMENT_BYTES: usize = 25 * 1024 * 1024;
const GROK_MAX_ATTACHMENT_REDIRECTS: usize = 5;
const GROK_IMAGINE_STREAM_TIMEOUT_MS: u64 = 10_000;
const GROK_IMAGINE_ROUND_TIMEOUT_MS: u64 = 120_000;
@@ -642,7 +641,7 @@ fn grok_success_frame_stream(
let chunk = match item {
Ok(chunk) => chunk,
Err(message) => {
match encode_grok_error_frame(status_code, message) {
match encode_grok_error_frame(message) {
Ok(frame) => yield Ok(frame),
Err(err) => {
yield Err(err);
@@ -872,17 +871,17 @@ fn encode_grok_telemetry_frame(
})
}
fn encode_grok_error_frame(status_code: u16, message: String) -> Result<Bytes, IoError> {
fn encode_grok_error_frame(message: String) -> Result<Bytes, IoError> {
encode_stream_frame_ndjson(&StreamFrame {
frame_type: StreamFrameType::Error,
payload: StreamFramePayload::Error {
error: aether_contracts::ExecutionError {
kind: aether_contracts::ExecutionErrorKind::Internal,
kind: aether_contracts::ExecutionErrorKind::ProtocolError,
phase: aether_contracts::ExecutionPhase::StreamRead,
message,
upstream_status: Some(status_code),
retryable: false,
failover_recommended: false,
upstream_status: None,
retryable: true,
failover_recommended: true,
},
},
})
@@ -896,7 +895,7 @@ fn encode_grok_first_byte_timeout_frame(timeout: Duration) -> Result<Bytes, IoEr
kind: aether_contracts::ExecutionErrorKind::FirstByteTimeout,
phase: aether_contracts::ExecutionPhase::FirstByte,
message: stream_first_byte_timeout_message(timeout),
upstream_status: Some(504),
upstream_status: None,
retryable: true,
failover_recommended: true,
},
@@ -1422,20 +1421,15 @@ fn grok_attachment_payload_from_data_uri(
})
.unwrap_or_else(|| "application/octet-stream".to_string());
let normalized_b64 = content_b64.split_whitespace().collect::<String>();
let decoded_len = base64::engine::general_purpose::STANDARD
.decode(&normalized_b64)
.map_err(|err| {
ExecutionRuntimeTransportError::UpstreamRequest(format!(
"Grok attachment data URI base64 is invalid: {err}"
))
})?
.len();
if decoded_len > GROK_MAX_ATTACHMENT_BYTES {
return Err(ExecutionRuntimeTransportError::UpstreamRequest(format!(
"Grok attachment exceeds {} byte limit",
GROK_MAX_ATTACHMENT_BYTES
)));
}
drop(
base64::engine::general_purpose::STANDARD
.decode(&normalized_b64)
.map_err(|err| {
ExecutionRuntimeTransportError::UpstreamRequest(format!(
"Grok attachment data URI base64 is invalid: {err}"
))
})?,
);
Ok(GrokAttachmentPayload {
filename: input
.filename
@@ -1635,12 +1629,6 @@ async fn collect_grok_attachment_url_bytes(
let chunk = chunk.map_err(|err| {
ExecutionRuntimeTransportError::UpstreamRequest(format_upstream_request_error(&err))
})?;
if bytes.len().saturating_add(chunk.len()) > GROK_MAX_ATTACHMENT_BYTES {
return Err(ExecutionRuntimeTransportError::UpstreamRequest(format!(
"Grok attachment exceeds {} byte limit",
GROK_MAX_ATTACHMENT_BYTES
)));
}
bytes.extend_from_slice(&chunk);
}
Ok(bytes)
@@ -3073,12 +3061,6 @@ async fn grok_download_image_asset(
if bytes.is_empty() {
return Ok(None);
}
if bytes.len() > GROK_MAX_ATTACHMENT_BYTES {
return Err(ExecutionRuntimeTransportError::UpstreamRequest(format!(
"Grok image asset exceeds {} byte limit",
GROK_MAX_ATTACHMENT_BYTES
)));
}
Ok(Some(format!(
"data:{content_type};base64,{}",
base64::engine::general_purpose::STANDARD.encode(bytes)
@@ -3236,7 +3218,10 @@ fn push_sse_event(body: &mut String, event: &str, data: &Value) {
mod tests {
use std::collections::BTreeMap;
use aether_contracts::{ExecutionPlan, RequestBody, StreamFrame, StreamFramePayload};
use aether_contracts::{
ExecutionErrorKind, ExecutionPhase, ExecutionPlan, RequestBody, StreamFrame,
StreamFramePayload,
};
use axum::body::{Body, Bytes};
use axum::extract::Request;
use axum::routing::any;
@@ -3246,16 +3231,18 @@ mod tests {
use http::{Method, StatusCode};
use super::{
encode_grok_error_frame, encode_grok_first_byte_timeout_frame,
extract_grok_attachment_inputs, grok_aspect_ratio_from_provider_body, grok_asset_url,
grok_attachment_ip_is_public, grok_client_json_body, grok_client_stream_body,
grok_handle_imagine_ws_message, grok_image_count_from_provider_body,
grok_image_prompt_from_provider_body, grok_imagine_request_message,
grok_imagine_reset_message, grok_media_post_url,
grok_attachment_ip_is_public, grok_attachment_payload_from_data_uri, grok_client_json_body,
grok_client_stream_body, grok_handle_imagine_ws_message,
grok_image_count_from_provider_body, grok_image_prompt_from_provider_body,
grok_imagine_request_message, grok_imagine_reset_message, grok_media_post_url,
grok_plan_uses_structured_image_generation, grok_should_collect_image_stream,
grok_should_use_imagine_websocket, grok_success_frame_stream, grok_upload_url,
grok_upstream_model_name, grok_usage_estimate, grok_user_id_from_cookie_header,
materialize_grok_image_assets, openai_chat_body, openai_image_body, openai_responses_body,
set_grok_image_edit_config, GrokCollected, GrokImagineImage, GrokStreamAdapter,
set_grok_image_edit_config, GrokAttachmentInput, GrokCollected, GrokImagineImage,
GrokStreamAdapter,
};
fn sample_plan(body: serde_json::Value, client_api_format: &str) -> ExecutionPlan {
@@ -3326,6 +3313,45 @@ mod tests {
out
}
fn decode_encoded_frame(encoded: Bytes) -> StreamFrame {
let line = String::from_utf8(encoded.to_vec()).expect("frame should be utf8");
serde_json::from_str(line.trim()).expect("frame should deserialize")
}
#[test]
fn grok_stream_read_error_is_retryable_transport_without_upstream_status() {
let frame = decode_encoded_frame(
encode_grok_error_frame("connection reset while reading response body".to_string())
.expect("error frame should encode"),
);
let StreamFramePayload::Error { error } = frame.payload else {
panic!("encoded frame should contain an execution error");
};
assert_eq!(error.kind, ExecutionErrorKind::ProtocolError);
assert_eq!(error.phase, ExecutionPhase::StreamRead);
assert_eq!(error.upstream_status, None);
assert!(error.retryable);
assert!(error.failover_recommended);
}
#[test]
fn grok_first_byte_timeout_is_retryable_transport_without_upstream_status() {
let frame = decode_encoded_frame(
encode_grok_first_byte_timeout_frame(std::time::Duration::from_millis(250))
.expect("timeout frame should encode"),
);
let StreamFramePayload::Error { error } = frame.payload else {
panic!("encoded frame should contain an execution error");
};
assert_eq!(error.kind, ExecutionErrorKind::FirstByteTimeout);
assert_eq!(error.phase, ExecutionPhase::FirstByte);
assert_eq!(error.upstream_status, None);
assert!(error.retryable);
assert!(error.failover_recommended);
}
#[tokio::test]
async fn grok_success_stream_forwards_token_chunks_incrementally() {
let plan = sample_plan(
@@ -4062,6 +4088,26 @@ mod tests {
assert_eq!(inputs[1].source.as_str(), "data:text/plain;base64,bm90ZXM=");
}
#[test]
fn grok_data_uri_attachment_accepts_content_above_previous_size_cap() {
const PREVIOUS_CAP_BYTES: usize = 25 * 1024 * 1024;
let base64_blocks = PREVIOUS_CAP_BYTES / 3 + 1;
let mut source = String::from("data:application/octet-stream;base64,");
source.extend(std::iter::repeat_n('A', base64_blocks * 4));
let input = GrokAttachmentInput {
source,
filename: Some("large.bin".to_string()),
mime_type: None,
};
let payload = grok_attachment_payload_from_data_uri(&input, 0)
.expect("attachment above the previous size cap should be accepted");
assert_eq!(payload.filename, "large.bin");
assert_eq!(payload.mime_type, "application/octet-stream");
assert_eq!(payload.content_b64.len(), base64_blocks * 4);
}
#[test]
fn extracts_responses_and_claude_attachment_inputs() {
let responses = extract_grok_attachment_inputs(
@@ -20,6 +20,7 @@ mod stream_pump;
pub(crate) mod submission;
pub(crate) mod sync;
pub(crate) mod transport;
mod transport_failure;
mod windsurf;
pub(crate) use self::chatgpt_web_image::maybe_execute_chatgpt_web_image_sync;
@@ -133,6 +134,10 @@ pub(crate) use transport::{
execute_sync_plan as execute_execution_runtime_sync_plan, DirectSyncExecutionRuntime,
DirectUpstreamStreamExecution, ExecutionRuntimeTransportError,
};
pub(crate) use transport_failure::{
build_transport_error_stop_response, mark_stream_candidate_watchdog_terminal_started,
StreamCandidateWatchdogProgress,
};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct ClientIntent {
@@ -370,6 +370,9 @@ impl IntoResponse for ExecutionRuntimeAppError {
| ExecutionRuntimeTransportError::UnsupportedTransportProfile(_)
| ExecutionRuntimeTransportError::BodyEncode(_),
) => StatusCode::BAD_REQUEST,
ExecutionRuntimeServerError::Transport(
ExecutionRuntimeTransportError::UpstreamHttpStatus { status_code, .. },
) => StatusCode::from_u16(status_code).unwrap_or(StatusCode::BAD_GATEWAY),
ExecutionRuntimeServerError::Transport(
ExecutionRuntimeTransportError::ClientBuild(_)
| ExecutionRuntimeTransportError::BrowserClientBuild(_)
@@ -24,7 +24,7 @@ use aether_scheduler_core::{
use aether_usage_runtime::{
build_lifecycle_usage_seed, build_stream_terminal_usage_payload_seed,
build_sync_terminal_usage_payload_seed, build_terminal_usage_context_seed, LifecycleUsageSeed,
SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed,
SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed, UsageRequestRecordLevel,
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
};
use async_stream::stream;
@@ -56,8 +56,9 @@ use super::error::{
mod execution_failures;
use self::execution_failures::{
build_stream_failure_from_execution_error, build_stream_failure_from_provider_error_body,
build_stream_failure_report, handle_prefetch_provider_private_stream_error,
handle_prefetch_stream_failure, submit_midstream_stream_failure, StreamFailureReport,
build_stream_failure_report, build_stream_transport_failure_report,
handle_prefetch_provider_private_stream_error, handle_prefetch_stream_failure,
submit_midstream_stream_failure, StreamFailureReport,
};
use crate::ai_serving::api::{
extract_provider_private_stream_error_body, maybe_bridge_standard_sync_json_to_stream,
@@ -129,7 +130,8 @@ use crate::request_candidate_runtime::{
};
use crate::request_diagnostics::{
attach_current_request_diagnostics_to_report_context,
attach_request_diagnostics_to_report_context, current_request_diagnostics, RequestDiagnostics,
attach_request_diagnostics_and_candidate_start_timing_to_report_context,
current_request_diagnostics, RequestDiagnostics,
};
use crate::stage_metrics::{
attach_stage_trace_to_report_context, observe_gateway_stage_ms, observe_gateway_stage_trace_ms,
@@ -148,6 +150,7 @@ const SSE_CONTROL_FILTER_MAX_BUFFER_BYTES: usize = 1024 * 1024;
const SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES: usize = 1024 * 1024;
const SSE_TERMINAL_DETECTOR_MAX_RECORD_BYTES: usize = SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES;
const PROVIDER_STREAM_ERROR_INSPECTION_MAX_BYTES: usize = SSE_TERMINAL_DETECTOR_MAX_LINE_BYTES;
const BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES: usize = 5 * 1024 * 1024;
const STREAM_IDLE_LOG_INTERVAL: Duration = Duration::from_secs(60);
const STREAM_IDLE_LOG_INTERVAL_MS: u64 = 60_000;
const REWRITTEN_STREAM_PREFETCH_TIMEOUT: Duration = Duration::from_millis(750);
@@ -281,8 +284,15 @@ fn report_context_with_stage_trace(
fn report_context_with_request_diagnostics(
report_context: Option<Value>,
diagnostics: Option<&Arc<RequestDiagnostics>>,
candidate_started_at: Instant,
terminal_telemetry: Option<&ExecutionTelemetry>,
) -> Option<Value> {
attach_request_diagnostics_to_report_context(report_context, diagnostics)
attach_request_diagnostics_and_candidate_start_timing_to_report_context(
report_context,
diagnostics,
Some(candidate_started_at),
terminal_telemetry.and_then(|telemetry| telemetry.ttfb_ms),
)
}
fn request_accepted_elapsed_ms(diagnostics: Option<&Arc<RequestDiagnostics>>) -> Option<u64> {
@@ -332,6 +342,37 @@ fn direct_passthrough_mode() -> DirectPassthroughMode {
.unwrap_or(DirectPassthroughMode::Inline)
}
fn stream_body_buffer_limit_for_record_level(record_level: UsageRequestRecordLevel) -> usize {
match record_level {
UsageRequestRecordLevel::Basic => BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES,
UsageRequestRecordLevel::Full => DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
}
}
async fn resolve_stream_body_buffer_limit(state: &AppState) -> usize {
if !state.usage_runtime.is_enabled() {
return BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES;
}
match state
.usage_runtime
.body_capture_policy_for(state.usage_lifecycle_data_state().as_ref())
.await
{
Ok(policy) => stream_body_buffer_limit_for_record_level(policy.record_level),
Err(error) => {
warn!(
event_name = "stream_body_capture_policy_read_failed",
log_type = "ops",
error = %error,
fallback = "full",
"gateway could not resolve stream body capture policy"
);
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES
}
}
}
fn parse_direct_passthrough_mode(value: &str) -> DirectPassthroughMode {
match value.trim().to_ascii_lowercase().as_str() {
"legacy" | "pump" | "mpsc" => DirectPassthroughMode::Legacy,
@@ -379,6 +420,7 @@ async fn record_sync_terminal_usage_with_handoff_after_spawn<F>(
) where
F: Future<Output = ()> + Send + 'static,
{
crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started();
// Capture request task-local diagnostics before handing the work to a spawned task. Tokio
// task-local values do not propagate across spawn boundaries.
let (context_seed, payload_seed) =
@@ -481,6 +523,7 @@ async fn record_stream_terminal_usage(
payload: &GatewayStreamReportRequest,
cancelled: bool,
) {
crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started();
let context_seed = build_terminal_usage_context_seed(plan, report_context);
let payload_seed = build_stream_terminal_usage_payload_seed(payload);
state
@@ -1719,6 +1762,7 @@ struct DirectPassthroughFinalizerCore {
stream_usage_observer: Option<StreamingStandardTerminalObserver>,
stream_usage_observer_buffered: Vec<u8>,
provider_error_inspection: ProviderStreamErrorInspection,
max_stream_body_buffer_bytes: usize,
provider_buffered_body: Vec<u8>,
buffered_body: Vec<u8>,
provider_body_truncated: bool,
@@ -1869,7 +1913,7 @@ impl DirectPassthroughFinalizer {
append_stream_capture_bytes(
&mut core.provider_buffered_body,
chunk.as_ref(),
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
core.max_stream_body_buffer_bytes,
&mut core.provider_body_truncated,
);
if let (Some(observer), Some(report_context)) = (
@@ -1907,7 +1951,7 @@ impl DirectPassthroughFinalizer {
append_stream_capture_bytes(
&mut core.buffered_body,
chunk.as_ref(),
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
core.max_stream_body_buffer_bytes,
&mut core.client_body_truncated,
);
if !core.requires_anthropic_message_stop {
@@ -2094,6 +2138,7 @@ impl DirectPassthroughFinalizerCore {
stream_usage_observer: _,
stream_usage_observer_buffered: _,
provider_error_inspection: _,
max_stream_body_buffer_bytes: _,
provider_buffered_body,
buffered_body,
provider_body_truncated,
@@ -2150,6 +2195,8 @@ impl DirectPassthroughFinalizerCore {
let report_context_for_payload = report_context_with_request_diagnostics(
report_context_for_payload,
request_diagnostics.as_ref(),
stream_started_at,
terminal_telemetry.as_ref(),
);
submit_midstream_stream_failure(
&state,
@@ -2183,6 +2230,8 @@ impl DirectPassthroughFinalizerCore {
let report_context_for_payload = report_context_with_request_diagnostics(
report_context_for_payload,
request_diagnostics.as_ref(),
stream_started_at,
terminal_telemetry.as_ref(),
);
let usage_payload = build_stream_usage_payload(
trace_id,
@@ -2275,6 +2324,8 @@ impl DirectPassthroughFinalizerCore {
let report_context_for_payload = report_context_with_request_diagnostics(
report_context_for_payload,
request_diagnostics.as_ref(),
stream_started_at,
terminal_telemetry.as_ref(),
);
let usage_payload = build_stream_usage_payload(
trace_id.clone(),
@@ -2582,7 +2633,7 @@ impl DirectPassthroughInlineBodyState {
Ok(item) => item,
Err(timeout) => {
if let Some(finalizer) = self.finalizer.as_mut() {
finalizer.set_terminal_failure(build_stream_failure_report(
finalizer.set_terminal_failure(build_stream_transport_failure_report(
"first_byte_timeout",
stream_first_byte_timeout_message(timeout),
504,
@@ -2637,7 +2688,7 @@ impl DirectPassthroughInlineBodyState {
error = %message,
"gateway direct passthrough upstream body read failed"
);
finalizer.set_terminal_failure(build_stream_failure_report(
finalizer.set_terminal_failure(build_stream_transport_failure_report(
"execution_runtime_stream_read_error",
message,
502,
@@ -2794,6 +2845,7 @@ async fn execute_stream_from_direct_passthrough(
}
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
let max_stream_body_buffer_bytes = resolve_stream_body_buffer_limit(state).await;
let request_candidate_status_snapshot =
snapshot_local_request_candidate_status(&plan, report_context.as_ref());
let passthrough_mode = direct_passthrough_mode();
@@ -2889,6 +2941,7 @@ async fn execute_stream_from_direct_passthrough(
stream_usage_observer,
stream_usage_observer_buffered: Vec::new(),
provider_error_inspection: ProviderStreamErrorInspection::default(),
max_stream_body_buffer_bytes,
provider_buffered_body: Vec::new(),
buffered_body: Vec::new(),
provider_body_truncated: false,
@@ -2958,7 +3011,6 @@ async fn execute_stream_from_direct_passthrough(
StageElapsedGuard::from_started_at("stream_total", stream_started_at_for_report);
let _provider_pool_in_flight_guard = provider_pool_in_flight_guard_for_report;
let _upstream_target_permit = upstream_target_permit;
let max_stream_body_buffer_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES;
let stream_usage_report_context =
normalized_stream_report_context_owned.clone().or_else(|| {
Some(serde_json::json!({
@@ -3014,7 +3066,7 @@ async fn execute_stream_from_direct_passthrough(
match result {
Ok(item) => item,
Err(timeout) => {
terminal_failure = Some(build_stream_failure_report(
terminal_failure = Some(build_stream_transport_failure_report(
"first_byte_timeout",
stream_first_byte_timeout_message(timeout),
504,
@@ -3063,7 +3115,7 @@ async fn execute_stream_from_direct_passthrough(
error = %message,
"gateway direct passthrough upstream body read failed"
);
terminal_failure = Some(build_stream_failure_report(
terminal_failure = Some(build_stream_transport_failure_report(
"execution_runtime_stream_read_error",
message,
502,
@@ -3308,6 +3360,8 @@ async fn execute_stream_from_direct_passthrough(
let report_context_for_payload = report_context_with_request_diagnostics(
report_context_for_payload,
request_diagnostics_for_report.as_ref(),
stream_started_at_for_report,
terminal_telemetry.as_ref(),
);
let usage_payload = build_stream_usage_payload(
trace_id_owned,
@@ -3368,6 +3422,8 @@ async fn execute_stream_from_direct_passthrough(
let report_context_for_payload = report_context_with_request_diagnostics(
report_context_for_payload,
request_diagnostics_for_report.as_ref(),
stream_started_at_for_report,
terminal_telemetry.as_ref(),
);
submit_midstream_stream_failure(
&state_for_report,
@@ -3433,6 +3489,8 @@ async fn execute_stream_from_direct_passthrough(
let report_context_for_payload = report_context_with_request_diagnostics(
report_context_for_payload,
request_diagnostics_for_report.as_ref(),
stream_started_at_for_report,
terminal_telemetry.as_ref(),
);
let usage_payload = build_stream_usage_payload(
trace_id_owned.clone(),
@@ -3627,6 +3685,41 @@ pub(crate) fn execute_execution_runtime_stream_with_retry_scope<'a>(
})
}
async fn maybe_build_stream_transport_error_stop_response(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
trace_id: &str,
decision: &GatewayControlDecision,
error_type: &str,
error_message: &str,
elapsed_ms: u64,
) -> Result<Option<Response<Body>>, GatewayError> {
let analysis = crate::orchestration::resolve_local_transport_failover_analysis_for_attempt(
state,
plan,
report_context,
)
.await;
if !matches!(analysis.decision, LocalFailoverDecision::StopLocalFailover) {
return Ok(None);
}
crate::execution_runtime::build_transport_error_stop_response(
state,
plan,
report_context,
trace_id,
decision,
http::StatusCode::BAD_GATEWAY.as_u16(),
error_type,
error_message,
elapsed_ms,
)
.await
.map(Some)
}
async fn execute_execution_runtime_stream_inner(
state: &AppState,
mut plan: ExecutionPlan,
@@ -3731,6 +3824,7 @@ async fn execute_execution_runtime_stream_inner(
}
Ok(None) => {}
Err(err) => {
let transport_error_message = err.to_string();
info!(
event_name = "grok_execution_unavailable",
log_type = "ops",
@@ -3754,13 +3848,27 @@ async fn execute_execution_runtime_stream_inner(
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("grok_execution_unavailable".to_string()),
error_message: Some(format!("{err:?}")),
latency_ms: None,
error_message: Some(transport_error_message.clone()),
latency_ms: Some(stream_elapsed_ms_since(stream_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_stream_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
"grok_execution_unavailable",
transport_error_message.as_str(),
stream_elapsed_ms_since(stream_started_at),
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
}
@@ -3788,6 +3896,7 @@ async fn execute_execution_runtime_stream_inner(
}
Ok(None) => {}
Err(err) => {
let transport_error_message = err.to_string();
info!(
event_name = "windsurf_native_execution_unavailable",
log_type = "ops",
@@ -3811,13 +3920,27 @@ async fn execute_execution_runtime_stream_inner(
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("windsurf_native_execution_unavailable".to_string()),
error_message: Some(err.to_string()),
latency_ms: None,
error_message: Some(transport_error_message.clone()),
latency_ms: Some(stream_elapsed_ms_since(stream_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_stream_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
"windsurf_native_execution_unavailable",
transport_error_message.as_str(),
stream_elapsed_ms_since(stream_started_at),
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
}
@@ -3845,6 +3968,7 @@ async fn execute_execution_runtime_stream_inner(
}
Ok(None) => {}
Err(err) => {
let transport_error_message = err.to_string();
info!(
event_name = "kiro_web_search_mcp_unavailable",
log_type = "ops",
@@ -3868,13 +3992,27 @@ async fn execute_execution_runtime_stream_inner(
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("kiro_web_search_mcp_unavailable".to_string()),
error_message: Some(format!("{err:?}")),
latency_ms: None,
error_message: Some(transport_error_message.clone()),
latency_ms: Some(stream_elapsed_ms_since(stream_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_stream_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
"kiro_web_search_mcp_unavailable",
transport_error_message.as_str(),
stream_elapsed_ms_since(stream_started_at),
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
}
@@ -3902,6 +4040,7 @@ async fn execute_execution_runtime_stream_inner(
}
Ok(None) => {}
Err(err) => {
let transport_error_message = err.to_string();
info!(
event_name = "chatgpt_web_image_execution_unavailable",
log_type = "ops",
@@ -3925,13 +4064,27 @@ async fn execute_execution_runtime_stream_inner(
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("chatgpt_web_image_execution_unavailable".to_string()),
error_message: Some(format!("{err:?}")),
latency_ms: None,
error_message: Some(transport_error_message.clone()),
latency_ms: Some(stream_elapsed_ms_since(stream_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_stream_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
"chatgpt_web_image_execution_unavailable",
transport_error_message.as_str(),
stream_elapsed_ms_since(stream_started_at),
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
}
@@ -3961,6 +4114,7 @@ async fn execute_execution_runtime_stream_inner(
return Err(err);
}
Err(InProcessStreamExecutionError::Transport(err)) => {
let transport_error_message = err.to_string();
info!(
event_name = "stream_execution_runtime_unavailable",
log_type = "ops",
@@ -3984,13 +4138,27 @@ async fn execute_execution_runtime_stream_inner(
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("execution_runtime_unavailable".to_string()),
error_message: Some(format!("{err:?}")),
latency_ms: None,
error_message: Some(transport_error_message.clone()),
latency_ms: Some(stream_elapsed_ms_since(stream_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_stream_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
"execution_runtime_unavailable",
transport_error_message.as_str(),
stream_elapsed_ms_since(stream_started_at),
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
};
@@ -4076,6 +4244,7 @@ async fn execute_execution_runtime_stream_inner(
return Err(err);
}
Err(InProcessStreamExecutionError::Transport(err)) => {
let transport_error_message = err.to_string();
info!(
event_name = "stream_execution_runtime_unavailable",
log_type = "ops",
@@ -4099,13 +4268,27 @@ async fn execute_execution_runtime_stream_inner(
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("execution_runtime_unavailable".to_string()),
error_message: Some(err.to_string()),
latency_ms: None,
error_message: Some(transport_error_message.clone()),
latency_ms: Some(stream_elapsed_ms_since(stream_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_stream_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
"execution_runtime_unavailable",
transport_error_message.as_str(),
stream_elapsed_ms_since(stream_started_at),
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
};
@@ -4177,6 +4360,7 @@ async fn execute_execution_runtime_stream_inner(
{
Ok(response) => response,
Err(err) => {
let transport_error_message = format!("{err:?}");
warn!(
event_name = "stream_execution_runtime_remote_unavailable",
log_type = "ops",
@@ -4195,13 +4379,27 @@ async fn execute_execution_runtime_stream_inner(
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("execution_runtime_unavailable".to_string()),
error_message: Some(format!("{err:?}")),
latency_ms: None,
error_message: Some(transport_error_message.clone()),
latency_ms: Some(stream_elapsed_ms_since(stream_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_stream_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
"execution_runtime_unavailable",
transport_error_message.as_str(),
stream_elapsed_ms_since(stream_started_at),
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
};
@@ -5315,6 +5513,7 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
if !lifecycle_pending_recorded {
record_stream_pending_lifecycle(state, &lifecycle_seed, &mut stage_trace).await;
}
let max_stream_body_buffer_bytes = resolve_stream_body_buffer_limit(state).await;
let request_candidate_status_snapshot =
snapshot_local_request_candidate_status(&plan, report_context.as_ref());
let candidate_index = parse_request_candidate_report_context(report_context.as_ref())
@@ -5912,7 +6111,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
headers,
prefetched_usage_telemetry.clone(),
&provider_prefetched_body,
candidate_started_unix_secs,
stream_elapsed_ms_since(stream_started_at),
failure,
None,
)
.await;
}
@@ -5984,7 +6186,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
headers,
prefetched_usage_telemetry.clone(),
&prefetched_body,
candidate_started_unix_secs,
stream_elapsed_ms_since(stream_started_at),
failure,
None,
)
.await;
}
@@ -6110,10 +6315,17 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
provider_prefetched_body_bytes = provider_prefetched_body.len(),
"gateway detected embedded error while prefetching execution runtime stream"
);
let request_diagnostics = current_request_diagnostics();
let terminal_report_context = report_context_with_request_diagnostics(
report_context,
request_diagnostics.as_ref(),
stream_started_at,
prefetched_usage_telemetry.as_ref(),
);
let payload = build_stream_sync_payload(
trace_id,
report_kind.clone(),
report_context,
terminal_report_context,
status_code,
headers,
Some(body_json),
@@ -6187,7 +6399,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
headers,
prefetched_usage_telemetry.clone(),
&provider_prefetched_body,
candidate_started_unix_secs,
stream_elapsed_ms_since(stream_started_at),
failure,
None,
)
.await;
}
@@ -6220,7 +6435,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
headers,
prefetched_usage_telemetry.clone(),
&provider_prefetched_body,
candidate_started_unix_secs,
stream_elapsed_ms_since(stream_started_at),
failure,
None,
)
.await;
}
@@ -6251,7 +6469,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
headers,
prefetched_usage_telemetry.clone(),
&provider_prefetched_body,
candidate_started_unix_secs,
stream_elapsed_ms_since(stream_started_at),
failure,
None,
)
.await;
}
@@ -6335,7 +6556,10 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
headers,
prefetched_usage_telemetry.clone(),
&provider_prefetched_body,
candidate_started_unix_secs,
stream_elapsed_ms_since(stream_started_at),
build_stream_failure_from_execution_error(&error),
retry_scope_out.as_deref_mut(),
)
.await;
}
@@ -6429,7 +6653,6 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
let _stream_total_guard =
StageElapsedGuard::from_started_at("stream_total", stream_started_at_for_report);
let _provider_pool_in_flight_guard = provider_pool_in_flight_guard_for_report;
let max_stream_body_buffer_bytes = DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES;
let mut provider_buffered_body = Vec::new();
let mut buffered_body = Vec::new();
let mut provider_body_truncated = false;
@@ -7373,6 +7596,8 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
let report_context_for_payload = report_context_with_request_diagnostics(
report_context_for_payload,
request_diagnostics_for_report.as_ref(),
stream_started_at_for_report,
terminal_telemetry.as_ref(),
);
let usage_payload = build_stream_usage_payload(
trace_id_owned,
@@ -7433,6 +7658,8 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
let report_context_for_payload = report_context_with_request_diagnostics(
report_context_for_payload,
request_diagnostics_for_report.as_ref(),
stream_started_at_for_report,
terminal_telemetry.as_ref(),
);
submit_midstream_stream_failure(
&state_for_report,
@@ -7499,6 +7726,8 @@ async fn execute_stream_from_frame_stream_with_retry_scope(
let report_context_for_payload = report_context_with_request_diagnostics(
report_context_for_payload,
request_diagnostics_for_report.as_ref(),
stream_started_at_for_report,
terminal_telemetry.as_ref(),
);
let usage_payload = build_stream_usage_payload(
trace_id_owned.clone(),
@@ -7686,12 +7915,14 @@ mod tests {
StoredUsageSettlement, UsageSettlementInput,
};
use aether_data_contracts::repository::usage::UsageReadRepository;
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::repository::usage::{
StoredRequestUsageAudit, UpsertUsageRecord, UsageBodyCaptureState,
};
use aether_data_contracts::DataLayerError;
use aether_usage_runtime::{
UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageEvent, UsageEventData,
UsageEventType, UsageRecordWriter, UsageRuntimeAccess, UsageRuntimeConfig,
UsageSettlementWriter,
apply_usage_body_capture_policy_to_event, UsageBillingEventEnricher,
UsageBodyCapturePolicy, UsageEvent, UsageEventData, UsageEventType, UsageRecordWriter,
UsageRequestRecordLevel, UsageRuntimeAccess, UsageRuntimeConfig, UsageSettlementWriter,
};
use async_stream::stream;
use async_trait::async_trait;
@@ -8058,6 +8289,185 @@ mod tests {
.expect("execution should succeed")
}
async fn execute_prefetched_transport_failure(
stop_on_transport_errors: bool,
) -> AiAttemptExecutionOutcome<axum::http::Response<Body>> {
let request_id = if stop_on_transport_errors {
"req-prefetch-transport-stop"
} else {
"req-prefetch-transport-retry"
};
let plan = native_anthropic_stream_plan(request_id);
let provider_config = stop_on_transport_errors.then(|| {
json!({
"failover_rules": {
"stop_on_transport_errors": true,
}
})
});
let provider_catalog = provider_catalog_for_plan(&plan, provider_config);
let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests(
Arc::new(provider_catalog),
"development-key",
);
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(data_state);
let frame_stream = stream! {
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code: 200,
headers: BTreeMap::from([(
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
frame_type: StreamFrameType::Error,
payload: StreamFramePayload::Error {
error: ExecutionError {
kind: ExecutionErrorKind::Internal,
phase: ExecutionPhase::StreamRead,
message: "connection reset before first body byte".to_string(),
upstream_status: None,
retryable: true,
failover_recommended: true,
},
},
}));
}
.boxed();
let mut retry_scope = AiAttemptRetryScope::Provider;
let response = execute_stream_from_frame_stream_with_retry_scope(
&state,
plan,
&format!("trace-{request_id}"),
&test_decision(),
"claude_chat_stream",
Some("claude_chat_stream_success".to_string()),
Some(json!({
"request_id": request_id,
"candidate_id": format!("candidate-{request_id}"),
"candidate_index": 0,
"retry_index": 0,
"provider_api_format": "claude:messages",
"client_api_format": "claude:messages"
})),
crate::clock::current_unix_ms(),
Instant::now(),
RequestStageTrace::from_env(),
true,
frame_stream,
false,
None,
Some(&mut retry_scope),
None,
)
.await
.expect("prefetch transport execution should resolve");
match response {
Some(response) => AiAttemptExecutionOutcome::Responded(response),
None => AiAttemptExecutionOutcome::Retry {
scope: retry_scope,
fallback_response: None,
},
}
}
async fn execute_prefetched_http_status_failure(
continue_failover: bool,
) -> AiAttemptExecutionOutcome<axum::http::Response<Body>> {
let request_id = if continue_failover {
"req-prefetch-http-continue"
} else {
"req-prefetch-http-stop"
};
let plan = native_anthropic_stream_plan(request_id);
let failover_rules = if continue_failover {
json!({"continue_status_codes": [500]})
} else {
json!({"stop_status_codes": [500]})
};
let provider_catalog = provider_catalog_for_plan(
&plan,
Some(json!({
"failover_rules": failover_rules,
})),
);
let data_state = crate::data::GatewayDataState::with_provider_transport_reader_for_tests(
Arc::new(provider_catalog),
"development-key",
);
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(data_state);
let frame_stream = stream! {
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
frame_type: StreamFrameType::Headers,
payload: StreamFramePayload::Headers {
status_code: 200,
headers: BTreeMap::from([(
"content-type".to_string(),
"text/event-stream".to_string(),
)]),
},
}));
yield Ok::<Bytes, std::io::Error>(ndjson_frame(StreamFrame {
frame_type: StreamFrameType::Error,
payload: StreamFramePayload::Error {
error: ExecutionError {
kind: ExecutionErrorKind::Internal,
phase: ExecutionPhase::StreamRead,
message: "upstream returned 500 before the first body byte".to_string(),
upstream_status: Some(500),
retryable: true,
failover_recommended: true,
},
},
}));
}
.boxed();
let mut retry_scope = AiAttemptRetryScope::Provider;
let response = execute_stream_from_frame_stream_with_retry_scope(
&state,
plan,
&format!("trace-{request_id}"),
&test_decision(),
"claude_chat_stream",
Some("claude_chat_stream_success".to_string()),
Some(json!({
"request_id": request_id,
"candidate_id": format!("candidate-{request_id}"),
"candidate_index": 0,
"retry_index": 0,
"provider_api_format": "claude:messages",
"client_api_format": "claude:messages"
})),
crate::clock::current_unix_ms(),
Instant::now(),
RequestStageTrace::from_env(),
true,
frame_stream,
false,
None,
Some(&mut retry_scope),
None,
)
.await
.expect("prefetch HTTP status execution should resolve");
match response {
Some(response) => AiAttemptExecutionOutcome::Responded(response),
None => AiAttemptExecutionOutcome::Retry {
scope: retry_scope,
fallback_response: None,
},
}
}
fn native_anthropic_stream_plan(request_id: &str) -> ExecutionPlan {
ExecutionPlan {
request_id: request_id.to_string(),
@@ -8126,6 +8536,7 @@ mod tests {
stream_usage_observer: None,
stream_usage_observer_buffered: Vec::new(),
provider_error_inspection: ProviderStreamErrorInspection::default(),
max_stream_body_buffer_bytes: super::DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
provider_buffered_body: Vec::new(),
buffered_body: Vec::new(),
provider_body_truncated: false,
@@ -9240,6 +9651,7 @@ mod tests {
stream_usage_observer: None,
stream_usage_observer_buffered: Vec::new(),
provider_error_inspection: ProviderStreamErrorInspection::default(),
max_stream_body_buffer_bytes: super::DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
provider_buffered_body: Vec::new(),
buffered_body: Vec::new(),
provider_body_truncated: false,
@@ -9655,6 +10067,79 @@ mod tests {
assert!(truncated);
}
#[test]
fn stream_capture_policy_keeps_full_unbounded_and_caps_basic_analysis_buffer() {
let oversized_chunk = vec![b'x'; super::BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES + 1];
let full_limit =
super::stream_body_buffer_limit_for_record_level(UsageRequestRecordLevel::Full);
assert_eq!(full_limit, usize::MAX);
let mut full_buffer = Vec::new();
let mut full_truncated = false;
super::append_stream_capture_bytes(
&mut full_buffer,
&oversized_chunk,
full_limit,
&mut full_truncated,
);
assert_eq!(full_buffer, oversized_chunk);
assert!(!full_truncated);
let (full_body, full_state) =
super::build_stream_body_capture(&full_buffer, full_truncated);
assert!(full_body.is_some());
assert_eq!(full_state, Some(UsageBodyCaptureState::Inline));
drop(full_body);
let basic_limit =
super::stream_body_buffer_limit_for_record_level(UsageRequestRecordLevel::Basic);
assert_eq!(basic_limit, super::BASIC_STREAM_BODY_ANALYSIS_LIMIT_BYTES);
let mut basic_buffer = Vec::new();
let mut basic_truncated = false;
super::append_stream_capture_bytes(
&mut basic_buffer,
&oversized_chunk,
basic_limit,
&mut basic_truncated,
);
assert_eq!(basic_buffer.len(), basic_limit);
assert!(basic_truncated);
let (basic_body, basic_state) =
super::build_stream_body_capture(&basic_buffer, basic_truncated);
assert!(basic_body.is_some());
assert_eq!(basic_state, Some(UsageBodyCaptureState::Truncated));
let mut event = UsageEvent::new(
UsageEventType::Completed,
"req-basic-stream-capture",
UsageEventData {
provider_name: "provider".to_string(),
model: "model".to_string(),
response_body: basic_body.map(Value::String),
response_body_state: basic_state,
client_response_body: Some(json!("captured client body")),
client_response_body_state: Some(UsageBodyCaptureState::Truncated),
..UsageEventData::default()
},
);
apply_usage_body_capture_policy_to_event(
UsageBodyCapturePolicy {
record_level: UsageRequestRecordLevel::Basic,
},
&mut event,
);
assert_eq!(event.data.response_body, None);
assert_eq!(
event.data.response_body_state,
Some(UsageBodyCaptureState::Disabled)
);
assert_eq!(event.data.client_response_body, None);
assert_eq!(
event.data.client_response_body_state,
Some(UsageBodyCaptureState::Disabled)
);
}
#[test]
fn provider_error_inspection_detects_response_failed_at_every_chunk_boundary() {
let body = concat!(
@@ -9719,6 +10204,45 @@ mod tests {
);
}
#[tokio::test]
async fn prefetched_transport_failure_retries_by_default() {
assert!(matches!(
execute_prefetched_transport_failure(false).await,
AiAttemptExecutionOutcome::Retry {
scope: AiAttemptRetryScope::Candidate,
fallback_response: None,
}
));
}
#[tokio::test]
async fn prefetched_transport_failure_can_stop_without_matching_http_status_rules() {
let AiAttemptExecutionOutcome::Responded(response) =
execute_prefetched_transport_failure(true).await
else {
panic!("transport stop policy should return a local response");
};
assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
}
#[tokio::test]
async fn prefetched_http_error_frame_honors_continue_status_codes() {
assert!(matches!(
execute_prefetched_http_status_failure(true).await,
AiAttemptExecutionOutcome::Retry { .. }
));
}
#[tokio::test]
async fn prefetched_http_error_frame_honors_stop_status_codes() {
let AiAttemptExecutionOutcome::Responded(response) =
execute_prefetched_http_status_failure(false).await
else {
panic!("HTTP stop policy should return the upstream error");
};
assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
}
fn tunnel_proxy_snapshot(base_url: String) -> aether_contracts::ProxySnapshot {
aether_contracts::ProxySnapshot {
enabled: Some(true),
@@ -1,5 +1,7 @@
use aether_ai_serving::AiAttemptRetryScope;
use aether_contracts::{ExecutionError, ExecutionPlan, ExecutionTelemetry};
use aether_contracts::{
ExecutionError, ExecutionErrorKind, ExecutionPhase, ExecutionPlan, ExecutionTelemetry,
};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
use aether_scheduler_core::SchedulerRequestCandidateStatusUpdate;
use aether_usage_runtime::{
@@ -22,13 +24,14 @@ use crate::execution_runtime::submission::{
use crate::log_ids::short_request_id;
use crate::orchestration::{
apply_local_execution_effect, classify_failure_disposition,
resolve_local_failover_analysis_for_attempt, with_upstream_response_report_context,
resolve_local_failover_analysis_for_attempt,
resolve_local_transport_failover_analysis_for_attempt, with_upstream_response_report_context,
LocalAdaptiveRateLimitEffect, LocalAttemptFailureEffect, LocalExecutionEffect,
LocalExecutionEffectContext, LocalFailoverAnalysis, LocalFailoverDecision,
LocalHealthFailureEffect, LocalOAuthInvalidationEffect, LocalPoolErrorEffect,
};
use crate::request_candidate_runtime::record_report_request_candidate_status;
use crate::request_diagnostics::attach_current_request_diagnostics_to_report_context;
use crate::request_diagnostics::attach_current_request_diagnostics_and_candidate_timing_to_report_context;
use crate::usage::submit_sync_report;
use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
@@ -37,6 +40,9 @@ pub(super) struct StreamFailureReport {
pub(super) status_code: u16,
pub(super) error_type: String,
pub(super) error_message: String,
upstream_status_code: Option<u16>,
transport_error: bool,
honor_http_failover: bool,
extra_error_fields: Map<String, Value>,
provider_body_json: Option<Value>,
}
@@ -68,6 +74,9 @@ impl StreamFailureReport {
status_code,
error_type,
error_message,
upstream_status_code: _,
transport_error: _,
honor_http_failover: _,
mut extra_error_fields,
provider_body_json,
} = self;
@@ -110,6 +119,26 @@ pub(super) fn build_stream_failure_report(
status_code,
error_type,
error_message,
upstream_status_code: Some(status_code),
transport_error: false,
honor_http_failover: false,
extra_error_fields: Map::new(),
provider_body_json: None,
}
}
pub(super) fn build_stream_transport_failure_report(
error_type: impl Into<String>,
error_message: impl Into<String>,
status_code: u16,
) -> StreamFailureReport {
StreamFailureReport {
status_code,
error_type: error_type.into(),
error_message: error_message.into(),
upstream_status_code: None,
transport_error: true,
honor_http_failover: false,
extra_error_fields: Map::new(),
provider_body_json: None,
}
@@ -118,7 +147,19 @@ pub(super) fn build_stream_failure_report(
pub(super) fn build_stream_failure_from_execution_error(
error: &ExecutionError,
) -> StreamFailureReport {
let status_code = error.upstream_status.unwrap_or(502);
let transport_error = execution_error_is_transport(error);
let status_code = error.upstream_status.unwrap_or_else(|| {
if matches!(
error.kind,
ExecutionErrorKind::ConnectTimeout
| ExecutionErrorKind::FirstByteTimeout
| ExecutionErrorKind::ReadTimeout
) {
504
} else {
502
}
});
let error_type = serde_json::to_value(&error.kind)
.ok()
.and_then(|value| value.as_str().map(ToOwned::to_owned))
@@ -141,6 +182,9 @@ pub(super) fn build_stream_failure_from_execution_error(
status_code,
error_type,
error_message,
upstream_status_code: error.upstream_status,
transport_error,
honor_http_failover: error.upstream_status.is_some(),
extra_error_fields: error_object,
provider_body_json: None,
}
@@ -168,11 +212,40 @@ pub(super) fn build_stream_failure_from_provider_error_body(
status_code,
error_type,
error_message,
upstream_status_code: Some(status_code),
transport_error: false,
honor_http_failover: true,
extra_error_fields: Map::new(),
provider_body_json: Some(body_json.clone()),
}
}
fn execution_error_is_transport(error: &ExecutionError) -> bool {
if error.upstream_status.is_some() {
return false;
}
let explicit_transport_kind = matches!(
error.kind,
ExecutionErrorKind::ConnectTimeout
| ExecutionErrorKind::FirstByteTimeout
| ExecutionErrorKind::ReadTimeout
| ExecutionErrorKind::TlsError
| ExecutionErrorKind::ProxyError
| ExecutionErrorKind::ProtocolError
);
let retryable_internal_transport_phase = matches!(error.kind, ExecutionErrorKind::Internal)
&& (error.retryable || error.failover_recommended)
&& matches!(
error.phase,
ExecutionPhase::Connect
| ExecutionPhase::Handshake
| ExecutionPhase::Write
| ExecutionPhase::FirstByte
| ExecutionPhase::StreamRead
);
explicit_transport_kind || retryable_internal_transport_phase
}
fn first_non_empty_error_text(
error_object: Option<&Map<String, Value>>,
body_object: Option<&Map<String, Value>>,
@@ -205,6 +278,8 @@ fn build_stream_failure_sync_payload(
failure: StreamFailureReport,
) -> GatewaySyncReportRequest {
let status_code = failure.status_code;
let upstream_status_code = failure.upstream_status_code;
let transport_error = failure.transport_error;
let (body, client_body) = failure.into_body_jsons();
headers.retain(|name, _| {
!name.eq_ignore_ascii_case("content-encoding")
@@ -212,23 +287,31 @@ fn build_stream_failure_sync_payload(
&& !name.eq_ignore_ascii_case("content-type")
});
headers.insert("content-type".to_string(), "application/json".to_string());
let report_context = with_upstream_response_report_context(
report_context.as_ref(),
status_code,
Some(&headers),
Some(&body),
None,
None,
)
.or(report_context);
let report_context = upstream_status_code
.and_then(|upstream_status_code| {
with_upstream_response_report_context(
report_context.as_ref(),
upstream_status_code,
Some(&headers),
Some(&body),
None,
None,
)
})
.or(report_context);
let report_context = report_context.map(|mut context| {
if let Some(object) = context.as_object_mut() {
let response_headers = serde_json::to_value(&headers).unwrap_or(Value::Null);
object.insert(
"provider_response_headers".to_string(),
response_headers.clone(),
);
if upstream_status_code.is_some() {
object.insert(
"provider_response_headers".to_string(),
response_headers.clone(),
);
}
object.insert("client_response_headers".to_string(), response_headers);
if transport_error {
object.insert("transport_error".to_string(), Value::Bool(true));
}
}
context
});
@@ -265,6 +348,7 @@ async fn record_stream_sync_failure(
plan: &ExecutionPlan,
report_context: Option<&Value>,
payload: &GatewaySyncReportRequest,
candidate_status_code: Option<u16>,
started_at_unix_ms: Option<u64>,
handling: StreamFailureHandling,
) -> LocalFailoverAnalysis {
@@ -361,8 +445,19 @@ async fn record_stream_sync_failure(
LocalFailoverDecision::RetryNextCandidate
);
if !matches!(handling, StreamFailureHandling::HonorLocalFailover) || !retrying_next_candidate {
crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started();
let report_context_with_diagnostics =
attach_current_request_diagnostics_to_report_context(report_context);
attach_current_request_diagnostics_and_candidate_timing_to_report_context(
report_context,
payload
.telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms),
payload
.telemetry
.as_ref()
.and_then(|telemetry| telemetry.ttfb_ms),
);
let context_seed = build_terminal_usage_context_seed(
plan,
report_context_with_diagnostics.as_ref().or(report_context),
@@ -383,7 +478,7 @@ async fn record_stream_sync_failure(
report_context,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(payload.status_code),
status_code: candidate_status_code,
error_type: Some(error_type.to_string()),
error_message: Some(error_message.to_string()),
latency_ms: payload
@@ -439,6 +534,7 @@ pub(super) async fn handle_prefetch_provider_private_stream_error(
plan,
payload.report_context.as_ref(),
&payload,
Some(status_code),
None,
StreamFailureHandling::HonorLocalFailover,
)
@@ -505,9 +601,15 @@ pub(super) async fn handle_prefetch_stream_failure(
headers: std::collections::BTreeMap<String, String>,
telemetry: Option<ExecutionTelemetry>,
buffered_body: &[u8],
candidate_started_unix_ms: u64,
candidate_elapsed_ms: u64,
failure: StreamFailureReport,
retry_scope_out: Option<&mut AiAttemptRetryScope>,
) -> Result<Option<Response<Body>>, GatewayError> {
let payload = build_stream_failure_sync_payload(
let transport_error = failure.transport_error;
let candidate_status_code = failure.upstream_status_code;
let honor_http_failover = failure.honor_http_failover;
let mut payload = build_stream_failure_sync_payload(
trace_id,
report_kind.to_string(),
report_context,
@@ -516,15 +618,179 @@ pub(super) async fn handle_prefetch_stream_failure(
buffered_body,
failure,
);
record_stream_sync_failure(
if transport_error {
let telemetry = payload.telemetry.get_or_insert(ExecutionTelemetry {
ttfb_ms: None,
elapsed_ms: None,
upstream_bytes: None,
});
telemetry.elapsed_ms.get_or_insert(candidate_elapsed_ms);
return handle_prefetch_transport_stream_failure(
state,
trace_id,
decision,
plan,
request_id,
candidate_id,
payload,
candidate_started_unix_ms,
candidate_elapsed_ms,
retry_scope_out,
)
.await;
}
let failure_analysis = record_stream_sync_failure(
state,
plan,
payload.report_context.as_ref(),
&payload,
candidate_status_code,
None,
StreamFailureHandling::Terminal,
if honor_http_failover {
StreamFailureHandling::HonorLocalFailover
} else {
StreamFailureHandling::Terminal
},
)
.await;
if honor_http_failover
&& matches!(
failure_analysis.decision,
LocalFailoverDecision::RetryNextCandidate
)
{
let failure_disposition = classify_failure_disposition(
&plan.provider_api_format,
failure_analysis.classification,
payload.status_code,
);
if let Some(retry_scope) = retry_scope_out {
*retry_scope = ai_attempt_retry_scope_from_failure_disposition(failure_disposition);
}
warn!(
event_name = "local_stream_candidate_retry_scheduled",
log_type = "event",
trace_id = %trace_id,
request_id = %request_id,
candidate_id = ?candidate_id,
status_code = payload.status_code,
failover_classification = failure_analysis.classification.as_str(),
"gateway local stream decision retrying next candidate after prefetched execution error"
);
return Ok(None);
}
let response =
submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?;
Ok(Some(attach_control_metadata_headers(
response,
Some(request_id),
candidate_id,
)?))
}
#[allow(clippy::too_many_arguments)]
async fn handle_prefetch_transport_stream_failure(
state: &AppState,
trace_id: &str,
decision: &GatewayControlDecision,
plan: &ExecutionPlan,
request_id: &str,
candidate_id: Option<&str>,
payload: GatewaySyncReportRequest,
candidate_started_unix_ms: u64,
candidate_elapsed_ms: u64,
retry_scope_out: Option<&mut AiAttemptRetryScope>,
) -> Result<Option<Response<Body>>, GatewayError> {
let error_type = stream_failure_body_field(&payload, "type").unwrap_or("internal");
let error_message = stream_failure_body_field(&payload, "message").unwrap_or_default();
if matches!(error_type, "first_byte_timeout" | "read_timeout") {
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan,
report_context: payload.report_context.as_ref(),
},
LocalExecutionEffect::PoolStreamTimeout,
)
.await;
}
let analysis = resolve_local_transport_failover_analysis_for_attempt(
state,
plan,
payload.report_context.as_ref(),
)
.await;
let retrying_next_candidate =
matches!(analysis.decision, LocalFailoverDecision::RetryNextCandidate);
if !retrying_next_candidate {
crate::execution_runtime::mark_stream_candidate_watchdog_terminal_started();
let report_context_with_diagnostics =
attach_current_request_diagnostics_and_candidate_timing_to_report_context(
payload.report_context.as_ref(),
payload
.telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms)
.or(Some(candidate_elapsed_ms)),
payload
.telemetry
.as_ref()
.and_then(|telemetry| telemetry.ttfb_ms),
);
let context_seed = build_terminal_usage_context_seed(
plan,
report_context_with_diagnostics
.as_ref()
.or(payload.report_context.as_ref()),
);
let payload_seed = build_sync_terminal_usage_payload_seed(&payload);
state
.usage_runtime
.record_sync_terminal(
state.usage_lifecycle_data_state().as_ref(),
context_seed,
payload_seed,
)
.await;
}
let terminal_unix_ms = current_request_candidate_unix_ms();
record_report_request_candidate_status(
state,
payload.report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some(error_type.to_string()),
error_message: Some(error_message.to_string()),
latency_ms: payload
.telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms)
.or(Some(candidate_elapsed_ms)),
started_at_unix_ms: Some(candidate_started_unix_ms),
finished_at_unix_ms: Some(terminal_unix_ms),
},
)
.await;
if retrying_next_candidate {
if let Some(retry_scope) = retry_scope_out {
*retry_scope = AiAttemptRetryScope::Candidate;
}
warn!(
event_name = "local_stream_transport_retry_scheduled",
log_type = "event",
trace_id = %trace_id,
request_id = %request_id,
candidate_id = ?candidate_id,
transport_classification = analysis.classification.as_str(),
"gateway retrying next candidate after precommit transport failure"
);
return Ok(None);
}
let response =
submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?;
@@ -553,6 +819,7 @@ pub(super) async fn submit_midstream_stream_failure(
return;
};
let candidate_status_code = failure.upstream_status_code;
let payload = build_stream_failure_sync_payload(
trace_id,
report_kind,
@@ -567,6 +834,7 @@ pub(super) async fn submit_midstream_stream_failure(
plan,
payload.report_context.as_ref(),
&payload,
candidate_status_code,
Some(started_at_unix_ms),
StreamFailureHandling::Terminal,
)
@@ -590,10 +858,60 @@ pub(super) async fn submit_midstream_stream_failure(
mod tests {
use std::collections::BTreeMap;
use aether_contracts::{ExecutionError, ExecutionErrorKind, ExecutionPhase};
use base64::Engine as _;
use serde_json::json;
use super::{build_stream_failure_from_provider_error_body, build_stream_failure_sync_payload};
use super::{
build_stream_failure_from_execution_error, build_stream_failure_from_provider_error_body,
build_stream_failure_sync_payload, build_stream_transport_failure_report,
};
#[test]
fn committed_transport_failure_has_no_upstream_status() {
for status_code in [502, 504] {
let failure = build_stream_transport_failure_report(
"execution_runtime_stream_read_error",
"upstream disconnected",
status_code,
);
assert_eq!(failure.status_code, status_code);
assert_eq!(failure.upstream_status_code, None);
assert!(failure.transport_error);
assert!(!failure.honor_http_failover);
}
}
#[test]
fn precommit_protocol_error_is_transport_without_upstream_status() {
let failure = build_stream_failure_from_execution_error(&ExecutionError {
kind: ExecutionErrorKind::ProtocolError,
phase: ExecutionPhase::StreamRead,
message: "connection reset".to_string(),
upstream_status: None,
retryable: true,
failover_recommended: true,
});
assert!(failure.transport_error);
assert_eq!(failure.upstream_status_code, None);
assert_eq!(failure.status_code, 502);
}
#[test]
fn cancelled_stream_is_not_reclassified_as_transport_retry() {
let failure = build_stream_failure_from_execution_error(&ExecutionError {
kind: ExecutionErrorKind::Cancelled,
phase: ExecutionPhase::StreamRead,
message: "downstream cancelled".to_string(),
upstream_status: None,
retryable: true,
failover_recommended: true,
});
assert!(!failure.transport_error);
}
#[test]
fn midstream_failure_trace_uses_terminal_error_instead_of_buffered_sse() {
@@ -163,7 +163,7 @@ pub(crate) fn build_direct_execution_frame_stream(
let error_frame = if let Some(timeout) = first_byte_timeout {
encode_first_byte_timeout_frame(timeout)
} else {
encode_error_frame(status_code, message)
encode_error_frame(message)
};
match error_frame {
Ok(frame) => yield Ok(frame),
@@ -245,7 +245,7 @@ pub(crate) fn build_direct_execution_frame_stream(
error = %message,
"upstream body stream read error"
);
match encode_error_frame(status_code, message) {
match encode_error_frame(message) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
@@ -329,7 +329,7 @@ pub(crate) fn build_direct_execution_frame_stream(
error = %message,
"upstream body stream read error"
);
match encode_error_frame(status_code, message) {
match encode_error_frame(message) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
@@ -411,7 +411,7 @@ pub(crate) fn build_direct_execution_frame_stream(
error = %message,
"upstream body stream read error"
);
match encode_error_frame(status_code, message) {
match encode_error_frame(message) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
@@ -493,7 +493,7 @@ pub(crate) fn build_direct_execution_frame_stream(
error = %message,
"upstream body stream read error"
);
match encode_error_frame(status_code, message) {
match encode_error_frame(message) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
@@ -570,7 +570,7 @@ pub(crate) fn build_direct_execution_frame_stream(
error = %message,
"upstream body stream read error"
);
match encode_error_frame(status_code, message) {
match encode_error_frame(message) {
Ok(frame) => yield Ok(frame),
Err(encode_err) => {
yield Err(encode_err);
@@ -648,17 +648,17 @@ fn encode_data_frame(chunk: &Bytes) -> Result<Bytes, IoError> {
})
}
fn encode_error_frame(status_code: u16, message: String) -> Result<Bytes, IoError> {
fn encode_error_frame(message: String) -> Result<Bytes, IoError> {
encode_stream_frame_ndjson(&StreamFrame {
frame_type: StreamFrameType::Error,
payload: StreamFramePayload::Error {
error: ExecutionError {
kind: ExecutionErrorKind::Internal,
kind: ExecutionErrorKind::ProtocolError,
phase: ExecutionPhase::StreamRead,
message,
upstream_status: Some(status_code),
retryable: false,
failover_recommended: false,
upstream_status: None,
retryable: true,
failover_recommended: true,
},
},
})
@@ -672,7 +672,7 @@ fn encode_first_byte_timeout_frame(timeout: Duration) -> Result<Bytes, IoError>
kind: ExecutionErrorKind::FirstByteTimeout,
phase: ExecutionPhase::FirstByte,
message: stream_first_byte_timeout_message(timeout),
upstream_status: Some(504),
upstream_status: None,
retryable: true,
failover_recommended: true,
},
@@ -1490,6 +1490,13 @@ mod tests {
.is_some_and(
|message| message.contains("provider stream first byte timeout after 50 ms")
));
let error = error_frame
.get("payload")
.and_then(|payload| payload.get("error"))
.expect("timeout error should exist");
assert_eq!(error.get("upstream_status"), None);
assert_eq!(error.get("retryable"), Some(&Value::Bool(true)));
assert_eq!(error.get("failover_recommended"), Some(&Value::Bool(true)));
}
#[tokio::test]
@@ -82,7 +82,11 @@ use crate::request_candidate_runtime::{
record_local_request_candidate_status, record_local_request_candidate_status_snapshot,
snapshot_local_request_candidate_status,
};
use crate::request_diagnostics::attach_current_request_diagnostics_to_report_context;
use crate::request_diagnostics::{
attach_current_request_diagnostics_and_candidate_start_timing_to_report_context,
attach_request_diagnostics_to_report_context, calibrate_candidate_first_byte_elapsed_ms,
current_request_diagnostics, RequestDiagnostics,
};
use crate::usage::{spawn_sync_report, submit_sync_report};
use crate::video_tasks::VideoTaskSyncReportMode;
use crate::{usage::GatewaySyncReportRequest, AppState, GatewayError};
@@ -108,6 +112,22 @@ const OPENAI_IMAGE_SYNC_JSON_HEARTBEAT_BYTES: &[u8] = b"\n";
const OPENAI_IMAGE_SYNC_PROGRESS_WRITE_INTERVAL: Duration = Duration::from_secs(5);
const INVALID_GEMINI_PROVIDER_SUCCESS_MESSAGE: &str = "Provider returned HTTP 200 but the Gemini response did not contain visible model output; refusing to finalize it as a successful response.";
fn elapsed_ms_since(started_at: Instant) -> u64 {
started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64
}
fn calibrated_sync_candidate_first_byte_elapsed_ms(
candidate_started_at: Instant,
result: &ExecutionResult,
) -> Option<u64> {
let telemetry = result.telemetry.as_ref()?;
calibrate_candidate_first_byte_elapsed_ms(
elapsed_ms_since(candidate_started_at),
telemetry.elapsed_ms,
telemetry.ttfb_ms,
)
}
#[derive(Debug)]
struct SyncExecutionFailure {
error_type: &'static str,
@@ -143,7 +163,9 @@ struct SyncAttemptTerminalGuard {
state: AppState,
plan: ExecutionPlan,
report_context: Option<Value>,
request_diagnostics: Option<Arc<RequestDiagnostics>>,
candidate_started_unix_ms: u64,
candidate_started_at: Instant,
armed: bool,
}
@@ -153,12 +175,15 @@ impl SyncAttemptTerminalGuard {
plan: &ExecutionPlan,
report_context: Option<Value>,
candidate_started_unix_ms: u64,
candidate_started_at: Instant,
) -> Self {
Self {
state: state.clone(),
plan: plan.clone(),
report_context,
request_diagnostics: current_request_diagnostics(),
candidate_started_unix_ms,
candidate_started_at,
armed: true,
}
}
@@ -176,7 +201,9 @@ impl SyncAttemptTerminalGuard {
self.state.clone(),
self.plan.clone(),
self.report_context.clone(),
self.request_diagnostics.clone(),
self.candidate_started_unix_ms,
self.candidate_started_at,
UsageEventType::Failed,
RequestCandidateStatus::Failed,
StatusCode::INTERNAL_SERVER_ERROR.as_u16(),
@@ -196,14 +223,18 @@ impl Drop for SyncAttemptTerminalGuard {
let state = self.state.clone();
let plan = self.plan.clone();
let report_context = self.report_context.clone();
let request_diagnostics = self.request_diagnostics.clone();
let candidate_started_unix_ms = self.candidate_started_unix_ms;
let candidate_started_at = self.candidate_started_at;
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
record_sync_attempt_forced_terminal_state(
state,
plan,
report_context,
request_diagnostics,
candidate_started_unix_ms,
candidate_started_at,
UsageEventType::Cancelled,
RequestCandidateStatus::Cancelled,
499,
@@ -229,7 +260,9 @@ async fn record_sync_attempt_forced_terminal_state(
state: AppState,
plan: ExecutionPlan,
report_context: Option<Value>,
request_diagnostics: Option<Arc<RequestDiagnostics>>,
candidate_started_unix_ms: u64,
candidate_started_at: Instant,
usage_event_type: UsageEventType,
candidate_status: RequestCandidateStatus,
status_code: u16,
@@ -237,8 +270,10 @@ async fn record_sync_attempt_forced_terminal_state(
error_message: impl Into<String>,
) {
let error_message = error_message.into();
let report_context =
attach_request_diagnostics_to_report_context(report_context, request_diagnostics.as_ref());
let terminal_unix_ms = current_request_candidate_unix_ms();
let latency_ms = terminal_unix_ms.saturating_sub(candidate_started_unix_ms);
let latency_ms = elapsed_ms_since(candidate_started_at);
record_local_request_candidate_status(
&state,
&plan,
@@ -409,6 +444,41 @@ fn maybe_store_sync_execution_failure_fallback(
Ok(())
}
async fn maybe_build_sync_transport_error_stop_response(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
trace_id: &str,
decision: &GatewayControlDecision,
error_type: &str,
error_message: &str,
elapsed_ms: u64,
) -> Result<Option<Response<Body>>, GatewayError> {
let analysis = crate::orchestration::resolve_local_transport_failover_analysis_for_attempt(
state,
plan,
report_context,
)
.await;
if !matches!(analysis.decision, LocalFailoverDecision::StopLocalFailover) {
return Ok(None);
}
crate::execution_runtime::build_transport_error_stop_response(
state,
plan,
report_context,
trace_id,
decision,
StatusCode::BAD_GATEWAY.as_u16(),
error_type,
error_message,
elapsed_ms,
)
.await
.map(Some)
}
struct ImplicitSyncFinalizeOutcome {
payload: GatewaySyncReportRequest,
outcome: LocalCoreSyncFinalizeOutcome,
@@ -496,9 +566,15 @@ async fn record_sync_terminal_usage(
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
payload: &GatewaySyncReportRequest,
candidate_started_at: Instant,
candidate_first_byte_elapsed_ms: Option<u64>,
) {
let report_context_with_diagnostics =
attach_current_request_diagnostics_to_report_context(report_context);
attach_current_request_diagnostics_and_candidate_start_timing_to_report_context(
report_context,
candidate_started_at,
candidate_first_byte_elapsed_ms,
);
let context_seed = build_terminal_usage_context_seed(
plan,
report_context_with_diagnostics.as_ref().or(report_context),
@@ -519,9 +595,19 @@ async fn record_sync_terminal_usage_and_disarm_guard(
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
payload: &GatewaySyncReportRequest,
candidate_started_at: Instant,
candidate_first_byte_elapsed_ms: Option<u64>,
terminal_guard: &mut SyncAttemptTerminalGuard,
) {
record_sync_terminal_usage(state, plan, report_context, payload).await;
record_sync_terminal_usage(
state,
plan,
report_context,
payload,
candidate_started_at,
candidate_first_byte_elapsed_ms,
)
.await;
terminal_guard.disarm();
}
@@ -1850,6 +1936,7 @@ async fn execute_execution_runtime_sync_impl(
.and_then(|context| context.candidate_index)
.map(|value| value.to_string())
.unwrap_or_else(|| "-".to_string());
let candidate_started_at = Instant::now();
let candidate_started_unix_secs = current_request_candidate_unix_ms();
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
let usage_data = state.usage_lifecycle_data_state().as_ref().clone();
@@ -1877,6 +1964,7 @@ async fn execute_execution_runtime_sync_impl(
&plan,
report_context.clone(),
candidate_started_unix_secs,
candidate_started_at,
);
let result = (async {
let _provider_pool_in_flight_guard = acquire_provider_pool_in_flight_guard(
@@ -1922,6 +2010,11 @@ async fn execute_execution_runtime_sync_impl(
{
Ok(result) => result,
Err(err) => {
let failure_error_type = err.error_type;
let failure_message = err.message.clone();
let failure_latency_ms = err
.latency_ms
.unwrap_or_else(|| elapsed_ms_since(candidate_started_at));
maybe_store_sync_execution_failure_fallback(
&err,
&plan,
@@ -1952,19 +2045,34 @@ async fn execute_execution_runtime_sync_impl(
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: err.status_code,
error_type: Some(err.error_type.to_string()),
status_code: None,
error_type: Some(failure_error_type.to_string()),
error_message: Some(err.message),
latency_ms: err.latency_ms,
latency_ms: Some(failure_latency_ms),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_sync_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
failure_error_type,
failure_message.as_str(),
failure_latency_ms,
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
},
Err(err) => {
let transport_error_message = err.to_string();
warn!(
event_name = "chatgpt_web_image_execution_unavailable",
log_type = "ops",
@@ -1990,18 +2098,33 @@ async fn execute_execution_runtime_sync_impl(
error_type: Some(
"chatgpt_web_image_execution_unavailable".to_string(),
),
error_message: Some(err.to_string()),
latency_ms: None,
error_message: Some(transport_error_message.clone()),
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_sync_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
"chatgpt_web_image_execution_unavailable",
transport_error_message.as_str(),
elapsed_ms_since(candidate_started_at),
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
}
}
Err(err) => {
let transport_error_message = err.to_string();
warn!(
event_name = "grok_execution_unavailable",
log_type = "ops",
@@ -2025,13 +2148,27 @@ async fn execute_execution_runtime_sync_impl(
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("grok_execution_unavailable".to_string()),
error_message: Some(err.to_string()),
latency_ms: None,
error_message: Some(transport_error_message.clone()),
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_sync_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
"grok_execution_unavailable",
transport_error_message.as_str(),
elapsed_ms_since(candidate_started_at),
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
}
@@ -2042,6 +2179,7 @@ async fn execute_execution_runtime_sync_impl(
match (override_fn.0)(&plan) {
Ok(result) => result,
Err(err) => {
let transport_error_message = format!("{err:?}");
warn!(
event_name = "sync_execution_runtime_test_override_failed",
log_type = "ops",
@@ -2065,13 +2203,27 @@ async fn execute_execution_runtime_sync_impl(
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("execution_runtime_unavailable".to_string()),
error_message: Some(format!("{err:?}")),
latency_ms: None,
error_message: Some(transport_error_message.clone()),
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_sync_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
"execution_runtime_unavailable",
transport_error_message.as_str(),
elapsed_ms_since(candidate_started_at),
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
}
@@ -2111,6 +2263,11 @@ async fn execute_execution_runtime_sync_impl(
{
Ok(result) => result,
Err(err) => {
let failure_error_type = err.error_type;
let failure_message = err.message.clone();
let failure_latency_ms = err
.latency_ms
.unwrap_or_else(|| elapsed_ms_since(candidate_started_at));
maybe_store_sync_execution_failure_fallback(
&err,
&plan,
@@ -2141,19 +2298,34 @@ async fn execute_execution_runtime_sync_impl(
report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: err.status_code,
error_type: Some(err.error_type.to_string()),
status_code: None,
error_type: Some(failure_error_type.to_string()),
error_message: Some(err.message),
latency_ms: err.latency_ms,
latency_ms: Some(failure_latency_ms),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_sync_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
failure_error_type,
failure_message.as_str(),
failure_latency_ms,
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
},
Err(err) => {
let transport_error_message = err.to_string();
warn!(
event_name = "chatgpt_web_image_execution_unavailable",
log_type = "ops",
@@ -2179,17 +2351,32 @@ async fn execute_execution_runtime_sync_impl(
error_type: Some(
"chatgpt_web_image_execution_unavailable".to_string(),
),
error_message: Some(err.to_string()),
latency_ms: None,
error_message: Some(transport_error_message.clone()),
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_sync_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
"chatgpt_web_image_execution_unavailable",
transport_error_message.as_str(),
elapsed_ms_since(candidate_started_at),
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
},
Err(err) => {
let transport_error_message = err.to_string();
warn!(
event_name = "grok_execution_unavailable",
log_type = "ops",
@@ -2213,13 +2400,27 @@ async fn execute_execution_runtime_sync_impl(
status: RequestCandidateStatus::Failed,
status_code: None,
error_type: Some("grok_execution_unavailable".to_string()),
error_message: Some(err.to_string()),
latency_ms: None,
error_message: Some(transport_error_message.clone()),
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
)
.await;
if let Some(response) = maybe_build_sync_transport_error_stop_response(
state,
&plan,
report_context.as_ref(),
trace_id,
decision,
"grok_execution_unavailable",
transport_error_message.as_str(),
elapsed_ms_since(candidate_started_at),
)
.await?
{
return Ok(Some(response));
}
return Ok(None);
}
}
@@ -2237,6 +2438,7 @@ async fn execute_execution_runtime_sync_impl(
plan_candidate_id.as_deref(),
report_context.as_ref(),
candidate_started_unix_secs,
candidate_started_at,
)
.await?;
match remote_outcome {
@@ -2246,6 +2448,8 @@ async fn execute_execution_runtime_sync_impl(
}
}
};
let mut candidate_first_byte_elapsed_ms =
calibrated_sync_candidate_first_byte_elapsed_ms(candidate_started_at, &result);
let mut oauth_retry_attempted = false;
let (
result_error_type,
@@ -2331,6 +2535,11 @@ async fn execute_execution_runtime_sync_impl(
.await
{
Ok(retry_result) => {
candidate_first_byte_elapsed_ms =
calibrated_sync_candidate_first_byte_elapsed_ms(
candidate_started_at,
&retry_result,
);
result = retry_result;
continue;
}
@@ -2657,6 +2866,8 @@ async fn execute_execution_runtime_sync_impl(
&plan,
implicit_finalize.payload.report_context.as_ref(),
usage_payload,
candidate_started_at,
candidate_first_byte_elapsed_ms,
&mut terminal_guard,
)
.await;
@@ -2711,6 +2922,8 @@ async fn execute_execution_runtime_sync_impl(
&plan,
payload.report_context.as_ref(),
usage_payload,
candidate_started_at,
candidate_first_byte_elapsed_ms,
&mut terminal_guard,
)
.await;
@@ -2758,6 +2971,8 @@ async fn execute_execution_runtime_sync_impl(
&plan,
original_report_context.as_ref(),
&report_payload,
candidate_started_at,
candidate_first_byte_elapsed_ms,
&mut terminal_guard,
)
.await;
@@ -2793,6 +3008,8 @@ async fn execute_execution_runtime_sync_impl(
&plan,
payload.report_context.as_ref(),
&payload,
candidate_started_at,
candidate_first_byte_elapsed_ms,
&mut terminal_guard,
)
.await;
@@ -2840,6 +3057,8 @@ async fn execute_execution_runtime_sync_impl(
&plan,
payload.report_context.as_ref(),
&payload,
candidate_started_at,
candidate_first_byte_elapsed_ms,
&mut terminal_guard,
)
.await;
@@ -2867,6 +3086,8 @@ async fn execute_execution_runtime_sync_impl(
&plan,
payload.report_context.as_ref(),
&payload,
candidate_started_at,
candidate_first_byte_elapsed_ms,
&mut terminal_guard,
)
.await;
@@ -2903,6 +3124,8 @@ async fn execute_execution_runtime_sync_impl(
&plan,
usage_payload.report_context.as_ref(),
&usage_payload,
candidate_started_at,
candidate_first_byte_elapsed_ms,
&mut terminal_guard,
)
.await;
@@ -2995,6 +3218,7 @@ async fn execute_sync_via_remote_execution_runtime(
plan_candidate_id: Option<&str>,
report_context: Option<&serde_json::Value>,
candidate_started_unix_secs: u64,
candidate_started_at: Instant,
) -> Result<RemoteSyncFallbackOutcome, GatewayError> {
let response = match post_sync_plan_to_remote_execution_runtime(
state,
@@ -3025,7 +3249,7 @@ async fn execute_sync_via_remote_execution_runtime(
status_code: None,
error_type: Some("execution_runtime_unavailable".to_string()),
error_message: Some(format!("{err:?}")),
latency_ms: None,
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
@@ -3049,7 +3273,7 @@ async fn execute_sync_via_remote_execution_runtime(
"execution runtime returned HTTP {}",
response.status()
)),
latency_ms: None,
latency_ms: Some(elapsed_ms_since(candidate_started_at)),
started_at_unix_ms: Some(candidate_started_unix_secs),
finished_at_unix_ms: Some(terminal_unix_secs),
},
@@ -3373,6 +3597,7 @@ mod tests {
}));
ensure_execution_request_candidate_slot(&state, &mut plan, &mut report_context).await;
let candidate_started_at = Instant::now();
let started_at = current_request_candidate_unix_ms();
state.usage_runtime.record_pending(
state.usage_lifecycle_data_state().as_ref(),
@@ -3394,10 +3619,19 @@ mod tests {
)
.await;
{
let _guard =
SyncAttemptTerminalGuard::new(&state, &plan, report_context.clone(), started_at);
}
crate::request_diagnostics::scope_request_diagnostics(async {
crate::request_diagnostics::record_request_accepted_at(
Instant::now() - Duration::from_millis(25),
);
let _guard = SyncAttemptTerminalGuard::new(
&state,
&plan,
report_context.clone(),
started_at,
candidate_started_at,
);
})
.await;
let mut stored_usage = None;
for _ in 0..50 {
@@ -3418,6 +3652,17 @@ mod tests {
assert_eq!(stored_usage.billing_status, "void");
assert_eq!(stored_usage.status_code, Some(499));
assert_eq!(stored_usage.error_category.as_deref(), Some("cancelled"));
let request_metadata = stored_usage
.request_metadata
.as_ref()
.expect("cancelled usage should retain request diagnostics");
assert!(request_metadata
.get("end_to_end_time_ms")
.and_then(Value::as_u64)
.is_some());
assert!(request_metadata
.get("end_to_end_first_byte_time_ms")
.is_none());
let stored_candidates = request_candidate_repository
.list_by_request_id("sync-cancel-guard-request")
@@ -575,6 +575,8 @@ pub(crate) enum ExecutionRuntimeTransportError {
BrowserClientBuild(wreq::Error),
#[error("browser impersonation response body failed: {0}")]
BrowserBody(String),
#[error("{message}")]
UpstreamHttpStatus { status_code: u16, message: String },
#[error("failed to execute upstream request: {0}")]
UpstreamRequest(String),
#[error("upstream response {phase} body exceeds {limit_bytes} bytes")]
@@ -0,0 +1,139 @@
use std::collections::BTreeMap;
use std::future::Future;
use std::sync::{
atomic::{AtomicBool, Ordering},
Arc,
};
use aether_usage_runtime::{build_usage_event_data_seed, UsageEvent, UsageEventType};
use axum::body::Body;
use axum::http::Response;
use serde_json::{json, Value};
use crate::ai_serving::{build_core_error_body_for_client_format, LocalCoreSyncErrorKind};
use crate::api::response::{attach_control_metadata_headers, build_client_response_from_parts};
use crate::control::GatewayControlDecision;
use crate::request_diagnostics::attach_current_request_diagnostics_and_candidate_timing_to_report_context;
use crate::{AppState, GatewayError};
const TRANSPORT_ERROR_CLIENT_MESSAGE: &str =
"Upstream transport failed before an HTTP response was received";
#[derive(Debug, Default)]
pub(crate) struct StreamCandidateWatchdogProgress {
terminal_started: AtomicBool,
}
tokio::task_local! {
static STREAM_CANDIDATE_WATCHDOG_PROGRESS: Arc<StreamCandidateWatchdogProgress>;
}
impl StreamCandidateWatchdogProgress {
pub(crate) fn shared() -> Arc<Self> {
Arc::new(Self::default())
}
pub(crate) fn terminal_started(&self) -> bool {
self.terminal_started.load(Ordering::Acquire)
}
pub(crate) async fn scope<F>(self: Arc<Self>, future: F) -> F::Output
where
F: Future,
{
STREAM_CANDIDATE_WATCHDOG_PROGRESS.scope(self, future).await
}
}
pub(crate) fn mark_stream_candidate_watchdog_terminal_started() {
let _ = STREAM_CANDIDATE_WATCHDOG_PROGRESS.try_with(|progress| {
progress.terminal_started.store(true, Ordering::Release);
});
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn build_transport_error_stop_response(
state: &AppState,
plan: &aether_contracts::ExecutionPlan,
report_context: Option<&Value>,
trace_id: &str,
decision: &GatewayControlDecision,
client_status_code: u16,
error_type: &str,
error_message: &str,
elapsed_ms: u64,
) -> Result<Response<Body>, GatewayError> {
mark_stream_candidate_watchdog_terminal_started();
let client_body = build_core_error_body_for_client_format(
&plan.client_api_format,
TRANSPORT_ERROR_CLIENT_MESSAGE,
Some("upstream_transport_error"),
LocalCoreSyncErrorKind::ServerError,
)
.unwrap_or_else(|| {
json!({
"error": {
"type": "server_error",
"message": TRANSPORT_ERROR_CLIENT_MESSAGE,
"code": "upstream_transport_error",
}
})
});
let body_bytes =
serde_json::to_vec(&client_body).map_err(|err| GatewayError::Internal(err.to_string()))?;
let headers = BTreeMap::from([
("content-type".to_string(), "application/json".to_string()),
("content-length".to_string(), body_bytes.len().to_string()),
]);
if state.usage_runtime.is_enabled() {
let report_context_with_diagnostics =
attach_current_request_diagnostics_and_candidate_timing_to_report_context(
report_context,
Some(elapsed_ms),
None,
);
let mut usage_data = build_usage_event_data_seed(
plan,
report_context_with_diagnostics.as_ref().or(report_context),
);
usage_data.status_code = Some(client_status_code);
usage_data.error_message = Some(error_message.to_string());
usage_data.error_category = Some("server_error".to_string());
usage_data.response_time_ms = Some(elapsed_ms);
usage_data.response_headers = None;
usage_data.response_body = None;
usage_data.client_response_headers = Some(json!({"content-type": "application/json"}));
usage_data.client_response_body = Some(client_body);
let mut request_metadata = match usage_data.request_metadata.take() {
Some(Value::Object(object)) => object,
Some(other) => serde_json::Map::from_iter([("seed".to_string(), other)]),
None => serde_json::Map::new(),
};
request_metadata.insert("transport_error".to_string(), Value::Bool(true));
request_metadata.insert(
"transport_error_type".to_string(),
Value::String(error_type.to_string()),
);
usage_data.request_metadata = Some(Value::Object(request_metadata));
state
.usage_runtime
.record_terminal_event_direct(
state.usage_lifecycle_data_state().as_ref(),
UsageEvent::new(UsageEventType::Failed, plan.request_id.clone(), usage_data),
)
.await;
}
attach_control_metadata_headers(
build_client_response_from_parts(
client_status_code,
&headers,
Body::from(body_bytes),
trace_id,
Some(decision),
)?,
Some(plan.request_id.as_str()),
plan.candidate_id.as_deref(),
)
}
@@ -700,18 +700,9 @@ fn windsurf_execution_error_from_transport_error(
failover_recommended: false,
};
}
if is_windsurf_cascade_transport_error(err) {
return ExecutionError {
kind: ExecutionErrorKind::Upstream5xx,
phase,
message: format!("{message}; Windsurf IDE language server is unavailable"),
upstream_status: Some(503),
retryable: true,
failover_recommended: true,
};
}
let (kind, phase) = classify_windsurf_transport_execution_error(&lower, phase);
ExecutionError {
kind: ExecutionErrorKind::ProtocolError,
kind,
phase,
message,
upstream_status: None,
@@ -720,6 +711,57 @@ fn windsurf_execution_error_from_transport_error(
}
}
fn classify_windsurf_transport_execution_error(
message: &str,
phase: ExecutionPhase,
) -> (ExecutionErrorKind, ExecutionPhase) {
if message.contains("proxy") {
return (ExecutionErrorKind::ProxyError, ExecutionPhase::Connect);
}
if ["tls", "ssl", "certificate", "handshake"]
.iter()
.any(|needle| message.contains(needle))
{
return (ExecutionErrorKind::TlsError, ExecutionPhase::Handshake);
}
if message.contains("first byte") {
return (
ExecutionErrorKind::FirstByteTimeout,
ExecutionPhase::FirstByte,
);
}
if ["timed out", "timeout", "deadline exceeded"]
.iter()
.any(|needle| message.contains(needle))
{
let kind = match &phase {
ExecutionPhase::Connect | ExecutionPhase::Handshake | ExecutionPhase::Write => {
ExecutionErrorKind::ConnectTimeout
}
ExecutionPhase::FirstByte => ExecutionErrorKind::FirstByteTimeout,
_ => ExecutionErrorKind::ReadTimeout,
};
return (kind, phase);
}
if [
"dns",
"failed to lookup address",
"name or service not known",
"no such host",
"connection refused",
"tcp connect",
"kind=connect",
"connect error",
"failed to connect",
]
.iter()
.any(|needle| message.contains(needle))
{
return (ExecutionErrorKind::ProtocolError, ExecutionPhase::Connect);
}
(ExecutionErrorKind::ProtocolError, phase)
}
async fn poll_windsurf_cascade_with_transport_recovery<F>(
prepared: &PreparedCascade,
mut on_event: F,
@@ -4554,6 +4596,56 @@ mod tests {
assert!(execution_error.failover_recommended);
}
#[test]
fn windsurf_transport_errors_do_not_synthesize_provider_status() {
let cases = [
(
"tcp connect error: connection refused",
ExecutionErrorKind::ProtocolError,
ExecutionPhase::Connect,
),
(
"dns lookup failed: no such host",
ExecutionErrorKind::ProtocolError,
ExecutionPhase::Connect,
),
(
"TLS certificate handshake failed",
ExecutionErrorKind::TlsError,
ExecutionPhase::Handshake,
),
(
"proxy connection failed",
ExecutionErrorKind::ProxyError,
ExecutionPhase::Connect,
),
(
"cascade polling timed out",
ExecutionErrorKind::ReadTimeout,
ExecutionPhase::StreamRead,
),
(
"connection reset by peer",
ExecutionErrorKind::ProtocolError,
ExecutionPhase::StreamRead,
),
];
for (message, expected_kind, expected_phase) in cases {
let err = ExecutionRuntimeTransportError::UpstreamRequest(message.to_string());
let execution_error = super::windsurf_execution_error_from_transport_error(
&err,
ExecutionPhase::StreamRead,
);
assert_eq!(execution_error.kind, expected_kind, "{message}");
assert_eq!(execution_error.phase, expected_phase, "{message}");
assert_eq!(execution_error.upstream_status, None, "{message}");
assert!(execution_error.retryable, "{message}");
assert!(execution_error.failover_recommended, "{message}");
}
}
#[test]
fn windsurf_sanitizer_redacts_workspace_paths_in_text_and_tool_args() {
let text = super::sanitize_windsurf_text(