i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-01 17:00:21 +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

3
apps/aether-gateway/src/execution_runtime/mod.rs
View File

@@ -42,7 +42,8 @@ pub(crate) use sync::{
execute_execution_runtime_sync, maybe_build_local_sync_finalize_response,
maybe_build_local_video_error_response, maybe_build_local_video_success_outcome,
resolve_local_sync_error_background_report_kind,
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessOutcome,
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessBuild,
LocalVideoSyncSuccessOutcome,
};
pub(crate) use transport::{
execute_sync_plan as execute_execution_runtime_sync_plan, DirectSyncExecutionRuntime,

4
apps/aether-gateway/src/execution_runtime/server.rs
View File

@@ -221,7 +221,7 @@ async fn execute_sync(
let plan = parse_request_json::<ExecutionPlan>(request).await?;
let result = state
.execution_runtime
.execute_sync(plan)
.execute_sync(&plan)
.await
.map_err(|err| ExecutionRuntimeAppError(ExecutionRuntimeServerError::Transport(err)))?;
Ok(maybe_hold_axum_response_permit(
@@ -238,7 +238,7 @@ async fn execute_stream(
let plan = parse_request_json::<ExecutionPlan>(request).await?;
let execution = state
.execution_runtime
.execute_stream(plan)
.execute_stream(&plan)
.await
.map_err(|err| ExecutionRuntimeAppError(ExecutionRuntimeServerError::Transport(err)))?;

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",

8
apps/aether-gateway/src/execution_runtime/stream_pump.rs
View File

@@ -274,7 +274,7 @@ fn format_error_chain(err: &(dyn std::error::Error + 'static)) -> String {
fn observe_stream_chunk(
observer: &mut StreamingStandardTerminalObserver,
report_context: &Value,
private_stream_normalizer: Option<&mut crate::ai_pipeline::ProviderPrivateStreamNormalizer>,
private_stream_normalizer: Option<&mut crate::ai_pipeline::ProviderPrivateStreamNormalizer<'_>>,
observer_buffered: &mut Vec<u8>,
chunk: &[u8],
) {
@@ -298,7 +298,7 @@ fn observe_stream_chunk(
fn finalize_stream_terminal_summary(
observer: &mut StreamingStandardTerminalObserver,
report_context: &Value,
private_stream_normalizer: Option<&mut crate::ai_pipeline::ProviderPrivateStreamNormalizer>,
private_stream_normalizer: Option<&mut crate::ai_pipeline::ProviderPrivateStreamNormalizer<'_>>,
observer_buffered: &mut Vec<u8>,
) -> Option<ExecutionStreamTerminalSummary> {
if let Some(normalizer) = private_stream_normalizer {
@@ -412,7 +412,7 @@ mod tests {
});
let execution = DirectSyncExecutionRuntime::new()
.execute_stream(ExecutionPlan {
.execute_stream(&ExecutionPlan {
request_id: "req-stream-ttfb-1".into(),
candidate_id: Some("cand-stream-ttfb-1".into()),
provider_name: Some("openai".into()),
@@ -499,7 +499,7 @@ mod tests {
let runtime = DirectSyncExecutionRuntime::new();
let execution = runtime
.execute_stream(ExecutionPlan {
.execute_stream(&ExecutionPlan {
request_id: "req-telemetry-order".to_string(),
candidate_id: Some("cand-telemetry-order".to_string()),
provider_name: Some("OpenAI".to_string()),

2
apps/aether-gateway/src/execution_runtime/submission.rs
View File

@@ -545,7 +545,7 @@ pub(crate) async fn submit_local_core_error_or_sync_finalize(
{
let mut report_payload = payload.clone();
report_payload.report_kind = error_report_kind;
spawn_sync_report(state.clone(), trace_id.to_string(), report_payload);
spawn_sync_report(state.clone(), report_payload);
} else {
warn!(
event_name = "local_core_finalize_missing_error_report_mapping",

484
apps/aether-gateway/src/execution_runtime/sync/execution.rs
View File

@@ -23,7 +23,6 @@ use crate::api::response::{
attach_control_metadata_headers, build_client_response, build_client_response_from_parts,
};
use crate::clock::current_unix_ms as current_request_candidate_unix_ms;
use crate::constants::{CONTROL_CANDIDATE_ID_HEADER, CONTROL_REQUEST_ID_HEADER};
use crate::control::GatewayControlDecision;
#[cfg(test)]
use crate::execution_runtime::remote_compat::post_sync_plan_to_remote_execution_runtime;
@@ -57,7 +56,8 @@ use policy::decode_execution_result_body;
pub(crate) use response::{
maybe_build_local_sync_finalize_response, maybe_build_local_video_error_response,
maybe_build_local_video_success_outcome, resolve_local_sync_error_background_report_kind,
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessOutcome,
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessBuild,
LocalVideoSyncSuccessOutcome,
};
struct ImplicitSyncFinalizeOutcome {
@@ -75,7 +75,65 @@ 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_sync_report_payload(
trace_id: &str,
report_kind: String,
report_context: Option<serde_json::Value>,
status_code: u16,
headers: BTreeMap<String, String>,
body_json: Option<serde_json::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,
}
}
async fn apply_sync_success_effects(
state: &AppState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
payload: &GatewaySyncReportRequest,
) {
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan,
report_context,
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan,
report_context,
},
LocalExecutionEffect::AdaptiveSuccess(LocalAdaptiveSuccessEffect),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan,
report_context,
},
LocalExecutionEffect::PoolSuccessSync { payload },
)
.await;
}
#[cfg(test)]
@@ -112,7 +170,7 @@ pub(crate) async fn execute_execution_runtime_sync(
let lifecycle_seed = build_lifecycle_usage_seed(&plan, report_context.as_ref());
state
.usage_runtime
.record_pending(state.data.as_ref(), &lifecycle_seed);
.record_pending(state.data.as_ref(), lifecycle_seed);
record_local_request_candidate_status(
state,
&plan,
@@ -129,11 +187,8 @@ pub(crate) async fn execute_execution_runtime_sync(
)
.await;
#[cfg(not(test))]
let result = {
match DirectSyncExecutionRuntime::new()
.execute_sync(plan.clone())
.await
{
let mut result = {
match DirectSyncExecutionRuntime::new().execute_sync(&plan).await {
Ok(result) => result,
Err(err) => {
warn!(
@@ -171,7 +226,7 @@ pub(crate) async fn execute_execution_runtime_sync(
}
};
#[cfg(test)]
let result = {
let mut result = {
if let Some(override_fn) = state.execution_runtime_sync_override.as_ref() {
match (override_fn.0)(&plan) {
Ok(result) => result,
@@ -215,10 +270,7 @@ pub(crate) async fn execute_execution_runtime_sync(
.trim()
.is_empty()
{
match DirectSyncExecutionRuntime::new()
.execute_sync(plan.clone())
.await
{
match DirectSyncExecutionRuntime::new().execute_sync(&plan).await {
Ok(result) => result,
Err(err) => {
warn!(
@@ -287,8 +339,9 @@ pub(crate) async fn execute_execution_runtime_sync(
.telemetry
.as_ref()
.and_then(|telemetry| telemetry.elapsed_ms);
let mut headers = result.headers.clone();
let (body_bytes, body_json, body_base64) = decode_execution_result_body(&result, &mut headers)?;
let mut headers = std::mem::take(&mut result.headers);
let (body_bytes, body_json, body_base64) =
decode_execution_result_body(result.body.take(), &mut headers)?;
let local_failover_response_text = local_failover_response_text(
body_json.as_ref(),
&body_bytes,
@@ -403,38 +456,26 @@ pub(crate) async fn execute_execution_runtime_sync(
);
return Ok(None);
}
let request_id = (!result.request_id.trim().is_empty())
.then_some(result.request_id.as_str())
.or(Some(plan_request_id));
let request_id_for_log = short_request_id(request_id.unwrap_or("-"));
let candidate_id = result.candidate_id.as_deref().or(plan_candidate_id);
let status_code = result.status_code;
let has_body_bytes = body_base64.is_some();
let explicit_finalize = should_finalize_sync_response(report_kind.as_deref());
let mapped_error_finalize_kind =
resolve_core_sync_error_finalize_report_kind(plan_kind, &result, body_json.as_ref());
let implicit_finalize = if explicit_finalize || mapped_error_finalize_kind.is_some() {
None
} else {
let implicit_finalize = if !explicit_finalize && mapped_error_finalize_kind.is_none() {
maybe_build_implicit_sync_finalize_outcome(
trace_id,
decision,
plan_kind,
report_context.clone(),
result.status_code,
headers.clone(),
body_json.clone(),
body_base64.clone(),
result.telemetry.clone(),
&report_context,
status_code,
&headers,
&body_json,
&body_base64,
&result.telemetry,
)?
};
let finalize_report_kind = if explicit_finalize {
report_kind.clone()
} else if let Some(implicit_finalize) = implicit_finalize.as_ref() {
Some(implicit_finalize.payload.report_kind.clone())
} else {
mapped_error_finalize_kind.clone()
None
};
if !matches!(
local_failover_analysis.decision,
LocalFailoverDecision::StopLocalFailover
@@ -486,91 +527,83 @@ pub(crate) async fn execute_execution_runtime_sync(
)
.await;
let base_usage_payload = GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind: finalize_report_kind
.clone()
.or_else(|| report_kind.clone())
.unwrap_or_default(),
report_context: report_context.clone(),
status_code: result.status_code,
headers: headers.clone(),
body_json: body_json.clone(),
client_body_json: None,
body_base64: body_base64.clone(),
telemetry: result.telemetry.clone(),
};
if result.status_code < 400 {
apply_local_execution_effect(
let request_id_owned = result.request_id;
let candidate_id_owned = result.candidate_id;
let request_id = (!request_id_owned.trim().is_empty())
.then_some(request_id_owned.as_str())
.or(Some(plan_request_id));
let request_id_for_log = short_request_id(request_id.unwrap_or("-"));
let candidate_id = candidate_id_owned.as_deref().or(plan_candidate_id);
let report_context = report_context;
let headers = headers;
let body_json = body_json;
let telemetry = result.telemetry;
if let Some(implicit_finalize) = implicit_finalize {
let usage_payload = implicit_finalize
.outcome
.background_report
.as_ref()
.unwrap_or(&implicit_finalize.payload);
apply_sync_success_effects(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
&plan,
implicit_finalize.payload.report_context.as_ref(),
usage_payload,
)
.await;
apply_local_execution_effect(
record_sync_terminal_usage(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::AdaptiveSuccess(LocalAdaptiveSuccessEffect),
)
.await;
apply_local_execution_effect(
state,
LocalExecutionEffectContext {
plan: &plan,
report_context: report_context.as_ref(),
},
LocalExecutionEffect::PoolSuccessSync {
payload: &base_usage_payload,
},
)
.await;
&plan,
implicit_finalize.payload.report_context.as_ref(),
usage_payload,
);
if let Some(report_payload) = implicit_finalize.outcome.background_report {
spawn_sync_report(state.clone(), report_payload);
} else {
warn!(
event_name = "local_core_finalize_missing_success_report_mapping",
log_type = "event",
trace_id = %trace_id,
report_kind = %implicit_finalize.payload.report_kind,
"gateway implicit local core finalize produced response without background success report mapping"
);
}
return Ok(Some(attach_control_metadata_headers(
implicit_finalize.outcome.response,
request_id,
candidate_id,
)?));
}
if let Some(finalize_report_kind) = finalize_report_kind {
if let Some(implicit_finalize) = implicit_finalize {
let usage_payload = implicit_finalize
.outcome
.background_report
.as_ref()
.unwrap_or(&implicit_finalize.payload);
record_sync_terminal_usage(state, &plan, report_context.as_ref(), usage_payload);
if let Some(report_payload) = implicit_finalize.outcome.background_report {
spawn_sync_report(state.clone(), trace_id.to_string(), report_payload);
} else {
warn!(
event_name = "local_core_finalize_missing_success_report_mapping",
log_type = "event",
trace_id = %trace_id,
report_kind = %implicit_finalize.payload.report_kind,
"gateway implicit local core finalize produced response without background success report mapping"
);
}
return Ok(Some(attach_control_metadata_headers(
implicit_finalize.outcome.response,
request_id,
candidate_id,
)?));
}
let finalize_report_kind = if explicit_finalize {
report_kind.clone()
} else {
mapped_error_finalize_kind
};
let payload = GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind: finalize_report_kind,
if let Some(finalize_report_kind) = finalize_report_kind {
let mut payload = build_sync_report_payload(
trace_id,
finalize_report_kind,
report_context,
status_code: result.status_code,
headers: headers.clone(),
body_json: body_json.clone(),
client_body_json: None,
body_base64: body_base64.clone(),
telemetry: result.telemetry.clone(),
};
status_code,
headers,
body_json,
body_base64,
telemetry,
);
if let Some(outcome) = maybe_build_sync_finalize_outcome(trace_id, decision, &payload)? {
let usage_payload = outcome.background_report.as_ref().unwrap_or(&payload);
if status_code < 400 {
apply_sync_success_effects(
state,
&plan,
payload.report_context.as_ref(),
usage_payload,
)
.await;
}
record_sync_terminal_usage(
state,
&plan,
@@ -578,7 +611,7 @@ pub(crate) async fn execute_execution_runtime_sync(
usage_payload,
);
if let Some(report_payload) = outcome.background_report {
spawn_sync_report(state.clone(), trace_id.to_string(), report_payload);
spawn_sync_report(state.clone(), report_payload);
} else {
warn!(
event_name = "local_core_finalize_missing_success_report_mapping",
@@ -594,55 +627,62 @@ pub(crate) async fn execute_execution_runtime_sync(
candidate_id,
)?));
}
if let Some(outcome) = maybe_build_local_video_success_outcome(
let mut payload = match maybe_build_local_video_success_outcome(
trace_id,
decision,
&payload,
payload,
&state.video_tasks,
&plan,
)? {
record_sync_terminal_usage(
state,
&plan,
payload.report_context.as_ref(),
&outcome.report_payload,
);
if let Some(snapshot) = outcome.local_task_snapshot.clone() {
state.video_tasks.record_snapshot(snapshot.clone());
let _ = state.upsert_video_task_snapshot(&snapshot).await?;
}
match outcome.report_mode {
VideoTaskSyncReportMode::InlineSync => {
submit_sync_report(state, trace_id, outcome.report_payload).await?;
LocalVideoSyncSuccessBuild::Handled(outcome) => {
let LocalVideoSyncSuccessOutcome {
response,
report_payload,
original_report_context,
report_mode,
local_task_snapshot,
} = outcome;
apply_sync_success_effects(
state,
&plan,
original_report_context.as_ref(),
&report_payload,
)
.await;
record_sync_terminal_usage(
state,
&plan,
original_report_context.as_ref(),
&report_payload,
);
if let Some(snapshot) = local_task_snapshot {
let _ = state.upsert_video_task_snapshot(&snapshot).await?;
state.video_tasks.record_snapshot(snapshot);
}
VideoTaskSyncReportMode::Background => {
spawn_sync_report(state.clone(), trace_id.to_string(), outcome.report_payload);
match report_mode {
VideoTaskSyncReportMode::InlineSync => {
submit_sync_report(state, report_payload).await?;
}
VideoTaskSyncReportMode::Background => {
spawn_sync_report(state.clone(), report_payload);
}
}
return Ok(Some(attach_control_metadata_headers(
response,
request_id,
candidate_id,
)?));
}
return Ok(Some(attach_control_metadata_headers(
outcome.response,
request_id,
candidate_id,
)?));
}
LocalVideoSyncSuccessBuild::NotHandled(payload) => payload,
};
if let Some(response) =
maybe_build_local_sync_finalize_response(trace_id, decision, &payload)?
{
let usage_payload = if let Some(success_report_kind) =
resolve_local_sync_success_background_report_kind(payload.report_kind.as_str())
{
let mut report_payload = payload.clone();
report_payload.report_kind = success_report_kind.to_string();
report_payload
} else {
payload.clone()
};
record_sync_terminal_usage(
state,
&plan,
payload.report_context.as_ref(),
&usage_payload,
);
let background_success_report_kind =
resolve_local_sync_success_background_report_kind(payload.report_kind.as_str());
apply_sync_success_effects(state, &plan, payload.report_context.as_ref(), &payload)
.await;
record_sync_terminal_usage(state, &plan, payload.report_context.as_ref(), &payload);
state
.video_tasks
.apply_finalize_mutation(request_path, payload.report_kind.as_str());
@@ -652,12 +692,11 @@ pub(crate) async fn execute_execution_runtime_sync(
{
let _ = state.upsert_video_task_snapshot(&snapshot).await?;
}
if let Some(success_report_kind) =
resolve_local_sync_success_background_report_kind(payload.report_kind.as_str())
{
let mut report_payload = usage_payload;
report_payload.report_kind = success_report_kind.to_string();
spawn_sync_report(state.clone(), trace_id.to_string(), report_payload);
if let Some(success_report_kind) = background_success_report_kind {
payload.report_kind = success_report_kind.to_string();
}
if background_success_report_kind.is_some() {
spawn_sync_report(state.clone(), payload);
} else {
warn!(
event_name = "local_video_finalize_missing_success_report_mapping",
@@ -678,27 +717,14 @@ pub(crate) async fn execute_execution_runtime_sync(
if let Some(response) =
maybe_build_local_video_error_response(trace_id, decision, &payload)?
{
let usage_payload = if let Some(error_report_kind) =
resolve_local_sync_error_background_report_kind(payload.report_kind.as_str())
{
let mut report_payload = payload.clone();
report_payload.report_kind = error_report_kind.to_string();
report_payload
} else {
payload.clone()
};
record_sync_terminal_usage(
state,
&plan,
payload.report_context.as_ref(),
&usage_payload,
);
if let Some(error_report_kind) =
resolve_local_sync_error_background_report_kind(payload.report_kind.as_str())
{
let mut report_payload = usage_payload;
report_payload.report_kind = error_report_kind.to_string();
spawn_sync_report(state.clone(), trace_id.to_string(), report_payload);
let background_error_report_kind =
resolve_local_sync_error_background_report_kind(payload.report_kind.as_str());
if let Some(error_report_kind) = background_error_report_kind {
payload.report_kind = error_report_kind.to_string();
}
record_sync_terminal_usage(state, &plan, payload.report_context.as_ref(), &payload);
if background_error_report_kind.is_some() {
spawn_sync_report(state.clone(), payload);
} else {
warn!(
event_name = "local_video_finalize_missing_error_report_mapping",
@@ -726,49 +752,47 @@ pub(crate) async fn execute_execution_runtime_sync(
)?));
}
record_sync_terminal_usage(state, &plan, report_context.as_ref(), &base_usage_payload);
if let Some(report_kind) = report_kind {
let report = GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind,
report_context,
status_code: result.status_code,
headers: headers.clone(),
body_json: body_json.clone(),
client_body_json: None,
body_base64: body_base64.clone(),
telemetry: result.telemetry.clone(),
};
spawn_sync_report(state.clone(), trace_id.to_string(), report);
}
let request_id_header: Option<&str> = request_id
.map(str::trim)
.filter(|value: &&str| !value.is_empty());
if let Some(request_id) = request_id_header {
headers.insert(
CONTROL_REQUEST_ID_HEADER.to_string(),
request_id.to_string(),
);
}
let candidate_id_header: Option<&str> = candidate_id
.map(str::trim)
.filter(|value: &&str| !value.is_empty());
if let Some(candidate_id) = candidate_id_header {
headers.insert(
CONTROL_CANDIDATE_ID_HEADER.to_string(),
candidate_id.to_string(),
);
}
Ok(Some(build_client_response_from_parts(
result.status_code,
&headers,
Body::from(body_bytes),
let usage_payload = build_sync_report_payload(
trace_id,
Some(decision),
)?))
report_kind.unwrap_or_default(),
report_context,
status_code,
headers,
body_json,
body_base64,
telemetry,
);
if status_code < 400 {
apply_sync_success_effects(
state,
&plan,
usage_payload.report_context.as_ref(),
&usage_payload,
)
.await;
}
record_sync_terminal_usage(
state,
&plan,
usage_payload.report_context.as_ref(),
&usage_payload,
);
let response = attach_control_metadata_headers(
build_client_response_from_parts(
status_code,
&usage_payload.headers,
Body::from(body_bytes),
trace_id,
Some(decision),
)?,
request_id,
candidate_id,
)?;
if !usage_payload.report_kind.trim().is_empty() {
spawn_sync_report(state.clone(), usage_payload);
}
Ok(Some(response))
}
#[allow(clippy::too_many_arguments)] // mirrors sync execution context
@@ -776,12 +800,12 @@ fn maybe_build_implicit_sync_finalize_outcome(
trace_id: &str,
decision: &GatewayControlDecision,
plan_kind: &str,
report_context: Option<serde_json::Value>,
report_context: &Option<serde_json::Value>,
status_code: u16,
headers: BTreeMap<String, String>,
body_json: Option<serde_json::Value>,
body_base64: Option<String>,
telemetry: Option<ExecutionTelemetry>,
headers: &BTreeMap<String, String>,
body_json: &Option<serde_json::Value>,
body_base64: &Option<String>,
telemetry: &Option<ExecutionTelemetry>,
) -> Result<Option<ImplicitSyncFinalizeOutcome>, GatewayError> {
if status_code >= 400 || body_json.is_some() || body_base64.is_none() {
return Ok(None);
@@ -794,13 +818,13 @@ fn maybe_build_implicit_sync_finalize_outcome(
let payload = GatewaySyncReportRequest {
trace_id: trace_id.to_string(),
report_kind: report_kind.to_string(),
report_context,
report_context: report_context.clone(),
status_code,
headers,
body_json,
headers: headers.clone(),
body_json: body_json.clone(),
client_body_json: None,
body_base64,
telemetry,
body_base64: body_base64.clone(),
telemetry: telemetry.clone(),
};
let Some(outcome) = maybe_build_sync_finalize_outcome(trace_id, decision, &payload)? else {
return Ok(None);

10
apps/aether-gateway/src/execution_runtime/sync/execution/policy.rs
View File

@@ -1,6 +1,6 @@
use std::collections::BTreeMap;
use aether_contracts::ExecutionResult;
use aether_contracts::ResponseBody;
use base64::Engine as _;
use crate::GatewayError;
@@ -8,14 +8,14 @@ use crate::GatewayError;
type DecodedBody = (Vec<u8>, Option<serde_json::Value>, Option<String>);
pub(super) fn decode_execution_result_body(
result: &ExecutionResult,
body: Option<ResponseBody>,
headers: &mut BTreeMap<String, String>,
) -> Result<DecodedBody, GatewayError> {
let Some(body) = result.body.as_ref() else {
let Some(body) = body else {
return Ok((Vec::new(), None, None));
};
if let Some(json_body) = body.json_body.clone() {
if let Some(json_body) = body.json_body {
headers
.entry("content-type".to_string())
.or_insert_with(|| "application/json".to_string());
@@ -25,7 +25,7 @@ pub(super) fn decode_execution_result_body(
return Ok((bytes, Some(json_body), None));
}
if let Some(body_bytes_b64) = body.body_bytes_b64.clone() {
if let Some(body_bytes_b64) = body.body_bytes_b64 {
let bytes = base64::engine::general_purpose::STANDARD
.decode(&body_bytes_b64)
.map_err(|err| GatewayError::Internal(err.to_string()))?;

202
apps/aether-gateway/src/execution_runtime/sync/execution/response.rs
View File

@@ -2,10 +2,13 @@ use std::collections::BTreeMap;
use aether_contracts::ExecutionPlan;
use axum::body::Body;
use axum::http::header::HeaderValue;
use axum::http::Response;
use serde_json::json;
use crate::api::response::build_client_response_from_parts;
use crate::api::response::{
build_client_response_from_parts, build_client_response_from_parts_with_mutator,
};
use crate::async_task::VideoTaskService;
use crate::control::GatewayControlDecision;
use crate::video_tasks::{
@@ -17,9 +20,15 @@ pub(crate) use crate::video_tasks::{
};
use crate::{usage::GatewaySyncReportRequest, GatewayError};
pub(crate) enum LocalVideoSyncSuccessBuild {
Handled(LocalVideoSyncSuccessOutcome),
NotHandled(GatewaySyncReportRequest),
}
pub(crate) struct LocalVideoSyncSuccessOutcome {
pub(crate) response: Response<Body>,
pub(crate) report_payload: GatewaySyncReportRequest,
pub(crate) original_report_context: Option<serde_json::Value>,
pub(crate) report_mode: VideoTaskSyncReportMode,
pub(crate) local_task_snapshot: Option<LocalVideoTaskSnapshot>,
}
@@ -29,8 +38,9 @@ fn cloned_report_context_object(
) -> serde_json::Map<String, serde_json::Value> {
payload
.report_context
.clone()
.and_then(|value| value.as_object().cloned())
.as_ref()
.and_then(serde_json::Value::as_object)
.cloned()
.unwrap_or_default()
}
@@ -56,54 +66,53 @@ fn build_local_video_success_response(
pub(crate) fn maybe_build_local_video_success_outcome(
trace_id: &str,
decision: &GatewayControlDecision,
payload: &GatewaySyncReportRequest,
mut payload: GatewaySyncReportRequest,
video_tasks: &VideoTaskService,
plan: &ExecutionPlan,
) -> Result<Option<LocalVideoSyncSuccessOutcome>, GatewayError> {
) -> Result<LocalVideoSyncSuccessBuild, GatewayError> {
if payload.status_code >= 400 {
return Ok(None);
return Ok(LocalVideoSyncSuccessBuild::NotHandled(payload));
}
let provider_body = match payload
.body_json
.as_ref()
.and_then(serde_json::Value::as_object)
{
Some(value) => value,
None => return Ok(None),
let mut report_context = cloned_report_context_object(&payload);
let prepared_plan = {
let provider_body = match payload
.body_json
.as_ref()
.and_then(serde_json::Value::as_object)
{
Some(value) => value,
None => return Ok(LocalVideoSyncSuccessBuild::NotHandled(payload)),
};
video_tasks.prepare_sync_success(
payload.report_kind.as_str(),
provider_body,
&report_context,
plan,
)
};
let mut report_context = cloned_report_context_object(payload);
let Some(plan) = video_tasks.prepare_sync_success(
payload.report_kind.as_str(),
provider_body,
&report_context,
plan,
) else {
return Ok(None);
let Some(plan) = prepared_plan else {
return Ok(LocalVideoSyncSuccessBuild::NotHandled(payload));
};
plan.apply_to_report_context(&mut report_context);
let client_body_json = plan.client_body_json();
let response = build_local_video_success_response(trace_id, decision, &client_body_json)?;
let report_payload = GatewaySyncReportRequest {
trace_id: payload.trace_id.clone(),
report_kind: plan.success_report_kind().to_string(),
report_context: Some(serde_json::Value::Object(report_context)),
status_code: payload.status_code,
headers: payload.headers.clone(),
body_json: payload.body_json.clone(),
client_body_json: Some(client_body_json),
body_base64: None,
telemetry: payload.telemetry.clone(),
};
let original_report_context = payload.report_context.take();
payload.report_kind = plan.success_report_kind().to_string();
payload.report_context = Some(serde_json::Value::Object(report_context));
payload.client_body_json = Some(client_body_json);
Ok(Some(LocalVideoSyncSuccessOutcome {
response,
report_payload,
report_mode: plan.report_mode(),
local_task_snapshot: matches!(plan.report_mode(), VideoTaskSyncReportMode::Background)
.then(|| plan.to_snapshot()),
}))
Ok(LocalVideoSyncSuccessBuild::Handled(
LocalVideoSyncSuccessOutcome {
response,
report_payload: payload,
original_report_context,
report_mode: plan.report_mode(),
local_task_snapshot: matches!(plan.report_mode(), VideoTaskSyncReportMode::Background)
.then(|| plan.to_snapshot()),
},
))
}
pub(crate) fn maybe_build_local_sync_finalize_response(
@@ -147,21 +156,114 @@ pub(crate) fn maybe_build_local_video_error_response(
return Ok(None);
}
let response_body = payload.body_json.clone().unwrap_or_else(|| json!({}));
let body_bytes = serde_json::to_vec(&response_body)
.map_err(|err| GatewayError::Internal(err.to_string()))?;
let empty_body = json!({});
let response_body = payload.body_json.as_ref().unwrap_or(&empty_body);
let body_bytes =
serde_json::to_vec(response_body).map_err(|err| GatewayError::Internal(err.to_string()))?;
let body_len = body_bytes.len().to_string();
let mut response_headers = payload.headers.clone();
response_headers.remove("content-encoding");
response_headers.remove("content-length");
response_headers.insert("content-type".to_string(), "application/json".to_string());
response_headers.insert("content-length".to_string(), body_bytes.len().to_string());
Ok(Some(build_client_response_from_parts(
Ok(Some(build_client_response_from_parts_with_mutator(
payload.status_code,
&response_headers,
&payload.headers,
Body::from(body_bytes),
trace_id,
Some(decision),
|headers| {
headers.remove(http::header::CONTENT_ENCODING);
headers.remove(http::header::CONTENT_LENGTH);
headers.insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/json"),
);
headers.insert(
http::header::CONTENT_LENGTH,
HeaderValue::from_str(body_len.as_str())
.map_err(|err| GatewayError::Internal(err.to_string()))?,
);
Ok(())
},
)?))
}
#[cfg(test)]
mod tests {
use super::*;
use axum::body::to_bytes;
use serde_json::json;
#[tokio::test]
async fn local_video_error_response_rewrites_headers_without_mutating_payload() {
let decision = GatewayControlDecision::synthetic(
"/v1/videos",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("video".to_string()),
Some("openai:video".to_string()),
)
.with_execution_runtime_candidate(true);
let payload = GatewaySyncReportRequest {
trace_id: "trace-payload".to_string(),
report_kind: "openai_video_create_sync_finalize".to_string(),
report_context: Some(json!({
"request_id": "req_123",
})),
status_code: http::StatusCode::BAD_GATEWAY.as_u16(),
headers: BTreeMap::from([
("content-encoding".to_string(), "gzip".to_string()),
("content-length".to_string(), "999".to_string()),
("x-upstream-id".to_string(), "video-123".to_string()),
]),
body_json: Some(json!({
"error": {
"type": "video_backend_error",
"message": "backend failed",
}
})),
client_body_json: None,
body_base64: None,
telemetry: None,
};
let response =
maybe_build_local_video_error_response("trace-response", &decision, &payload)
.expect("video error response should build")
.expect("video error response should match local video error kinds");
assert_eq!(response.status(), http::StatusCode::BAD_GATEWAY);
assert_eq!(
response
.headers()
.get(http::header::CONTENT_TYPE)
.and_then(|value| value.to_str().ok()),
Some("application/json")
);
assert_eq!(response.headers().get(http::header::CONTENT_ENCODING), None);
assert_eq!(
response
.headers()
.get("x-upstream-id")
.and_then(|value| value.to_str().ok()),
Some("video-123")
);
assert_eq!(
payload.headers.get("content-encoding").map(String::as_str),
Some("gzip")
);
assert_eq!(
payload.headers.get("content-length").map(String::as_str),
Some("999")
);
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&body).expect("response body should parse"),
payload
.body_json
.clone()
.expect("payload body should exist")
);
}
}

3
apps/aether-gateway/src/execution_runtime/sync/mod.rs
View File

@@ -6,5 +6,6 @@ pub(crate) use execution::execute_execution_runtime_sync;
pub(crate) use execution::{
maybe_build_local_sync_finalize_response, maybe_build_local_video_error_response,
maybe_build_local_video_success_outcome, resolve_local_sync_error_background_report_kind,
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessOutcome,
resolve_local_sync_success_background_report_kind, LocalVideoSyncSuccessBuild,
LocalVideoSyncSuccessOutcome,
};

40
apps/aether-gateway/src/execution_runtime/transport.rs
View File

@@ -161,12 +161,12 @@ impl DirectSyncExecutionRuntime {
pub(crate) async fn execute_sync(
&self,
plan: ExecutionPlan,
plan: &ExecutionPlan,
) -> Result<ExecutionResult, ExecutionRuntimeTransportError> {
let body_bytes = build_request_body(&plan)?;
let body_bytes = build_request_body(plan)?;
let started_at = Instant::now();
let response = send_request(&plan, body_bytes).await?;
let response = send_request(plan, body_bytes).await?;
let ttfb_ms = started_at.elapsed().as_millis() as u64;
let status_code = response.status().as_u16();
let headers = collect_response_headers(response.headers());
@@ -200,8 +200,8 @@ impl DirectSyncExecutionRuntime {
};
Ok(ExecutionResult {
request_id: plan.request_id,
candidate_id: plan.candidate_id,
request_id: plan.request_id.clone(),
candidate_id: plan.candidate_id.clone(),
status_code,
headers,
body,
@@ -216,24 +216,24 @@ impl DirectSyncExecutionRuntime {
pub(crate) async fn execute_stream(
&self,
plan: ExecutionPlan,
plan: &ExecutionPlan,
) -> Result<DirectUpstreamStreamExecution, ExecutionRuntimeTransportError> {
if !plan.stream {
return Err(ExecutionRuntimeTransportError::StreamUnsupported);
}
let body_bytes = build_request_body(&plan)?;
let body_bytes = build_request_body(plan)?;
let started_at = Instant::now();
let response = send_request(&plan, body_bytes).await?;
let response = send_request(plan, body_bytes).await?;
let status_code = response.status().as_u16();
let headers = collect_response_headers(response.headers());
let stream_summary_report_context = build_stream_summary_report_context(&plan);
let stream_summary_report_context = build_stream_summary_report_context(plan);
Ok(DirectUpstreamStreamExecution {
request_id: plan.request_id,
candidate_id: plan.candidate_id,
request_id: plan.request_id.clone(),
candidate_id: plan.candidate_id.clone(),
status_code,
headers,
provider_api_format: plan.provider_api_format.clone(),
@@ -274,7 +274,7 @@ pub(crate) async fn execute_sync_plan(
let _ = state;
let _ = trace_id;
DirectSyncExecutionRuntime::new()
.execute_sync(plan.clone())
.execute_sync(plan)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
@@ -1101,7 +1101,7 @@ mod tests {
let execution_runtime = DirectSyncExecutionRuntime::new();
let result = execution_runtime
.execute_sync(ExecutionPlan {
.execute_sync(&ExecutionPlan {
request_id: "req-1".into(),
candidate_id: Some("cand-1".into()),
provider_name: Some("openai".into()),
@@ -1168,7 +1168,7 @@ mod tests {
let execution_runtime = DirectSyncExecutionRuntime::new();
let result = execution_runtime
.execute_sync(ExecutionPlan {
.execute_sync(&ExecutionPlan {
request_id: "req-1".into(),
candidate_id: None,
provider_name: None,
@@ -1369,7 +1369,7 @@ mod tests {
let execution_runtime = DirectSyncExecutionRuntime::new();
let result = execution_runtime
.execute_sync(ExecutionPlan {
.execute_sync(&ExecutionPlan {
request_id: "req-redirect-1".into(),
candidate_id: None,
provider_name: Some("provider_ops".into()),
@@ -1442,7 +1442,7 @@ mod tests {
let execution_runtime = DirectSyncExecutionRuntime::new();
let result = execution_runtime
.execute_sync(ExecutionPlan {
.execute_sync(&ExecutionPlan {
request_id: "req-redirect-2".into(),
candidate_id: None,
provider_name: Some("provider_oauth".into()),
@@ -1512,7 +1512,7 @@ mod tests {
let execution_runtime = DirectSyncExecutionRuntime::new();
let result = execution_runtime
.execute_sync(ExecutionPlan {
.execute_sync(&ExecutionPlan {
request_id: "req-relay-http1-1".into(),
candidate_id: None,
provider_name: Some("provider_ops".into()),
@@ -1579,7 +1579,7 @@ mod tests {
let execution_runtime = DirectSyncExecutionRuntime::new();
let result = execution_runtime
.execute_sync(ExecutionPlan {
.execute_sync(&ExecutionPlan {
request_id: "req-tls-1".into(),
candidate_id: Some("cand-1".into()),
provider_name: Some("claude".into()),
@@ -1654,7 +1654,7 @@ mod tests {
let execution_runtime = DirectSyncExecutionRuntime::new();
let result = execution_runtime
.execute_sync(ExecutionPlan {
.execute_sync(&ExecutionPlan {
request_id: "req-gzip-1".into(),
candidate_id: Some("cand-1".into()),
provider_name: Some("openai".into()),
@@ -1715,7 +1715,7 @@ mod tests {
let execution_runtime = DirectSyncExecutionRuntime::new();
let result = execution_runtime
.execute_sync(ExecutionPlan {
.execute_sync(&ExecutionPlan {
request_id: "req-ttfb-1".into(),
candidate_id: Some("cand-1".into()),
provider_name: Some("openai".into()),