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

Add Codex image progress heartbeat

Browse Source
This commit is contained in:
Entropy.Xu
2026-05-09 01:21:26 +08:00
parent 9a84a6ff6c
commit 3b4f27f767
44 changed files with 2676 additions and 208 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
apps/aether-gateway/src/ai_serving/api.rs
View File

@@ -30,6 +30,7 @@ pub(crate) use crate::ai_serving::{
maybe_build_stream_decision_payload, maybe_build_stream_plan_payload,
maybe_build_sync_decision_payload, maybe_build_sync_plan_payload,
set_local_openai_chat_execution_exhausted_diagnostic,
set_local_openai_image_execution_exhausted_diagnostic,
};
pub(crate) use crate::ai_serving::{
maybe_bridge_standard_sync_json_to_stream, maybe_build_provider_private_stream_normalizer,

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

@@ -49,7 +49,8 @@ pub(crate) use self::planner::{
extract_pool_sticky_session_token, maybe_build_stream_decision_payload,
maybe_build_stream_plan_payload, maybe_build_sync_decision_payload,
maybe_build_sync_plan_payload, planner_is_matching_stream_request,
set_local_openai_chat_execution_exhausted_diagnostic, CandidateFailureDiagnostic,
set_local_openai_chat_execution_exhausted_diagnostic,
set_local_openai_image_execution_exhausted_diagnostic, CandidateFailureDiagnostic,
CandidateFailureDiagnosticKind, GatewayAuthApiKeySnapshot, GatewayProviderTransportSnapshot,
LocalExecutionAttemptSource, LocalResolvedOAuthRequestAuth, PlannerAppState,
};

1
apps/aether-gateway/src/ai_serving/planner/mod.rs
View File

@@ -48,6 +48,7 @@ pub(crate) use self::specialized::{
build_local_image_sync_plan_and_reports_for_kind,
build_local_video_sync_attempt_source_for_kind,
build_local_video_sync_plan_and_reports_for_kind,
set_local_openai_image_execution_exhausted_diagnostic,
};
pub(crate) use self::standard::{
build_local_openai_chat_stream_attempt_source_for_kind,

30
apps/aether-gateway/src/ai_serving/planner/specialized/image.rs
View File

@@ -10,6 +10,7 @@ use crate::ai_serving::planner::plan_builders::{
build_passthrough_sync_plan_from_decision, build_standard_stream_plan_from_decision,
AiStreamAttempt, AiSyncAttempt,
};
use crate::ai_serving::planner::runtime_miss::set_local_runtime_execution_exhausted_diagnostic;
use crate::ai_serving::planner::spec_metadata::local_openai_image_spec_metadata;
use crate::ai_serving::GatewayControlDecision;
use crate::ai_serving::{
@@ -49,6 +50,35 @@ pub(crate) struct LocalOpenAiImageStreamAttemptSource<'a> {
candidates: LocalOpenAiImageCandidateAttemptSource<'a>,
}
pub(crate) fn set_local_openai_image_execution_exhausted_diagnostic(
state: &AppState,
trace_id: &str,
decision: &GatewayControlDecision,
plan_kind: &str,
body_json: &serde_json::Value,
candidate_count: usize,
) {
warn!(
event_name = "local_openai_image_candidates_exhausted",
log_type = "event",
trace_id = %trace_id,
plan_kind,
route_class = decision.route_class.as_deref().unwrap_or("passthrough"),
route_family = decision.route_family.as_deref().unwrap_or("unknown"),
candidate_count,
model = body_json.get("model").and_then(|value| value.as_str()).unwrap_or(""),
"gateway local openai image execution exhausted all candidates"
);
set_local_runtime_execution_exhausted_diagnostic(
state,
trace_id,
decision,
plan_kind,
body_json.get("model").and_then(|value| value.as_str()),
candidate_count,
);
}
pub(crate) async fn build_local_image_sync_plan_and_reports_for_kind(
state: &AppState,
parts: &http::request::Parts,

1
apps/aether-gateway/src/ai_serving/planner/specialized/mod.rs
View File

@@ -18,6 +18,7 @@ pub(crate) use self::image::{
build_local_image_sync_attempt_source_for_kind,
build_local_image_sync_plan_and_reports_for_kind,
maybe_build_stream_local_image_decision_payload, maybe_build_sync_local_image_decision_payload,
set_local_openai_image_execution_exhausted_diagnostic,
};
pub(crate) use self::video::{
build_local_video_sync_attempt_source_for_kind,

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

@@ -1,6 +1,10 @@
use std::collections::{BTreeMap, VecDeque};
use std::io::Error as IoError;
use std::time::Instant;
use std::sync::{
atomic::{AtomicBool, AtomicU64, Ordering},
Arc,
};
use std::time::{Duration, Instant};
use aether_contracts::{
ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry, StreamFrame,
@@ -25,6 +29,7 @@ use futures_util::stream::BoxStream;
use futures_util::{StreamExt, TryStreamExt};
use serde_json::Value;
use tokio::sync::mpsc;
use tokio::time::MissedTickBehavior;
use tokio_util::codec::{FramedRead, LinesCodec};
use tokio_util::io::StreamReader;
use tracing::{debug, info, warn};
@@ -88,6 +93,14 @@ use crate::usage::submit_stream_report;
use crate::usage::{GatewayStreamReportRequest, GatewaySyncReportRequest};
use crate::{AppState, GatewayError};
const OPENAI_IMAGE_STREAM_PLAN_KIND: &str = "openai_image_stream";
const SSE_KEEPALIVE_INTERVAL: Duration = Duration::from_secs(15);
const SSE_KEEPALIVE_BYTES: &[u8] = b": aether-keepalive\n\n";
const STREAM_IDLE_LOG_INTERVAL: Duration = Duration::from_secs(60);
const STREAM_IDLE_LOG_INTERVAL_MS: u64 = 60_000;
const REWRITTEN_STREAM_PREFETCH_TIMEOUT: Duration = Duration::from_millis(750);
const OPENAI_IMAGE_STREAM_DEFAULT_TOTAL_TIMEOUT_MS: u64 = 900_000;
fn record_sync_terminal_usage(
state: &AppState,
plan: &ExecutionPlan,
@@ -750,6 +763,114 @@ fn encode_terminal_sse_error_event(failure: &StreamFailureReport) -> Result<Byte
Ok(Bytes::from(event))
}
fn image_stream_failed_event_name(report_context: Option<&Value>) -> &'static str {
let operation = report_context
.and_then(|value| value.get("image_request"))
.and_then(|value| value.get("operation"))
.and_then(Value::as_str)
.map(str::trim)
.unwrap_or_default();
if operation == "edit" {
"image_edit.failed"
} else {
"image_generation.failed"
}
}
fn encode_openai_image_failed_event(
report_context: Option<&Value>,
failure: &StreamFailureReport,
) -> Result<Bytes, std::io::Error> {
let event_name = image_stream_failed_event_name(report_context);
let failure_body = failure
.to_json_string()
.map_err(|err| IoError::other(err.to_string()))?;
let failure_json: Value =
serde_json::from_str(&failure_body).map_err(|err| IoError::other(err.to_string()))?;
let error = failure_json.get("error").cloned().unwrap_or_else(|| {
serde_json::json!({
"type": failure.error_type.as_str(),
"message": failure.error_message.as_str(),
"code": failure.status_code,
})
});
let payload = serde_json::json!({
"type": event_name,
"error": error,
});
let payload = serde_json::to_string(&payload).map_err(|err| IoError::other(err.to_string()))?;
let mut event = format!("event: {event_name}\n");
for line in payload.lines() {
event.push_str("data: ");
event.push_str(line);
event.push('\n');
}
event.push('\n');
Ok(Bytes::from(event))
}
fn resolve_openai_image_stream_total_timeout_ms(
plan_kind: &str,
plan: &ExecutionPlan,
) -> Option<u64> {
if plan_kind != OPENAI_IMAGE_STREAM_PLAN_KIND {
return None;
}
Some(
plan.timeouts
.as_ref()
.and_then(|timeouts| timeouts.total_ms)
.unwrap_or(OPENAI_IMAGE_STREAM_DEFAULT_TOTAL_TIMEOUT_MS)
.max(1),
)
}
fn should_limit_direct_finalize_prefetch(plan_kind: &str, has_local_stream_rewriter: bool) -> bool {
plan_kind == OPENAI_IMAGE_STREAM_PLAN_KIND || has_local_stream_rewriter
}
fn build_sse_body_stream(
prefetched_chunks_for_body: Vec<Bytes>,
mut rx: mpsc::Receiver<Result<Bytes, IoError>>,
emit_keepalive: bool,
keepalive_interval: Duration,
) -> impl futures_util::Stream<Item = Result<Bytes, IoError>> + Send + 'static {
stream! {
let mut sent_prefetched_chunk = false;
for chunk in prefetched_chunks_for_body {
sent_prefetched_chunk = true;
yield Ok(chunk);
}
if emit_keepalive {
if !sent_prefetched_chunk {
yield Ok(Bytes::from_static(SSE_KEEPALIVE_BYTES));
}
let mut keepalive = tokio::time::interval(keepalive_interval);
keepalive.set_missed_tick_behavior(MissedTickBehavior::Delay);
keepalive.tick().await;
loop {
tokio::select! {
biased;
item = rx.recv() => {
let Some(item) = item else {
break;
};
yield item;
}
_ = keepalive.tick() => {
yield Ok(Bytes::from_static(SSE_KEEPALIVE_BYTES));
}
}
}
} else {
while let Some(item) = rx.recv().await {
yield item;
}
}
}
}
async fn next_stream_frame<R>(
buffered_frames: &mut VecDeque<StreamFrame>,
lines: &mut FramedRead<R, LinesCodec>,
@@ -1217,6 +1338,8 @@ async fn execute_stream_from_frame_stream(
private_stream_normalizer.is_some(),
local_stream_rewriter.is_some(),
);
let limit_direct_finalize_prefetch =
should_limit_direct_finalize_prefetch(plan_kind, local_stream_rewriter.is_some());
let mut prefetched_chunks: Vec<Bytes> = Vec::new();
let mut provider_prefetched_body = Vec::new();
let mut prefetched_body = Vec::new();
@@ -1250,7 +1373,38 @@ async fn execute_stream_from_frame_stream(
while prefetched_chunks.len() < MAX_STREAM_PREFETCH_FRAMES
&& prefetched_inspection_body.len() < MAX_STREAM_PREFETCH_BYTES
{
let Some(frame) = (match next_stream_frame(&mut buffered_frames, &mut lines).await {
let next_frame_result = if limit_direct_finalize_prefetch {
match tokio::time::timeout(
REWRITTEN_STREAM_PREFETCH_TIMEOUT,
next_stream_frame(&mut buffered_frames, &mut lines),
)
.await
{
Ok(result) => result,
Err(_) => {
debug!(
event_name = "execution_runtime_stream_prefetch_limited",
log_type = "debug",
trace_id = %trace_id,
request_id = %request_id_for_log,
candidate_id = ?candidate_id,
plan_kind,
report_kind,
provider_name,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
model_name,
candidate_index = candidate_index.as_str(),
timeout_ms = REWRITTEN_STREAM_PREFETCH_TIMEOUT.as_millis() as u64,
"gateway stopped rewritten stream prefetch before client-visible body"
);
break;
}
}
} else {
next_stream_frame(&mut buffered_frames, &mut lines).await
};
let Some(frame) = (match next_frame_result {
Ok(frame) => frame,
Err(err) => {
let failure = build_stream_failure_report(
@@ -1575,7 +1729,6 @@ async fn execute_stream_from_frame_stream(
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;
let trace_id_owned = trace_id.to_string();
let headers_for_report = headers.clone();
let report_kind_owned = report_kind;
@@ -1593,8 +1746,14 @@ async fn execute_stream_from_frame_stream(
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(&upstream_headers);
let candidate_index_for_report = candidate_index.clone();
let is_openai_image_stream_for_report = plan_kind == OPENAI_IMAGE_STREAM_PLAN_KIND;
let openai_image_stream_total_timeout_ms =
resolve_openai_image_stream_total_timeout_ms(plan_kind, &plan);
let plan_for_report = plan;
let emit_passthrough_sse_terminal_error = skip_direct_finalize_prefetch
&& response_headers_indicate_sse(&upstream_headers)
&& !is_openai_image_stream_for_report;
let body_capture_policy = match UsageRuntimeAccess::body_capture_policy(state.data.as_ref())
.await
{
@@ -1623,6 +1782,8 @@ async fn execute_stream_from_frame_stream(
.max_response_body_bytes
.unwrap_or(DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES)
};
let plan_kind_for_report = plan_kind.to_string();
let stream_started_at_for_report = stream_started_at;
tokio::spawn(async move {
let mut provider_buffered_body = Vec::new();
let mut buffered_body = Vec::new();
@@ -1667,6 +1828,115 @@ async fn execute_stream_from_frame_stream(
let reached_eof = initial_reached_eof;
let mut downstream_dropped = false;
let mut terminal_failure: Option<StreamFailureReport> = None;
let initial_elapsed_ms = stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX)) as u64;
let last_upstream_frame_elapsed_ms = Arc::new(AtomicU64::new(initial_elapsed_ms));
let last_client_chunk_elapsed_ms =
Arc::new(AtomicU64::new(if prefetched_body_for_report.is_empty() {
0
} else {
initial_elapsed_ms
}));
let provider_stream_bytes = Arc::new(AtomicU64::new(
u64::try_from(provider_prefetched_body_for_report.len()).unwrap_or(u64::MAX),
));
let client_stream_bytes = Arc::new(AtomicU64::new(
u64::try_from(prefetched_body_for_report.len()).unwrap_or(u64::MAX),
));
let idle_monitor_done = Arc::new(AtomicBool::new(false));
let idle_monitor_handle = {
let done = Arc::clone(&idle_monitor_done);
let last_upstream = Arc::clone(&last_upstream_frame_elapsed_ms);
let last_client = Arc::clone(&last_client_chunk_elapsed_ms);
let provider_bytes = Arc::clone(&provider_stream_bytes);
let client_bytes = Arc::clone(&client_stream_bytes);
let trace_id_for_idle = trace_id_owned.clone();
let request_id_for_idle = request_id_for_report_log.clone();
let candidate_id_for_idle = candidate_id_for_report.clone();
let candidate_index_for_idle = candidate_index_for_report.clone();
let plan_kind_for_idle = plan_kind_for_report.clone();
let provider_name_for_idle = plan_for_report
.provider_name
.clone()
.unwrap_or_else(|| "-".to_string());
let endpoint_id_for_idle = plan_for_report.endpoint_id.clone();
let key_id_for_idle = plan_for_report.key_id.clone();
let model_name_for_idle = plan_for_report
.model_name
.clone()
.unwrap_or_else(|| "-".to_string());
let has_local_stream_rewriter_for_idle = local_stream_rewriter.is_some();
tokio::spawn(async move {
let mut interval = tokio::time::interval(STREAM_IDLE_LOG_INTERVAL);
interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
interval.tick().await;
loop {
interval.tick().await;
if done.load(Ordering::Relaxed) {
break;
}
let elapsed_ms = stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX)) as u64;
let last_upstream_frame_elapsed_ms = last_upstream.load(Ordering::Relaxed);
let last_client_chunk_elapsed_ms = last_client.load(Ordering::Relaxed);
let upstream_idle_ms =
elapsed_ms.saturating_sub(last_upstream_frame_elapsed_ms);
let client_idle_ms = if last_client_chunk_elapsed_ms == 0 {
elapsed_ms
} else {
elapsed_ms.saturating_sub(last_client_chunk_elapsed_ms)
};
if upstream_idle_ms >= STREAM_IDLE_LOG_INTERVAL_MS {
warn!(
event_name = "stream_execution_upstream_idle",
log_type = "ops",
trace_id = %trace_id_for_idle,
request_id = %request_id_for_idle,
candidate_id = ?candidate_id_for_idle.as_deref(),
candidate_index = candidate_index_for_idle.as_str(),
plan_kind = plan_kind_for_idle.as_str(),
provider_name = provider_name_for_idle.as_str(),
endpoint_id = %endpoint_id_for_idle,
key_id = %key_id_for_idle,
model_name = model_name_for_idle.as_str(),
elapsed_ms,
provider_bytes = provider_bytes.load(Ordering::Relaxed),
client_bytes = client_bytes.load(Ordering::Relaxed),
last_upstream_frame_elapsed_ms,
last_client_chunk_elapsed_ms,
"gateway stream has not received an upstream frame within the idle window"
);
} else if client_idle_ms >= STREAM_IDLE_LOG_INTERVAL_MS
&& last_upstream_frame_elapsed_ms >= last_client_chunk_elapsed_ms
{
warn!(
event_name = "stream_execution_client_visible_idle",
log_type = "ops",
trace_id = %trace_id_for_idle,
request_id = %request_id_for_idle,
candidate_id = ?candidate_id_for_idle.as_deref(),
candidate_index = candidate_index_for_idle.as_str(),
plan_kind = plan_kind_for_idle.as_str(),
provider_name = provider_name_for_idle.as_str(),
endpoint_id = %endpoint_id_for_idle,
key_id = %key_id_for_idle,
model_name = model_name_for_idle.as_str(),
elapsed_ms,
provider_bytes = provider_bytes.load(Ordering::Relaxed),
client_bytes = client_bytes.load(Ordering::Relaxed),
last_upstream_frame_elapsed_ms,
last_client_chunk_elapsed_ms,
local_stream_rewriter = has_local_stream_rewriter_for_idle,
"gateway stream received upstream frames but has no recent client-visible chunk"
);
}
}
})
};
if !provider_prefetched_body_for_report.is_empty() {
let normalized_prefetched_chunk = if let Some(normalizer) =
private_stream_normalizer.as_mut()
@@ -1735,8 +2005,67 @@ async fn execute_stream_from_frame_stream(
}
if terminal_failure.is_none() && !reached_eof {
let mut image_stream_total_timeout = openai_image_stream_total_timeout_ms
.map(|timeout_ms| Box::pin(tokio::time::sleep(Duration::from_millis(timeout_ms))));
loop {
let next_frame = match next_stream_frame(&mut buffered_frames, &mut lines).await {
let next_frame_result = if let Some(timeout_sleep) =
image_stream_total_timeout.as_mut()
{
tokio::select! {
result = next_stream_frame(&mut buffered_frames, &mut lines) => result,
_ = timeout_sleep.as_mut() => {
let timeout_ms = openai_image_stream_total_timeout_ms
.unwrap_or(OPENAI_IMAGE_STREAM_DEFAULT_TOTAL_TIMEOUT_MS);
let elapsed_ms = stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX)) as u64;
warn!(
event_name = "openai_image_stream_total_timeout",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
candidate_index = candidate_index_for_report.as_str(),
plan_kind = plan_kind_for_report.as_str(),
provider_name = plan_for_report.provider_name.as_deref().unwrap_or("-"),
endpoint_id = %plan_for_report.endpoint_id,
key_id = %plan_for_report.key_id,
model_name = plan_for_report.model_name.as_deref().unwrap_or("-"),
elapsed_ms,
timeout_ms,
provider_bytes = provider_stream_bytes.load(Ordering::Relaxed),
client_bytes = client_stream_bytes.load(Ordering::Relaxed),
last_upstream_frame_elapsed_ms = last_upstream_frame_elapsed_ms.load(Ordering::Relaxed),
last_client_chunk_elapsed_ms = last_client_chunk_elapsed_ms.load(Ordering::Relaxed),
"gateway OpenAI image stream exceeded total timeout"
);
telemetry = Some(ExecutionTelemetry {
ttfb_ms: telemetry
.as_ref()
.and_then(|telemetry| telemetry.ttfb_ms)
.or_else(|| {
usage_stream_telemetry
.as_ref()
.and_then(|telemetry| telemetry.ttfb_ms)
}),
elapsed_ms: Some(elapsed_ms),
upstream_bytes: Some(provider_stream_bytes.load(Ordering::Relaxed)),
});
terminal_failure = Some(build_stream_failure_report(
"image_stream_total_timeout",
format!(
"OpenAI image stream exceeded total timeout of {timeout_ms}ms"
),
504,
));
break;
}
}
} else {
next_stream_frame(&mut buffered_frames, &mut lines).await
};
let next_frame = match next_frame_result {
Ok(frame) => frame,
Err(err) => {
warn!(
@@ -1759,6 +2088,11 @@ async fn execute_stream_from_frame_stream(
let Some(frame) = next_frame else {
break;
};
let frame_elapsed_ms = stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX)) as u64;
last_upstream_frame_elapsed_ms.store(frame_elapsed_ms, Ordering::Relaxed);
match frame.payload {
StreamFramePayload::Data { chunk_b64, text } => {
if sync_json_stream_bridge_active_for_report {
@@ -1792,6 +2126,10 @@ async fn execute_stream_from_frame_stream(
continue;
}
provider_stream_bytes.fetch_add(
u64::try_from(chunk.len()).unwrap_or(u64::MAX),
Ordering::Relaxed,
);
append_stream_capture_bytes(
&mut provider_buffered_body,
&chunk,
@@ -1870,7 +2208,7 @@ async fn execute_stream_from_frame_stream(
.and_then(|telemetry| telemetry.ttfb_ms)
.is_none()
{
let first_data_elapsed_ms = stream_started_at
let first_data_elapsed_ms = stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX))
@@ -1897,6 +2235,8 @@ async fn execute_stream_from_frame_stream(
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
let rewritten_chunk_len =
u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_disconnected",
@@ -1908,6 +2248,16 @@ async fn execute_stream_from_frame_stream(
);
downstream_dropped = true;
break;
} else {
client_stream_bytes.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX))
as u64,
Ordering::Relaxed,
);
}
}
StreamFramePayload::Telemetry {
@@ -2010,16 +2360,29 @@ async fn execute_stream_from_frame_stream(
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
let rewritten_chunk_len =
u64::try_from(rewritten_chunk.len()).unwrap_or(u64::MAX);
if tx.send(Ok(Bytes::from(rewritten_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_flush_disconnected",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway stream downstream dropped while flushing private stream normalization"
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway stream downstream dropped while flushing private stream normalization"
);
downstream_dropped = true;
} else {
client_stream_bytes
.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX))
as u64,
Ordering::Relaxed,
);
}
}
}
@@ -2054,16 +2417,28 @@ async fn execute_stream_from_frame_stream(
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
let flushed_chunk_len =
u64::try_from(flushed_chunk.len()).unwrap_or(u64::MAX);
if tx.send(Ok(Bytes::from(flushed_chunk))).await.is_err() {
warn!(
event_name = "stream_execution_downstream_rewrite_flush_disconnected",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway stream downstream dropped while flushing local stream rewrite"
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway stream downstream dropped while flushing local stream rewrite"
);
downstream_dropped = true;
} else {
client_stream_bytes.fetch_add(flushed_chunk_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX))
as u64,
Ordering::Relaxed,
);
}
}
Ok(_) => {}
@@ -2090,44 +2465,70 @@ async fn execute_stream_from_frame_stream(
}
}
if !downstream_dropped && emit_passthrough_sse_terminal_error {
if !downstream_dropped {
if let Some(failure) = terminal_failure.as_ref() {
match encode_terminal_sse_error_event(failure) {
Ok(error_event) => {
append_stream_capture_bytes(
&mut buffered_body,
error_event.as_ref(),
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
if tx.send(Ok(error_event)).await.is_err() {
warn!(
let terminal_event = if is_openai_image_stream_for_report {
Some(encode_openai_image_failed_event(
report_context_owned.as_ref(),
failure,
))
} else if emit_passthrough_sse_terminal_error {
Some(encode_terminal_sse_error_event(failure))
} else {
None
};
if let Some(terminal_event) = terminal_event {
match terminal_event {
Ok(error_event) => {
let error_event_len =
u64::try_from(error_event.len()).unwrap_or(u64::MAX);
append_stream_capture_bytes(
&mut buffered_body,
error_event.as_ref(),
max_stream_body_buffer_bytes,
&mut client_body_truncated,
);
if tx.send(Ok(error_event)).await.is_err() {
warn!(
event_name = "stream_execution_downstream_terminal_error_disconnected",
log_type = "ops",
trace_id = %trace_id_owned,
request_id = %request_id_for_report_log,
candidate_id = ?candidate_id_for_report.as_deref(),
"gateway stream downstream dropped while sending terminal SSE error event"
);
downstream_dropped = true;
);
downstream_dropped = true;
} else {
client_stream_bytes.fetch_add(error_event_len, Ordering::Relaxed);
last_client_chunk_elapsed_ms.store(
stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX))
as u64,
Ordering::Relaxed,
);
}
}
}
Err(err) => {
warn!(
Err(err) => {
warn!(
event_name = "stream_execution_terminal_error_event_encode_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 encode terminal SSE error event"
);
"gateway failed to encode terminal SSE error event"
);
}
}
}
}
}
drop(tx);
idle_monitor_done.store(true, Ordering::Relaxed);
idle_monitor_handle.abort();
stream_terminal_summary = merge_stream_terminal_summary(
stream_terminal_summary,
@@ -2292,15 +2693,6 @@ async fn execute_stream_from_frame_stream(
}
});
let body_stream = stream! {
for chunk in prefetched_chunks_for_body {
yield Ok(chunk);
}
while let Some(item) = rx.recv().await {
yield item;
}
};
headers.insert(CONTROL_REQUEST_ID_HEADER.to_string(), request_id.clone());
if let Some(candidate_id) = candidate_id
@@ -2314,6 +2706,17 @@ async fn execute_stream_from_frame_stream(
);
}
let emit_sse_keepalive = response_headers_indicate_sse(&headers);
if emit_sse_keepalive {
headers.remove("content-length");
}
let body_stream = build_sse_body_stream(
prefetched_chunks_for_body,
rx,
emit_sse_keepalive,
SSE_KEEPALIVE_INTERVAL,
);
Ok(Some(build_client_response_from_parts(
status_code,
&headers,
@@ -2337,7 +2740,7 @@ mod tests {
use std::collections::BTreeMap;
use std::convert::Infallible;
use std::sync::Arc;
use std::time::Duration;
use std::time::{Duration, Instant};
use aether_contracts::{
ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTimeouts, RequestBody,
@@ -2352,11 +2755,13 @@ mod tests {
use axum::extract::Request;
use axum::routing::any;
use axum::{http::header, http::HeaderValue, Router};
use futures_util::StreamExt as _;
use serde_json::{json, Value};
use tokio::sync::{watch, Notify};
use tokio::sync::{mpsc, watch, Notify};
use super::{
execute_execution_runtime_stream, merge_stream_terminal_summary,
build_sse_body_stream, execute_execution_runtime_stream, execute_stream_from_frame_stream,
merge_stream_terminal_summary, should_limit_direct_finalize_prefetch,
should_probe_success_failover_before_stream, should_skip_direct_finalize_prefetch,
};
use crate::control::GatewayControlDecision;
@@ -2482,6 +2887,132 @@ mod tests {
));
}
#[test]
fn limits_prefetch_for_openai_image_and_rewritten_streams() {
assert!(should_limit_direct_finalize_prefetch(
"openai_image_stream",
false
));
assert!(should_limit_direct_finalize_prefetch(
"openai_chat_stream",
true
));
assert!(!should_limit_direct_finalize_prefetch(
"openai_chat_stream",
false
));
}
#[tokio::test]
async fn sse_body_stream_emits_initial_and_periodic_keepalive_without_business_chunks() {
let (_tx, rx) = mpsc::channel::<Result<Bytes, std::io::Error>>(1);
let mut body_stream = Box::pin(build_sse_body_stream(
Vec::new(),
rx,
true,
Duration::from_millis(10),
));
let first = tokio::time::timeout(Duration::from_millis(50), body_stream.next())
.await
.expect("initial keepalive should be immediate")
.expect("stream should yield initial keepalive")
.expect("initial keepalive should be ok");
assert_eq!(first.as_ref(), b": aether-keepalive\n\n");
let second = tokio::time::timeout(Duration::from_millis(100), body_stream.next())
.await
.expect("periodic keepalive should arrive")
.expect("stream should yield periodic keepalive")
.expect("periodic keepalive should be ok");
assert_eq!(second.as_ref(), b": aether-keepalive\n\n");
}
#[tokio::test]
async fn openai_image_stream_total_timeout_emits_image_failed_event() {
let state = AppState::new().expect("app state should build");
let plan = ExecutionPlan {
request_id: "req-image-stream-timeout".into(),
candidate_id: Some("cand-image-stream-timeout".into()),
provider_name: Some("codex".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: "https://chatgpt.com/backend-api/codex/responses".into(),
headers: BTreeMap::from([
("content-type".into(), "application/json".into()),
("accept".into(), "text/event-stream".into()),
]),
content_type: Some("application/json".into()),
content_encoding: None,
body: RequestBody::from_json(json!({
"model": "gpt-image-1",
"prompt": "hello",
"stream": true
})),
stream: true,
client_api_format: "openai:image".into(),
provider_api_format: "openai:image".into(),
model_name: Some("gpt-image-1".into()),
proxy: None,
transport_profile: None,
timeouts: Some(ExecutionTimeouts {
total_ms: Some(25),
..ExecutionTimeouts::default()
}),
};
let decision = GatewayControlDecision::synthetic(
"/v1/images/generations",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("image".to_string()),
Some("openai:image".to_string()),
)
.with_execution_runtime_candidate(true);
let frame_stream = stream! {
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
));
std::future::pending::<()>().await;
}
.boxed();
let response = execute_stream_from_frame_stream(
&state,
plan,
"trace-image-stream-timeout",
&decision,
"openai_image_stream",
None,
Some(json!({
"provider_api_format": "openai:image",
"client_api_format": "openai:image",
"image_request": {
"operation": "generate"
}
})),
crate::clock::current_unix_ms(),
Instant::now(),
frame_stream,
)
.await
.expect("execution should succeed")
.expect("execution should return a client response");
let body = tokio::time::timeout(
Duration::from_secs(2),
to_bytes(response.into_body(), usize::MAX),
)
.await
.expect("timeout failure should close the response body")
.expect("response body should read");
let text = String::from_utf8(body.to_vec()).expect("response body should be utf8");
assert!(text.contains(": aether-keepalive\n\n"));
assert!(text.contains("event: image_generation.failed"));
assert!(text.contains("\"type\":\"image_stream_total_timeout\""));
}
#[tokio::test]
async fn execute_execution_runtime_stream_records_first_data_as_streaming_before_terminal_telemetry(
) {

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

File diff suppressed because it is too large Load Diff

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

@@ -551,7 +551,7 @@ fn build_direct_tunnel_request_meta(
}
}
async fn send_request(
pub(crate) async fn send_request(
plan: &ExecutionPlan,
body_bytes: Vec<u8>,
) -> Result<reqwest::Response, ExecutionRuntimeTransportError> {
@@ -728,7 +728,9 @@ async fn send_via_tunnel_relay(
Ok(response)
}
fn build_request_body(plan: &ExecutionPlan) -> Result<Vec<u8>, ExecutionRuntimeTransportError> {
pub(crate) fn build_request_body(
plan: &ExecutionPlan,
) -> Result<Vec<u8>, ExecutionRuntimeTransportError> {
let mut body_bytes = if let Some(json_body) = plan.body.json_body.clone() {
serde_json::to_vec(&json_body).map_err(ExecutionRuntimeTransportError::BodyEncode)?
} else if let Some(body_b64) = plan.body.body_bytes_b64.as_deref() {
@@ -1102,7 +1104,7 @@ fn is_hop_by_hop_header(name: &str) -> bool {
)
}
fn collect_response_headers(headers: &HeaderMap) -> BTreeMap<String, String> {
pub(crate) fn collect_response_headers(headers: &HeaderMap) -> BTreeMap<String, String> {
header_map_to_string_map(headers)
}
@@ -1130,7 +1132,7 @@ fn execution_log_url_host(url: &str) -> String {
.unwrap_or_else(|| "-".to_string())
}
fn decode_response_body_bytes(
pub(crate) fn decode_response_body_bytes(
headers: &BTreeMap<String, String>,
body_bytes: &[u8],
) -> Option<Vec<u8>> {
@@ -1157,7 +1159,7 @@ fn decode_response_body_bytes(
}
}
fn response_body_is_json(headers: &BTreeMap<String, String>, body_bytes: &[u8]) -> bool {
pub(crate) fn response_body_is_json(headers: &BTreeMap<String, String>, body_bytes: &[u8]) -> bool {
if headers
.get("content-type")
.map(|value| value.to_ascii_lowercase())

13
apps/aether-gateway/src/executor/candidate_loop.rs
View File

@@ -158,6 +158,19 @@ where
last_plan: aether_contracts::ExecutionPlan,
last_report_context: Option<serde_json::Value>,
) -> Result<Self::Exhaustion, Self::Error> {
warn!(
event_name = "candidate_loop_exhausted",
log_type = "ops",
trace_id = %self.trace_id,
plan_kind = self.plan_kind,
request_id = %short_request_id(last_plan.request_id.as_str()),
candidate_id = ?last_plan.candidate_id,
provider_name = last_plan.provider_name.as_deref().unwrap_or("-"),
endpoint_id = %last_plan.endpoint_id,
key_id = %last_plan.key_id,
model_name = last_plan.model_name.as_deref().unwrap_or("-"),
"candidate loop exhausted local sync candidates"
);
Ok(
build_local_execution_exhaustion(self.state, &last_plan, last_report_context.as_ref())
.await,

60
apps/aether-gateway/src/executor/orchestration.rs
View File

@@ -18,7 +18,8 @@ use crate::ai_serving::api::{
resolve_claude_stream_spec, resolve_claude_sync_spec, resolve_gemini_stream_spec,
resolve_gemini_sync_spec, resolve_local_same_format_stream_spec,
resolve_local_same_format_sync_spec, set_local_openai_chat_execution_exhausted_diagnostic,
AiStreamAttempt, AiSyncAttempt, LocalStandardSpec, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
set_local_openai_image_execution_exhausted_diagnostic, AiStreamAttempt, AiSyncAttempt,
LocalStandardSpec, EXECUTION_RUNTIME_STREAM_DECISION_ACTION,
EXECUTION_RUNTIME_SYNC_DECISION_ACTION,
};
use crate::control::GatewayControlDecision;
@@ -457,7 +458,7 @@ pub(crate) async fn maybe_execute_sync_via_local_image_decision(
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let Some((attempt_source, _candidate_count)) = build_local_image_sync_attempt_source_for_kind(
let Some((attempt_source, candidate_count)) = build_local_image_sync_attempt_source_for_kind(
state,
parts,
body_json,
@@ -471,7 +472,7 @@ pub(crate) async fn maybe_execute_sync_via_local_image_decision(
return Ok(LocalExecutionRequestOutcome::NoPath);
};
execute_sync_attempt_source::<AiSyncAttempt, _>(
let outcome = execute_sync_attempt_source::<AiSyncAttempt, _>(
state,
parts,
trace_id,
@@ -479,7 +480,20 @@ pub(crate) async fn maybe_execute_sync_via_local_image_decision(
plan_kind,
attempt_source,
)
.await
.await?;
if let LocalExecutionRequestOutcome::Exhausted(_) = &outcome {
set_local_openai_image_execution_exhausted_diagnostic(
state,
trace_id,
decision,
plan_kind,
body_json,
candidate_count,
);
}
Ok(outcome)
}
pub(crate) async fn maybe_execute_stream_via_local_gemini_files_decision(
@@ -517,29 +531,41 @@ pub(crate) async fn maybe_execute_stream_via_local_image_decision(
decision: &GatewayControlDecision,
plan_kind: &str,
) -> Result<LocalExecutionRequestOutcome, GatewayError> {
let Some((attempt_source, _candidate_count)) =
build_local_image_stream_attempt_source_for_kind(
state,
parts,
body_json,
body_base64,
trace_id,
decision,
plan_kind,
)
.await?
let Some((attempt_source, candidate_count)) = build_local_image_stream_attempt_source_for_kind(
state,
parts,
body_json,
body_base64,
trace_id,
decision,
plan_kind,
)
.await?
else {
return Ok(LocalExecutionRequestOutcome::NoPath);
};
execute_stream_attempt_source::<AiStreamAttempt, _>(
let outcome = execute_stream_attempt_source::<AiStreamAttempt, _>(
state,
trace_id,
decision,
plan_kind,
attempt_source,
)
.await
.await?;
if let LocalExecutionRequestOutcome::Exhausted(_) = &outcome {
set_local_openai_image_execution_exhausted_diagnostic(
state,
trace_id,
decision,
plan_kind,
body_json,
candidate_count,
);
}
Ok(outcome)
}
pub(crate) async fn maybe_execute_sync_via_local_video_decision(

55
apps/aether-gateway/src/handlers/admin/observability/usage/summary_routes.rs
View File

@@ -194,6 +194,57 @@ async fn resolve_admin_usage_attempt_flags_by_usage_id(
.collect())
}
async fn resolve_admin_usage_image_progress_by_request_id(
state: &AdminAppState<'_>,
items: &[StoredRequestUsageAudit],
) -> Result<BTreeMap<String, serde_json::Value>, GatewayError> {
if !state.has_request_candidate_data_reader() || items.is_empty() {
return Ok(BTreeMap::new());
}
let request_ids = items
.iter()
.map(|item| item.request_id.clone())
.collect::<BTreeSet<_>>();
let mut progress_by_request_id = BTreeMap::new();
for request_id in request_ids {
let candidates = state
.app()
.read_request_candidates_by_request_id(&request_id)
.await?;
if let Some(progress) = latest_admin_usage_image_progress(&candidates) {
progress_by_request_id.insert(request_id, progress);
}
}
Ok(progress_by_request_id)
}
fn latest_admin_usage_image_progress(
candidates: &[StoredRequestCandidate],
) -> Option<serde_json::Value> {
candidates
.iter()
.filter_map(|candidate| {
let progress = candidate
.extra_data
.as_ref()
.and_then(|value| value.get("image_progress"))?
.clone();
Some((
candidate
.started_at_unix_ms
.unwrap_or(candidate.created_at_unix_ms),
candidate.candidate_index,
candidate.retry_index,
progress,
))
})
.max_by_key(|(started_at, candidate_index, retry_index, _)| {
(*started_at, *candidate_index, *retry_index)
})
.map(|(_, _, _, progress)| progress)
}
fn admin_usage_matches_attempt_status(
item: &StoredRequestUsageAudit,
status: &str,
@@ -490,6 +541,7 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
&BTreeMap::new(),
state.has_auth_api_key_data_reader(),
&BTreeMap::new(),
&BTreeMap::new(),
)));
};
state
@@ -513,12 +565,15 @@ pub(super) async fn maybe_build_local_admin_usage_summary_response(
};
let api_key_names = admin_usage_api_key_names(state, &items).await?;
let provider_key_names = admin_usage_provider_key_names(state, &items).await?;
let image_progress_by_request_id =
resolve_admin_usage_image_progress_by_request_id(state, &items).await?;
return Ok(Some(build_admin_usage_active_requests_response(
&items,
&api_key_names,
state.has_auth_api_key_data_reader(),
&provider_key_names,
&image_progress_by_request_id,
)));
}
Some("records")

41
apps/aether-gateway/src/request_candidate_runtime.rs
View File

@@ -274,6 +274,47 @@ pub(crate) async fn record_local_request_candidate_status(
persist_local_request_candidate_status_record(state, record).await;
}
pub(crate) async fn record_local_request_candidate_extra_data(
state: &(impl RequestCandidateRuntimeWriter + ?Sized),
plan: &ExecutionPlan,
report_context: Option<&Value>,
status: RequestCandidateStatus,
status_code: Option<u16>,
latency_ms: Option<u64>,
extra_data: Value,
) {
let Some(snapshot) = snapshot_local_request_candidate_status(plan, report_context) else {
return;
};
let record = UpsertRequestCandidateRecord {
id: snapshot.candidate_id.clone(),
request_id: snapshot.request_id.clone(),
user_id: snapshot.user_id.clone(),
api_key_id: snapshot.api_key_id.clone(),
username: None,
api_key_name: None,
candidate_index: snapshot.candidate_index,
retry_index: snapshot.retry_index,
provider_id: Some(snapshot.provider_id.clone()),
endpoint_id: Some(snapshot.endpoint_id.clone()),
key_id: Some(snapshot.key_id.clone()),
status,
skip_reason: None,
is_cached: None,
status_code,
error_type: None,
error_message: None,
latency_ms,
concurrent_requests: None,
extra_data: Some(extra_data),
required_capabilities: None,
created_at_unix_ms: None,
started_at_unix_ms: None,
finished_at_unix_ms: None,
};
persist_local_request_candidate_status_record(state, record).await;
}
pub(crate) async fn record_local_request_candidate_status_snapshot(
state: &(impl RequestCandidateRuntimeWriter + ?Sized),
snapshot: &LocalRequestCandidateStatusSnapshot,

7
apps/aether-gateway/src/tests/ai_execute/lifecycle.rs
View File

@@ -20,6 +20,7 @@ use aether_data_contracts::repository::provider_catalog::{
use sha2::{Digest, Sha256};
use crate::data::GatewayDataState;
use crate::tests::next_non_keepalive_chunk;
fn hash_api_key(value: &str) -> String {
let mut hasher = Sha256::new();
@@ -370,11 +371,7 @@ async fn gateway_stops_execution_runtime_stream_when_client_disconnects() {
assert_eq!(response.status(), StatusCode::OK);
let mut response = response;
let first_chunk = response
.chunk()
.await
.expect("first chunk should read")
.expect("first chunk should exist");
let first_chunk = next_non_keepalive_chunk(&mut response).await;
assert_eq!(
first_chunk,
Bytes::from_static(b"data: {\"id\":\"chatcmpl-first\"}\n\n")

7
apps/aether-gateway/src/tests/ai_execute/mod.rs
View File

@@ -18,9 +18,10 @@ use crate::constants::{
use super::{
build_router, build_router_with_execution_runtime_override, build_router_with_state,
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
UsageRuntimeConfig, VideoTaskTruthSourceMode,
build_state_with_execution_runtime_override, next_non_keepalive_chunk, start_server,
strip_sse_keepalive_comments, wait_until, AppState, FrontdoorCorsConfig,
FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason, UsageRuntimeConfig,
VideoTaskTruthSourceMode,
};
mod control_execute;

12
apps/aether-gateway/src/tests/ai_execute/stream/decision.rs
View File

@@ -1,8 +1,8 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, AppState, Arc, Body, Bytes, HeaderName, HeaderValue, Json, Mutex, Request, Response,
Router, StatusCode, EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_HEADER,
TRACE_ID_HEADER,
strip_sse_keepalive_comments, to_bytes, AppState, Arc, Body, Bytes, HeaderName, HeaderValue,
Json, Mutex, Request, Response, Router, StatusCode, EXECUTION_PATH_EXECUTION_RUNTIME_STREAM,
EXECUTION_PATH_HEADER, TRACE_ID_HEADER,
};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
@@ -391,7 +391,7 @@ async fn gateway_executes_openai_chat_stream_via_local_decision_gate_without_exe
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"id\":\"chatcmpl-local-123\"}\n\ndata: [DONE]\n\n"
);
@@ -1352,7 +1352,7 @@ async fn gateway_executes_openai_chat_stream_with_custom_path_via_local_decision
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"id\":\"chatcmpl-local-custom-path-123\"}\n\ndata: [DONE]\n\n"
);
@@ -1863,7 +1863,7 @@ async fn gateway_retries_next_local_openai_chat_stream_candidate_after_retryable
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"id\":\"chatcmpl-local-stream-failover-123\"}\n\ndata: [DONE]\n\n"
);

7
apps/aether-gateway/src/tests/ai_execute/stream/mod.rs
View File

@@ -18,9 +18,10 @@ use crate::constants::{
use super::{
build_router, build_router_with_execution_runtime_override, build_router_with_state,
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
UsageRuntimeConfig, VideoTaskTruthSourceMode,
build_state_with_execution_runtime_override, next_non_keepalive_chunk, start_server,
strip_sse_keepalive_comments, wait_until, AppState, FrontdoorCorsConfig,
FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason, UsageRuntimeConfig,
VideoTaskTruthSourceMode,
};
mod decision;

6
apps/aether-gateway/src/tests/ai_execute/stream_cli/direct.rs
View File

@@ -1,7 +1,7 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Infallible, Json, Mutex, Request,
Response, Router, StatusCode, UsageRuntimeConfig, TRACE_ID_HEADER,
strip_sse_keepalive_comments, to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Infallible,
Json, Mutex, Request, Response, Router, StatusCode, UsageRuntimeConfig, TRACE_ID_HEADER,
};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
@@ -469,7 +469,7 @@ async fn gateway_executes_codex_cli_stream_via_local_decision_gate_after_oauth_r
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"event: response.completed\ndata: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_codex_cli_stream_local_123\",\"object\":\"response\",\"model\":\"gpt-5.4\",\"status\":\"completed\",\"usage\":{\"input_tokens\":1,\"output_tokens\":2,\"total_tokens\":3}}}\n\n"
);

6
apps/aether-gateway/src/tests/ai_execute/stream_cli/mod.rs
View File

@@ -18,9 +18,9 @@ use crate::constants::{
use super::{
build_router, build_router_with_execution_runtime_override, build_router_with_state,
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
UsageRuntimeConfig, VideoTaskTruthSourceMode,
build_state_with_execution_runtime_override, start_server, strip_sse_keepalive_comments,
wait_until, AppState, FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind,
GatewayFallbackReason, UsageRuntimeConfig, VideoTaskTruthSourceMode,
};
mod compact;

22
apps/aether-gateway/src/tests/ai_execute/stream_provider.rs
View File

@@ -1,7 +1,8 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, json, start_server,
to_bytes, Arc, Body, Bytes, HeaderName, HeaderValue, Json, Mutex, Request, Response, Router,
StatusCode, TRACE_ID_HEADER,
any, build_router_with_state, build_state_with_execution_runtime_override, json,
next_non_keepalive_chunk, start_server, strip_sse_keepalive_comments, to_bytes, Arc, Body,
Bytes, HeaderName, HeaderValue, Json, Mutex, Request, Response, Router, StatusCode,
TRACE_ID_HEADER,
};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{
@@ -952,11 +953,12 @@ async fn gateway_executes_claude_cli_stream_via_local_decision_gate_without_wait
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
tokio::time::timeout(std::time::Duration::from_millis(100), response.chunk())
.await
.expect("same-format passthrough should yield first chunk before eof")
.expect("first chunk should read")
.expect("first chunk should exist"),
tokio::time::timeout(
std::time::Duration::from_millis(100),
next_non_keepalive_chunk(&mut response),
)
.await
.expect("same-format passthrough should yield first chunk before eof"),
Bytes::from_static(b"event: message_start\ndata: {\"type\":\"message_start\"}\n\n")
);
assert_eq!(
@@ -1470,7 +1472,7 @@ async fn gateway_executes_claude_code_cli_stream_via_local_decision_gate_with_lo
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"event: message_start\ndata: {\"type\":\"message_start\"}\n\n"
);
@@ -1924,7 +1926,7 @@ async fn gateway_executes_claude_chat_stream_via_local_decision_gate_with_local_
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"event: message_start\ndata: {\"type\":\"message_start\"}\n\n"
);

6
apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_chat.rs
View File

@@ -1,7 +1,7 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override,
encrypt_python_fernet_plaintext, json, start_server, to_bytes, Arc, Body, Bytes, Digest,
HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
encrypt_python_fernet_plaintext, json, start_server, strip_sse_keepalive_comments, to_bytes,
Arc, Body, Bytes, Digest, HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
InMemoryRequestCandidateRepository, Json, Mutex, Request, RequestCandidateReadRepository,
RequestCandidateStatus, Response, Router, Sha256, StatusCode, StoredAuthApiKeySnapshot,
@@ -404,7 +404,7 @@ async fn gateway_executes_gemini_chat_stream_via_local_decision_gate_with_local_
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"candidates\":[]}\n\n"
);

13
apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/local_cli.rs
View File

@@ -1,7 +1,7 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override,
encrypt_python_fernet_plaintext, json, start_server, to_bytes, Arc, Body, Bytes, Digest,
HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
encrypt_python_fernet_plaintext, json, start_server, strip_sse_keepalive_comments, to_bytes,
Arc, Body, Bytes, Digest, HeaderName, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
InMemoryRequestCandidateRepository, Json, Mutex, Request, RequestCandidateReadRepository,
RequestCandidateStatus, Response, Router, Sha256, StatusCode, StoredAuthApiKeySnapshot,
@@ -381,7 +381,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_with_local_s
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"candidates\":[]}\n\n"
);
@@ -872,7 +872,7 @@ async fn gateway_executes_gemini_cli_stream_via_local_decision_gate_after_oauth_
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"candidates\":[]}\n\n"
);
@@ -1339,7 +1339,7 @@ async fn gateway_executes_vertex_ai_gemini_cli_stream_via_local_decision_gate_wi
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: {\"candidates\":[]}\n\n"
);
@@ -1849,7 +1849,8 @@ async fn gateway_executes_antigravity_gemini_cli_stream_via_local_decision_gate_
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let response_text = response.text().await.expect("body should read");
let response_text =
strip_sse_keepalive_comments(&response.text().await.expect("body should read"));
let payload = response_text
.trim()
.strip_prefix("data: ")

6
apps/aether-gateway/src/tests/ai_execute/stream_provider_gemini/mod.rs
View File

@@ -18,9 +18,9 @@ use crate::constants::{
use super::{
build_router, build_router_with_execution_runtime_override, build_router_with_state,
build_state_with_execution_runtime_override, start_server, wait_until, AppState,
FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind, GatewayFallbackReason,
UsageRuntimeConfig, VideoTaskTruthSourceMode,
build_state_with_execution_runtime_override, start_server, strip_sse_keepalive_comments,
wait_until, AppState, FrontdoorCorsConfig, FrontdoorUserRpmConfig, GatewayFallbackMetricKind,
GatewayFallbackReason, UsageRuntimeConfig, VideoTaskTruthSourceMode,
};
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
use aether_data::repository::auth::{

6
apps/aether-gateway/src/tests/frontdoor/internal.rs
View File

@@ -6,8 +6,8 @@ use super::{
};
use crate::tests::{
any, build_router, build_router_with_state, build_state_with_execution_runtime_override, json,
start_server, AppState, Arc, Body, HeaderValue, Json, Mutex, Request, Response, Router,
StatusCode, CONTROL_ACTION_PROXY_PUBLIC, CONTROL_EXECUTED_HEADER,
start_server, strip_sse_keepalive_comments, AppState, Arc, Body, HeaderValue, Json, Mutex,
Request, Response, Router, StatusCode, CONTROL_ACTION_PROXY_PUBLIC, CONTROL_EXECUTED_HEADER,
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
EXECUTION_PATH_HEADER,
};
@@ -458,7 +458,7 @@ async fn gateway_handles_internal_gateway_execute_stream_locally() {
Some(EXECUTION_PATH_EXECUTION_RUNTIME_STREAM)
);
assert_eq!(
response.text().await.expect("body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("body should read")),
"data: one\n\ndata: [DONE]\n\n"
);
assert_eq!(*upstream_hits.lock().expect("mutex should lock"), 0);

17
apps/aether-gateway/src/tests/mod.rs
View File

@@ -88,3 +88,20 @@ pub(super) async fn wait_until(timeout_ms: u64, mut predicate: impl FnMut() -> b
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
}
pub(crate) fn strip_sse_keepalive_comments(body: &str) -> String {
body.replace(": aether-keepalive\n\n", "")
}
pub(crate) async fn next_non_keepalive_chunk(response: &mut reqwest::Response) -> Bytes {
loop {
let chunk = response
.chunk()
.await
.expect("chunk should read")
.expect("chunk should exist");
if chunk.as_ref() != b": aether-keepalive\n\n" {
return chunk;
}
}
}

4
apps/aether-gateway/src/tests/usage.rs
View File

@@ -28,8 +28,8 @@ use sha2::{Digest, Sha256};
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override, send_request,
start_server, Body, HeaderValue, Json, Mutex, Request, Response, Router, StatusCode,
UsageRuntimeConfig, TRACE_ID_HEADER,
start_server, strip_sse_keepalive_comments, Body, HeaderValue, Json, Mutex, Request, Response,
Router, StatusCode, UsageRuntimeConfig, TRACE_ID_HEADER,
};
use crate::data::GatewayDataState;

105
apps/aether-gateway/src/tests/usage/local.rs
View File

@@ -2,8 +2,8 @@ use super::{
any, build_router_with_state, build_state_with_execution_runtime_override,
encrypt_python_fernet_plaintext, hash_api_key, json, sample_local_openai_auth_snapshot,
sample_local_openai_candidate_row, sample_local_openai_endpoint, sample_local_openai_key,
sample_local_openai_provider, send_request, start_server, Arc, Body, GatewayDataState,
HeaderValue, InMemoryAuthApiKeySnapshotRepository,
sample_local_openai_provider, send_request, start_server, strip_sse_keepalive_comments, Arc,
Body, GatewayDataState, HeaderValue, InMemoryAuthApiKeySnapshotRepository,
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
InMemoryRequestCandidateRepository, InMemoryUsageReadRepository, Json, Mutex, Request,
RequestCandidateReadRepository, RequestCandidateStatus, Response, Router, StatusCode,
@@ -840,6 +840,104 @@ async fn gateway_records_failed_usage_when_all_local_openai_chat_candidates_exha
assert_eq!(stored_candidates[0].status_code, Some(503));
}
#[tokio::test]
async fn gateway_records_failed_usage_when_sync_runtime_transport_is_unavailable_without_plan_fallback(
) {
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let execution_hits = Arc::new(Mutex::new(0usize));
let execution_hits_clone = Arc::clone(&execution_hits);
let auth_repository = Arc::new(InMemoryAuthApiKeySnapshotRepository::seed(vec![(
Some(hash_api_key("sk-client-openai-local-transport-unavailable")),
sample_local_openai_auth_snapshot(
"api-key-openai-usage-local-transport-unavailable-1",
"user-openai-usage-local-transport-unavailable-1",
),
)]));
let candidate_selection_repository =
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(vec![
sample_local_openai_candidate_row(),
]));
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_local_openai_provider()],
vec![sample_local_openai_endpoint()],
vec![sample_local_openai_key()],
));
let gateway_state = crate::AppState::new()
.expect("gateway should build")
.with_execution_runtime_sync_override_for_tests(move |_plan| {
*execution_hits_clone.lock().expect("mutex should lock") += 1;
Err(crate::GatewayError::Internal(
"simulated transport unavailable".to_string(),
))
})
.with_data_state_for_tests(
GatewayDataState::with_auth_candidate_selection_provider_catalog_request_candidates_and_usage_for_tests(
auth_repository,
candidate_selection_repository,
provider_catalog_repository,
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
DEVELOPMENT_ENCRYPTION_KEY,
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
});
let gateway = build_router_with_state(gateway_state);
let request = Request::builder()
.method(http::Method::POST)
.uri("/v1/chat/completions")
.header(http::header::CONTENT_TYPE, "application/json")
.header(
http::header::AUTHORIZATION,
"Bearer sk-client-openai-local-transport-unavailable",
)
.header(
TRACE_ID_HEADER,
"trace-openai-chat-local-transport-unavailable-123",
)
.body(Body::from("{\"model\":\"gpt-5\",\"messages\":[]}"))
.expect("request should build");
let response = send_request(gateway, request).await;
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(*execution_hits.lock().expect("mutex should lock"), 1);
let stored_usage = wait_for_usage_status(
usage_repository.as_ref(),
"trace-openai-chat-local-transport-unavailable-123",
"failed",
)
.await;
assert_eq!(stored_usage.status, "failed");
assert_eq!(stored_usage.billing_status, "void");
assert_eq!(stored_usage.status_code, Some(503));
assert_eq!(
stored_usage
.response_body
.as_ref()
.and_then(|value| value.get("error"))
.and_then(|value| value.get("type"))
.and_then(|value| value.as_str()),
Some("execution_runtime_unavailable")
);
let stored_candidates = request_candidate_repository
.list_by_request_id("trace-openai-chat-local-transport-unavailable-123")
.await
.expect("request candidate trace should read");
assert_eq!(stored_candidates.len(), 1);
assert_eq!(stored_candidates[0].status, RequestCandidateStatus::Failed);
assert_eq!(
stored_candidates[0].error_type.as_deref(),
Some("execution_runtime_unavailable")
);
}
#[test]
fn gateway_records_failed_usage_for_claude_runtime_miss_without_execution_exhaustion() {
run_async_test_on_large_stack(
@@ -1184,7 +1282,8 @@ async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_wh
.expect("request should succeed");
assert_eq!(response.status(), StatusCode::OK);
let body_text = response.text().await.expect("stream body should read");
let body_text =
strip_sse_keepalive_comments(&response.text().await.expect("stream body should read"));
assert_eq!(
body_text,
"data: {\"id\":\"chatcmpl-local-report-stream-123\",\"usage\":{\"input_tokens\":2,\"output_tokens\":4,\"total_tokens\":6}}\n\ndata: [DONE]\n\n"

23
apps/aether-gateway/src/tests/usage/pricing.rs
View File

@@ -1,13 +1,14 @@
use super::{
any, build_router_with_state, build_state_with_execution_runtime_override,
encrypt_python_fernet_plaintext, hash_api_key, json, start_server, Arc, Body, GatewayDataState,
HeaderValue, InMemoryAuthApiKeySnapshotRepository,
InMemoryMinimalCandidateSelectionReadRepository, InMemoryProviderCatalogReadRepository,
InMemoryRequestCandidateRepository, InMemoryUsageReadRepository, Json, Request,
RequestCandidateReadRepository, RequestCandidateStatus, Response, Router, StatusCode,
StoredAuthApiKeySnapshot, StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogProvider, StoredProviderModelMapping,
UsageReadRepository, UsageRuntimeConfig, DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
encrypt_python_fernet_plaintext, hash_api_key, json, start_server,
strip_sse_keepalive_comments, Arc, Body, GatewayDataState, HeaderValue,
InMemoryAuthApiKeySnapshotRepository, InMemoryMinimalCandidateSelectionReadRepository,
InMemoryProviderCatalogReadRepository, InMemoryRequestCandidateRepository,
InMemoryUsageReadRepository, Json, Request, RequestCandidateReadRepository,
RequestCandidateStatus, Response, Router, StatusCode, StoredAuthApiKeySnapshot,
StoredMinimalCandidateSelectionRow, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogProvider, StoredProviderModelMapping, UsageReadRepository,
UsageRuntimeConfig, DEVELOPMENT_ENCRYPTION_KEY, TRACE_ID_HEADER,
};
use aether_data::repository::billing::InMemoryBillingReadRepository;
use aether_data::repository::wallet::{InMemoryWalletRepository, StoredWalletSnapshot};
@@ -939,7 +940,7 @@ async fn gateway_records_openai_stream_usage_and_pricing_with_cache_tokens_impl(
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("stream body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("stream body should read")),
stream_body.concat()
);
@@ -1132,7 +1133,7 @@ async fn gateway_records_claude_stream_usage_and_pricing_with_cache_breakdown_im
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("stream body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("stream body should read")),
stream_body.concat()
);
@@ -1310,7 +1311,7 @@ async fn gateway_records_gemini_stream_usage_and_pricing_with_cache_read_tokens_
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.text().await.expect("stream body should read"),
strip_sse_keepalive_comments(&response.text().await.expect("stream body should read")),
stream_body.concat()
);