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

refactor: 优化调度候选排序与用量写入链路并改进 Fernet 缓存与前端批量列表

Browse Source
This commit is contained in:
fawney19
2026-04-21 16:19:07 +08:00
parent c5c56ff92f
commit 25a2b417be
77 changed files with 4801 additions and 2881 deletions
Download Patch File Download Diff File Expand all files Collapse all files

425
apps/aether-gateway/src/execution_runtime/stream/execution.rs
View File

@@ -74,6 +74,7 @@ use crate::orchestration::{
};
use crate::request_candidate_runtime::{
ensure_execution_request_candidate_slot, record_local_request_candidate_status,
record_local_request_candidate_status_snapshot, snapshot_local_request_candidate_status,
};
use crate::usage::submit_stream_report;
use crate::usage::{GatewayStreamReportRequest, GatewaySyncReportRequest};
@@ -89,7 +90,30 @@ fn record_sync_terminal_usage(
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
state
.usage_runtime
.record_sync_terminal(state.data.as_ref(), &context_seed, &payload_seed);
.record_sync_terminal(state.data.as_ref(), context_seed, payload_seed);
}
fn build_stream_sync_payload(
trace_id: &str,
report_kind: String,
report_context: Option<Value>,
status_code: u16,
headers: BTreeMap<String, String>,
body_json: Option<Value>,
body_base64: Option<String>,
telemetry: Option<ExecutionTelemetry>,
) -> GatewaySyncReportRequest {
GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind,
report_context,
status_code,
headers,
body_json,
client_body_json: None,
body_base64,
telemetry,
}
}
fn record_stream_terminal_usage(
@@ -103,12 +127,61 @@ fn record_stream_terminal_usage(
let payload_seed = build_stream_terminal_usage_payload_seed(payload);
state.usage_runtime.record_stream_terminal(
state.data.as_ref(),
&context_seed,
&payload_seed,
context_seed,
payload_seed,
cancelled,
);
}
fn build_stream_body_capture(
body: &[u8],
truncated: bool,
) -> (Option<String>, Option<UsageBodyCaptureState>) {
let body_base64 =
(!body.is_empty()).then(|| base64::engine::general_purpose::STANDARD.encode(body));
let body_state = Some(if truncated {
UsageBodyCaptureState::Truncated
} else if body.is_empty() {
UsageBodyCaptureState::None
} else {
UsageBodyCaptureState::Inline
});
(body_base64, body_state)
}
#[allow(clippy::too_many_arguments)] // stream report payload assembly mirrors runtime state
fn build_stream_usage_payload(
trace_id: String,
report_kind: String,
report_context: Option<Value>,
status_code: u16,
headers: BTreeMap<String, String>,
provider_body: &[u8],
provider_body_truncated: bool,
client_body: &[u8],
client_body_truncated: bool,
terminal_summary: Option<ExecutionStreamTerminalSummary>,
telemetry: Option<ExecutionTelemetry>,
) -> GatewayStreamReportRequest {
let (provider_body_base64, provider_body_state) =
build_stream_body_capture(provider_body, provider_body_truncated);
let (client_body_base64, client_body_state) =
build_stream_body_capture(client_body, client_body_truncated);
GatewayStreamReportRequest {
trace_id,
report_kind,
report_context,
status_code,
headers,
provider_body_base64,
provider_body_state,
client_body_base64,
client_body_state,
terminal_summary,
telemetry,
}
}
fn append_stream_capture_bytes(
buffer: &mut Vec<u8>,
chunk: &[u8],
@@ -138,9 +211,7 @@ async fn execute_in_process_stream(
return Ok(execution);
}
DirectSyncExecutionRuntime::new()
.execute_stream(plan.clone())
.await
DirectSyncExecutionRuntime::new().execute_stream(plan).await
}
#[allow(clippy::too_many_arguments)] // internal function, grouping would add unnecessary indirection
@@ -155,19 +226,18 @@ pub(crate) async fn execute_execution_runtime_stream(
) -> Result<Option<Response<Body>>, GatewayError> {
ensure_execution_request_candidate_slot(state, &mut plan, &mut report_context).await;
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
let request_candidate_status_snapshot =
snapshot_local_request_candidate_status(&plan, report_context.as_ref());
state
.usage_runtime
.record_pending(state.data.as_ref(), &lifecycle_seed);
.record_pending(state.data.as_ref(), lifecycle_seed.clone());
let candidate_started_unix_secs = current_request_candidate_unix_ms();
{
if let Some(snapshot) = request_candidate_status_snapshot.clone() {
let state_bg = state.clone();
let plan_bg = plan.clone();
let report_context_bg = report_context.clone();
tokio::spawn(async move {
record_local_request_candidate_status(
record_local_request_candidate_status_snapshot(
&state_bg,
&plan_bg,
report_context_bg.as_ref(),
&snapshot,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Pending,
status_code: None,
@@ -412,8 +482,9 @@ fn response_headers_indicate_sse(headers: &BTreeMap<String, String>) -> bool {
}
fn encode_terminal_sse_error_event(failure: &StreamFailureReport) -> Result<Bytes, std::io::Error> {
let payload =
serde_json::to_string(&failure.body_json).map_err(|err| IoError::other(err.to_string()))?;
let payload = failure
.to_json_string()
.map_err(|err| IoError::other(err.to_string()))?;
let mut event = String::from("event: aether.error\n");
for line in payload.lines() {
event.push_str("data: ");
@@ -527,6 +598,8 @@ async fn execute_stream_from_frame_stream(
let provider_name = plan.provider_name.as_deref().unwrap_or("-");
let model_name = plan.model_name.as_deref().unwrap_or("-");
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
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())
.and_then(|context| context.candidate_index)
.map(|value| value.to_string())
@@ -756,27 +829,26 @@ async fn execute_stream_from_frame_stream(
return Ok(None);
}
let usage_report_kind = stream_error_finalize_kind
.clone()
.or_else(|| report_kind.clone())
.unwrap_or_default();
let usage_payload = GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind: usage_report_kind,
report_context: report_context.clone(),
let payload = build_stream_sync_payload(
trace_id,
stream_error_finalize_kind
.as_deref()
.or(report_kind.as_deref())
.unwrap_or_default()
.to_string(),
report_context,
status_code,
headers: headers.clone(),
body_json: body_json.clone(),
client_body_json: None,
body_base64: body_base64.clone(),
telemetry: None,
};
record_sync_terminal_usage(state, &plan, report_context.as_ref(), &usage_payload);
headers,
body_json,
body_base64,
None,
);
record_sync_terminal_usage(state, &plan, payload.report_context.as_ref(), &payload);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_local_request_candidate_status(
state,
&plan,
report_context.as_ref(),
payload.report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(status_code),
@@ -790,18 +862,7 @@ async fn execute_stream_from_frame_stream(
},
)
.await;
if let Some(report_kind) = stream_error_finalize_kind {
let payload = GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind,
report_context,
status_code,
headers: headers.clone(),
body_json,
client_body_json: None,
body_base64,
telemetry: None,
};
if stream_error_finalize_kind.is_some() {
let response =
submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload)
.await?;
@@ -817,7 +878,7 @@ async fn execute_stream_from_frame_stream(
decision,
plan_kind,
status_code,
headers,
payload.headers,
error_body,
)?,
Some(request_id),
@@ -891,12 +952,12 @@ async fn execute_stream_from_frame_stream(
trace_id,
decision,
&plan,
report_context.clone(),
report_context,
request_id,
candidate_id,
report_kind,
&headers,
prefetched_telemetry.clone(),
headers,
prefetched_telemetry,
&provider_prefetched_body,
failure,
)
@@ -924,12 +985,12 @@ async fn execute_stream_from_frame_stream(
trace_id,
decision,
&plan,
report_context.clone(),
report_context,
request_id,
candidate_id,
report_kind,
&headers,
prefetched_telemetry.clone(),
headers,
prefetched_telemetry,
&prefetched_body,
failure,
)
@@ -964,21 +1025,20 @@ async fn execute_stream_from_frame_stream(
provider_prefetched_body_bytes = provider_prefetched_body.len(),
"gateway detected embedded error while prefetching execution runtime stream"
);
let payload = GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind: report_kind.clone(),
report_context: report_context.clone(),
let payload = build_stream_sync_payload(
trace_id,
report_kind.clone(),
report_context,
status_code,
headers: headers.clone(),
body_json: Some(body_json),
client_body_json: None,
body_base64: None,
telemetry: prefetched_telemetry.clone(),
};
headers,
Some(body_json),
None,
prefetched_telemetry,
);
record_sync_terminal_usage(
state,
&plan,
report_context.as_ref(),
payload.report_context.as_ref(),
&payload,
);
let response = submit_local_core_error_or_sync_finalize(
@@ -1013,12 +1073,12 @@ async fn execute_stream_from_frame_stream(
trace_id,
decision,
&plan,
report_context.clone(),
report_context,
request_id,
candidate_id,
report_kind,
&headers,
prefetched_telemetry.clone(),
headers,
prefetched_telemetry,
&provider_prefetched_body,
failure,
)
@@ -1044,12 +1104,12 @@ async fn execute_stream_from_frame_stream(
trace_id,
decision,
&plan,
report_context.clone(),
report_context,
request_id,
candidate_id,
report_kind,
&headers,
prefetched_telemetry.clone(),
headers,
prefetched_telemetry,
&provider_prefetched_body,
failure,
)
@@ -1093,12 +1153,12 @@ async fn execute_stream_from_frame_stream(
trace_id,
decision,
&plan,
report_context.clone(),
report_context,
request_id,
candidate_id,
report_kind,
&headers,
prefetched_telemetry.clone(),
headers,
prefetched_telemetry,
&provider_prefetched_body,
build_stream_failure_from_execution_error(&error),
)
@@ -1108,6 +1168,8 @@ async fn execute_stream_from_frame_stream(
}
}
}
drop(private_stream_normalizer);
drop(local_stream_rewriter);
state.usage_runtime.record_stream_started(
state.data.as_ref(),
@@ -1115,18 +1177,15 @@ async fn execute_stream_from_frame_stream(
status_code,
prefetched_telemetry.as_ref(),
);
{
if let Some(snapshot) = request_candidate_status_snapshot {
let state_bg = state.clone();
let plan_bg = plan.clone();
let report_context_bg = report_context.clone();
let latency_ms = prefetched_telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms);
tokio::spawn(async move {
record_local_request_candidate_status(
record_local_request_candidate_status_snapshot(
&state_bg,
&plan_bg,
report_context_bg.as_ref(),
&snapshot,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Streaming,
status_code: Some(status_code),
@@ -1141,24 +1200,27 @@ async fn execute_stream_from_frame_stream(
});
}
let request_id = request_id.to_string();
let candidate_id = candidate_id.map(ToOwned::to_owned);
let (tx, mut rx) = mpsc::channel::<Result<Bytes, IoError>>(16);
let state_for_report = state.clone();
let plan_for_report = plan.clone();
let plan_for_report = plan;
let trace_id_owned = trace_id.to_string();
let headers_for_report = headers.clone();
let report_kind_owned = report_kind.clone();
let report_context_owned = report_context.clone();
let lifecycle_seed_for_report = lifecycle_seed.clone();
let report_kind_owned = report_kind;
let report_context_owned = report_context;
let normalized_stream_report_context_owned = normalized_stream_report_context;
let lifecycle_seed_for_report = lifecycle_seed;
let provider_prefetched_body_for_report = provider_prefetched_body;
let prefetched_body_for_report = prefetched_body;
let prefetched_chunks_for_body = prefetched_chunks;
let initial_telemetry = prefetched_telemetry.clone();
let initial_telemetry = prefetched_telemetry;
let initial_reached_eof = reached_eof;
let direct_stream_finalize_kind_owned = direct_stream_finalize_kind.clone();
let direct_stream_finalize_kind_owned = direct_stream_finalize_kind;
let candidate_started_unix_secs_for_report = candidate_started_unix_secs;
let request_id_for_report = request_id.to_string();
let request_id_for_report_log = short_request_id(request_id);
let candidate_id_for_report = candidate_id.map(ToOwned::to_owned);
let request_id_for_report = request_id.clone();
let request_id_for_report_log = short_request_id(&request_id);
let candidate_id_for_report = candidate_id.clone();
let emit_passthrough_sse_terminal_error =
skip_direct_finalize_prefetch && response_headers_indicate_sse(&headers);
let body_capture_policy = match UsageRuntimeAccess::body_capture_policy(state.data.as_ref())
@@ -1194,6 +1256,10 @@ async fn execute_stream_from_frame_stream(
let mut buffered_body = Vec::new();
let mut provider_body_truncated = false;
let mut client_body_truncated = false;
let mut private_stream_normalizer =
maybe_build_provider_private_stream_normalizer(report_context_owned.as_ref());
let mut local_stream_rewriter =
maybe_build_stream_response_rewriter(normalized_stream_report_context_owned.as_ref());
append_stream_capture_bytes(
&mut provider_buffered_body,
&provider_prefetched_body_for_report,
@@ -1206,13 +1272,68 @@ async fn execute_stream_from_frame_stream(
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
let mut telemetry: Option<ExecutionTelemetry> = initial_telemetry.clone();
let mut usage_stream_telemetry: Option<ExecutionTelemetry> = initial_telemetry;
let mut usage_stream_telemetry: Option<ExecutionTelemetry> = initial_telemetry.clone();
let mut telemetry: Option<ExecutionTelemetry> = initial_telemetry;
let reached_eof = initial_reached_eof;
let mut downstream_dropped = false;
let mut terminal_failure: Option<StreamFailureReport> = None;
if !provider_prefetched_body_for_report.is_empty() {
let normalized_prefetched_chunk = if let Some(normalizer) =
private_stream_normalizer.as_mut()
{
match normalizer.push_chunk(&provider_prefetched_body_for_report) {
Ok(normalized_chunk) => Some(normalized_chunk),
Err(err) => {
warn!(
event_name = "stream_execution_prefetch_normalize_restore_failed",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
error = ?err,
"gateway failed to restore private stream normalization state after prefetch"
);
terminal_failure = Some(build_stream_failure_report(
"execution_runtime_stream_rewrite_error",
format!(
"failed to restore private stream normalization state after prefetch: {err:?}"
),
502,
));
None
}
}
} else {
None
};
let replay_chunk = normalized_prefetched_chunk
.as_deref()
.unwrap_or(provider_prefetched_body_for_report.as_slice());
if terminal_failure.is_none() {
if let Some(rewriter) = local_stream_rewriter.as_mut() {
if let Err(err) = rewriter.push_chunk(replay_chunk) {
warn!(
event_name = "stream_execution_prefetch_rewrite_restore_failed",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
error = ?err,
"gateway failed to restore local stream rewrite state after prefetch"
);
terminal_failure = Some(build_stream_failure_report(
"execution_runtime_stream_rewrite_error",
format!(
"failed to restore local stream rewrite state after prefetch: {err:?}"
),
502,
));
}
}
}
}
if !reached_eof {
if terminal_failure.is_none() && !reached_eof {
loop {
let next_frame = match next_stream_frame(&mut buffered_frames, &mut lines).await {
Ok(frame) => frame,
@@ -1355,7 +1476,6 @@ async fn execute_stream_from_frame_stream(
usage_stream_telemetry.as_ref(),
&frame_telemetry,
);
telemetry = Some(frame_telemetry.clone());
if should_refresh_stream_usage {
state_for_report.usage_runtime.record_stream_started(
state_for_report.data.as_ref(),
@@ -1363,8 +1483,9 @@ async fn execute_stream_from_frame_stream(
status_code,
Some(&frame_telemetry),
);
usage_stream_telemetry = Some(frame_telemetry);
usage_stream_telemetry = Some(frame_telemetry.clone());
}
telemetry = Some(frame_telemetry);
}
StreamFramePayload::Eof { summary } => {
stream_terminal_summary = summary;
@@ -1564,50 +1685,39 @@ async fn execute_stream_from_frame_stream(
trace_id = %trace_id_owned,
"gateway skipped stream report because downstream disconnected before completion"
);
let usage_payload = build_stream_usage_payload(
trace_id_owned,
report_kind_owned.unwrap_or_default(),
report_context_owned,
499,
headers_for_report,
&provider_buffered_body,
provider_body_truncated,
&buffered_body,
client_body_truncated,
stream_terminal_summary,
telemetry,
);
record_stream_terminal_usage(
&state_for_report,
&plan_for_report,
report_context_owned.as_ref(),
&GatewayStreamReportRequest {
trace_id: trace_id_owned.clone(),
report_kind: report_kind_owned.clone().unwrap_or_default(),
report_context: report_context_owned.clone(),
status_code: 499,
headers: headers_for_report.clone(),
provider_body_base64: (!provider_buffered_body.is_empty()).then(|| {
base64::engine::general_purpose::STANDARD.encode(&provider_buffered_body)
}),
provider_body_state: Some(if provider_body_truncated {
UsageBodyCaptureState::Truncated
} else if provider_buffered_body.is_empty() {
UsageBodyCaptureState::None
} else {
UsageBodyCaptureState::Inline
}),
client_body_base64: (!buffered_body.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(&buffered_body)),
client_body_state: Some(if client_body_truncated {
UsageBodyCaptureState::Truncated
} else if buffered_body.is_empty() {
UsageBodyCaptureState::None
} else {
UsageBodyCaptureState::Inline
}),
terminal_summary: stream_terminal_summary.clone(),
telemetry: telemetry.clone(),
},
usage_payload.report_context.as_ref(),
&usage_payload,
true,
);
record_local_request_candidate_status(
&state_for_report,
&plan_for_report,
report_context_owned.as_ref(),
usage_payload.report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Cancelled,
status_code: Some(499),
error_type: Some("downstream_disconnect".to_string()),
error_message: Some("client disconnected before stream completion".to_string()),
latency_ms: telemetry.as_ref().and_then(|value| value.elapsed_ms),
latency_ms: usage_payload
.telemetry
.as_ref()
.and_then(|value| value.elapsed_ms),
started_at_unix_ms: Some(candidate_started_unix_secs_for_report),
finished_at_unix_ms: Some(current_request_candidate_unix_ms()),
},
@@ -1622,9 +1732,9 @@ async fn execute_stream_from_frame_stream(
&trace_id_owned,
&plan_for_report,
direct_stream_finalize_kind_owned.as_deref(),
report_context_owned.as_ref(),
&headers_for_report,
telemetry.clone(),
report_context_owned,
headers_for_report,
telemetry,
&provider_buffered_body,
candidate_started_unix_secs_for_report,
failure,
@@ -1633,38 +1743,25 @@ async fn execute_stream_from_frame_stream(
return;
}
let usage_payload = GatewayStreamReportRequest {
trace_id: trace_id_owned.clone(),
report_kind: report_kind_owned.clone().unwrap_or_default(),
report_context: report_context_owned.clone(),
let should_submit_report = report_kind_owned.is_some();
let usage_payload = build_stream_usage_payload(
trace_id_owned.clone(),
report_kind_owned.unwrap_or_default(),
report_context_owned,
status_code,
headers: headers_for_report.clone(),
provider_body_base64: (!provider_buffered_body.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(&provider_buffered_body)),
provider_body_state: Some(if provider_body_truncated {
UsageBodyCaptureState::Truncated
} else if provider_buffered_body.is_empty() {
UsageBodyCaptureState::None
} else {
UsageBodyCaptureState::Inline
}),
client_body_base64: (!buffered_body.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(&buffered_body)),
client_body_state: Some(if client_body_truncated {
UsageBodyCaptureState::Truncated
} else if buffered_body.is_empty() {
UsageBodyCaptureState::None
} else {
UsageBodyCaptureState::Inline
}),
terminal_summary: stream_terminal_summary,
telemetry: telemetry.clone(),
};
headers_for_report,
&provider_buffered_body,
provider_body_truncated,
&buffered_body,
client_body_truncated,
stream_terminal_summary,
telemetry,
);
apply_local_execution_effect(
&state_for_report,
LocalExecutionEffectContext {
plan: &plan_for_report,
report_context: report_context_owned.as_ref(),
report_context: usage_payload.report_context.as_ref(),
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
@@ -1673,7 +1770,7 @@ async fn execute_stream_from_frame_stream(
&state_for_report,
LocalExecutionEffectContext {
plan: &plan_for_report,
report_context: report_context_owned.as_ref(),
report_context: usage_payload.report_context.as_ref(),
},
LocalExecutionEffect::AdaptiveSuccess(LocalAdaptiveSuccessEffect),
)
@@ -1682,7 +1779,7 @@ async fn execute_stream_from_frame_stream(
&state_for_report,
LocalExecutionEffectContext {
plan: &plan_for_report,
report_context: report_context_owned.as_ref(),
report_context: usage_payload.report_context.as_ref(),
},
LocalExecutionEffect::PoolSuccessStream {
payload: &usage_payload,
@@ -1692,31 +1789,31 @@ async fn execute_stream_from_frame_stream(
record_stream_terminal_usage(
&state_for_report,
&plan_for_report,
report_context_owned.as_ref(),
usage_payload.report_context.as_ref(),
&usage_payload,
false,
);
record_local_request_candidate_status(
&state_for_report,
&plan_for_report,
report_context_owned.as_ref(),
usage_payload.report_context.as_ref(),
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Success,
status_code: Some(status_code),
error_type: None,
error_message: None,
latency_ms: telemetry.as_ref().and_then(|value| value.elapsed_ms),
latency_ms: usage_payload
.telemetry
.as_ref()
.and_then(|value| value.elapsed_ms),
started_at_unix_ms: Some(candidate_started_unix_secs_for_report),
finished_at_unix_ms: Some(current_request_candidate_unix_ms()),
},
)
.await;
if let Some(report_kind) = report_kind_owned {
let mut report = usage_payload;
report.report_kind = report_kind;
if let Err(err) = submit_stream_report(&state_for_report, &trace_id_owned, report).await
{
if should_submit_report {
if let Err(err) = submit_stream_report(&state_for_report, usage_payload).await {
warn!(
event_name = "execution_report_submit_failed",
log_type = "ops",
@@ -1740,12 +1837,10 @@ async fn execute_stream_from_frame_stream(
}
};
headers.insert(
CONTROL_REQUEST_ID_HEADER.to_string(),
request_id.to_string(),
);
headers.insert(CONTROL_REQUEST_ID_HEADER.to_string(), request_id.clone());
if let Some(candidate_id) = candidate_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{

159
apps/aether-gateway/src/execution_runtime/stream/execution_failures.rs
View File

@@ -7,6 +7,7 @@ use aether_usage_runtime::{
use axum::body::Body;
use axum::http::Response;
use base64::Engine as _;
use serde::Serialize;
use serde_json::{Map, Value};
use tracing::warn;
@@ -32,7 +33,51 @@ pub(super) struct StreamFailureReport {
pub(super) status_code: u16,
pub(super) error_type: String,
pub(super) error_message: String,
pub(super) body_json: Value,
extra_error_fields: Map<String, Value>,
}
#[derive(Serialize)]
struct StreamFailureBody<'a> {
error: StreamFailureBodyFields<'a>,
}
#[derive(Serialize)]
struct StreamFailureBodyFields<'a> {
#[serde(rename = "type")]
error_type: &'a str,
message: &'a str,
code: u16,
#[serde(flatten)]
extra_error_fields: &'a Map<String, Value>,
}
impl StreamFailureReport {
fn into_body_json(self) -> Value {
let Self {
status_code,
error_type,
error_message,
mut extra_error_fields,
} = self;
extra_error_fields.insert("type".to_string(), Value::String(error_type));
extra_error_fields.insert("message".to_string(), Value::String(error_message));
extra_error_fields.insert("code".to_string(), Value::from(status_code));
Value::Object(Map::from_iter([(
"error".to_string(),
Value::Object(extra_error_fields),
)]))
}
pub(super) fn to_json_string(&self) -> serde_json::Result<String> {
serde_json::to_string(&StreamFailureBody {
error: StreamFailureBodyFields {
error_type: self.error_type.as_str(),
message: self.error_message.as_str(),
code: self.status_code,
extra_error_fields: &self.extra_error_fields,
},
})
}
}
pub(super) fn build_stream_failure_report(
@@ -44,16 +89,9 @@ pub(super) fn build_stream_failure_report(
let error_message = error_message.into();
StreamFailureReport {
status_code,
body_json: Value::Object(Map::from_iter([(
"error".to_string(),
Value::Object(Map::from_iter([
("type".to_string(), Value::String(error_type.clone())),
("message".to_string(), Value::String(error_message.clone())),
("code".to_string(), Value::from(status_code)),
])),
)])),
error_type,
error_message,
extra_error_fields: Map::new(),
}
}
@@ -65,11 +103,9 @@ pub(super) fn build_stream_failure_from_execution_error(
.ok()
.and_then(|value| value.as_str().map(ToOwned::to_owned))
.unwrap_or_else(|| "internal".to_string());
let error_message = error.message.trim().to_string();
let phase = serde_json::to_value(&error.phase).unwrap_or(Value::Null);
let mut error_object = Map::from_iter([
("type".to_string(), Value::String(error_type.clone())),
("message".to_string(), Value::String(error.message.clone())),
("code".to_string(), Value::from(status_code)),
("phase".to_string(), phase),
("retryable".to_string(), Value::Bool(error.retryable)),
(
@@ -84,11 +120,8 @@ pub(super) fn build_stream_failure_from_execution_error(
StreamFailureReport {
status_code,
error_type,
error_message: error.message.trim().to_string(),
body_json: Value::Object(Map::from_iter([(
"error".to_string(),
Value::Object(error_object),
)])),
error_message,
extra_error_fields: error_object,
}
}
@@ -96,23 +129,23 @@ fn build_stream_failure_sync_payload(
trace_id: &str,
report_kind: String,
report_context: Option<Value>,
headers: &std::collections::BTreeMap<String, String>,
mut headers: std::collections::BTreeMap<String, String>,
telemetry: Option<ExecutionTelemetry>,
provider_buffered_body: &[u8],
failure: &StreamFailureReport,
failure: StreamFailureReport,
) -> GatewaySyncReportRequest {
let mut response_headers = headers.clone();
response_headers.remove("content-encoding");
response_headers.remove("content-length");
response_headers.insert("content-type".to_string(), "application/json".to_string());
let status_code = failure.status_code;
headers.remove("content-encoding");
headers.remove("content-length");
headers.insert("content-type".to_string(), "application/json".to_string());
GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind,
report_context,
status_code: failure.status_code,
headers: response_headers,
body_json: Some(failure.body_json.clone()),
status_code,
headers,
body_json: Some(failure.into_body_json()),
client_body_json: None,
body_base64: (!provider_buffered_body.is_empty())
.then(|| base64::engine::general_purpose::STANDARD.encode(provider_buffered_body)),
@@ -120,27 +153,40 @@ fn build_stream_failure_sync_payload(
}
}
fn stream_failure_body_field<'a>(
payload: &'a GatewaySyncReportRequest,
field: &str,
) -> Option<&'a str> {
payload
.body_json
.as_ref()
.and_then(|body_json| body_json.get("error"))
.and_then(|value| value.get(field))
.and_then(Value::as_str)
}
async fn record_stream_sync_failure(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&Value>,
payload: &GatewaySyncReportRequest,
failure: &StreamFailureReport,
started_at_unix_ms: Option<u64>,
) {
let error_body = serde_json::to_string(&failure.body_json).ok();
let error_type = stream_failure_body_field(payload, "type").unwrap_or("internal");
let error_message = stream_failure_body_field(payload, "message").unwrap_or_default();
let error_body = payload
.body_json
.as_ref()
.and_then(|body_json| serde_json::to_string(body_json).ok());
let failure_analysis = resolve_local_failover_analysis_for_attempt(
state,
plan,
report_context,
failure.status_code,
payload.status_code,
error_body.as_deref(),
)
.await;
if matches!(
failure.error_type.as_str(),
"first_byte_timeout" | "read_timeout"
) {
if matches!(error_type, "first_byte_timeout" | "read_timeout") {
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
@@ -158,7 +204,7 @@ async fn record_stream_sync_failure(
report_context,
},
LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
status_code: failure.status_code,
status_code: payload.status_code,
classification: failure_analysis.classification,
}),
)
@@ -170,7 +216,7 @@ async fn record_stream_sync_failure(
report_context,
},
LocalExecutionEffect::AdaptiveRateLimit(LocalAdaptiveRateLimitEffect {
status_code: failure.status_code,
status_code: payload.status_code,
classification: failure_analysis.classification,
headers: Some(&payload.headers),
}),
@@ -183,7 +229,7 @@ async fn record_stream_sync_failure(
report_context,
},
LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
status_code: failure.status_code,
status_code: payload.status_code,
classification: failure_analysis.classification,
}),
)
@@ -195,7 +241,7 @@ async fn record_stream_sync_failure(
report_context,
},
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
status_code: failure.status_code,
status_code: payload.status_code,
response_text: error_body.as_deref(),
}),
)
@@ -207,7 +253,7 @@ async fn record_stream_sync_failure(
report_context,
},
LocalExecutionEffect::PoolError(LocalPoolErrorEffect {
status_code: failure.status_code,
status_code: payload.status_code,
classification: failure_analysis.classification,
headers: &payload.headers,
error_body: error_body.as_deref(),
@@ -218,16 +264,16 @@ async fn record_stream_sync_failure(
let payload_seed = build_sync_terminal_usage_payload_seed(payload);
state
.usage_runtime
.record_sync_terminal(state.data.as_ref(), &context_seed, &payload_seed);
.record_sync_terminal(state.data.as_ref(), context_seed, payload_seed);
let terminal_unix_secs = current_request_candidate_unix_ms();
record_report_request_candidate_status(
state,
report_context,
SchedulerRequestCandidateStatusUpdate {
status: RequestCandidateStatus::Failed,
status_code: Some(failure.status_code),
error_type: Some(failure.error_type.clone()),
error_message: Some(failure.error_message.clone()),
status_code: Some(payload.status_code),
error_type: Some(error_type.to_string()),
error_message: Some(error_message.to_string()),
latency_ms: payload
.telemetry
.as_ref()
@@ -249,7 +295,7 @@ pub(super) async fn handle_prefetch_stream_failure(
request_id: &str,
candidate_id: Option<&str>,
report_kind: &str,
headers: &std::collections::BTreeMap<String, String>,
headers: std::collections::BTreeMap<String, String>,
telemetry: Option<ExecutionTelemetry>,
buffered_body: &[u8],
failure: StreamFailureReport,
@@ -257,21 +303,13 @@ pub(super) async fn handle_prefetch_stream_failure(
let payload = build_stream_failure_sync_payload(
trace_id,
report_kind.to_string(),
report_context.clone(),
report_context,
headers,
telemetry,
buffered_body,
&failure,
failure,
);
record_stream_sync_failure(
state,
plan,
report_context.as_ref(),
&payload,
&failure,
None,
)
.await;
record_stream_sync_failure(state, plan, payload.report_context.as_ref(), &payload, None).await;
let response =
submit_local_core_error_or_sync_finalize(state, trace_id, decision, payload).await?;
@@ -287,8 +325,8 @@ pub(super) async fn submit_midstream_stream_failure(
trace_id: &str,
plan: &ExecutionPlan,
direct_stream_finalize_kind: Option<&str>,
report_context: Option<&Value>,
headers: &std::collections::BTreeMap<String, String>,
report_context: Option<Value>,
headers: std::collections::BTreeMap<String, String>,
telemetry: Option<ExecutionTelemetry>,
buffered_body: &[u8],
started_at_unix_ms: u64,
@@ -303,22 +341,21 @@ pub(super) async fn submit_midstream_stream_failure(
let payload = build_stream_failure_sync_payload(
trace_id,
report_kind,
report_context.cloned(),
report_context,
headers,
telemetry,
buffered_body,
&failure,
failure,
);
record_stream_sync_failure(
state,
plan,
report_context,
payload.report_context.as_ref(),
&payload,
&failure,
Some(started_at_unix_ms),
)
.await;
if let Err(err) = submit_sync_report(state, trace_id, payload).await {
if let Err(err) = submit_sync_report(state, payload).await {
let request_id = short_request_id(plan.request_id.as_str());
warn!(
event_name = "execution_report_submit_failed",