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

fix(usage): record visible stream first byte timing

Browse Source
This commit is contained in:
fawney19
2026-05-27 02:48:01 +08:00
parent f03550415b
commit cf8372c8cb
10 changed files with 656 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/execution_runtime/stream/execution.rs
+568 -50
View File
@@ -17,7 +17,7 @@ use aether_scheduler_core::{
};
use aether_usage_runtime::{
build_lifecycle_usage_seed, build_stream_terminal_usage_payload_seed,
build_sync_terminal_usage_payload_seed, build_terminal_usage_context_seed,
build_sync_terminal_usage_payload_seed, build_terminal_usage_context_seed, LifecycleUsageSeed,
UsageBodyCapturePolicy, UsageRequestRecordLevel, UsageRuntimeAccess,
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
};
@@ -605,6 +605,209 @@ fn append_stream_capture_bytes(
}
}
const CLIENT_VISIBLE_TEXT_TRACKER_MAX_BUFFER_BYTES: usize = 256 * 1024;
struct ClientVisibleTextTracker {
client_api_format: String,
buffered: Vec<u8>,
observed: bool,
}
impl ClientVisibleTextTracker {
fn new(client_api_format: &str) -> Self {
Self {
client_api_format: client_api_format.trim().to_ascii_lowercase(),
buffered: Vec::new(),
observed: false,
}
}
fn observe_chunk(&mut self, chunk: &[u8]) -> bool {
if self.observed || chunk.is_empty() {
return false;
}
self.buffered.extend_from_slice(chunk);
while let Some((block_end, separator_len)) = find_sse_block_boundary(&self.buffered) {
let block_len = block_end + separator_len;
let block = self.buffered.drain(..block_len).collect::<Vec<_>>();
if sse_block_has_client_visible_text(&self.client_api_format, &block) {
self.observed = true;
self.buffered.clear();
return true;
}
}
if self.buffered.len() > CLIENT_VISIBLE_TEXT_TRACKER_MAX_BUFFER_BYTES {
self.buffered.clear();
}
false
}
}
fn sse_block_has_client_visible_text(client_api_format: &str, block: &[u8]) -> bool {
let text = String::from_utf8_lossy(block);
let mut data_lines = Vec::new();
for raw_line in text.lines() {
let line = raw_line.trim_end_matches('\r');
let Some(data) = line.strip_prefix("data:") else {
continue;
};
data_lines.push(data.strip_prefix(' ').unwrap_or(data));
}
if data_lines.is_empty() {
return false;
}
let data = data_lines.join("\n");
let data = data.trim();
if data.is_empty() || data == "[DONE]" {
return false;
}
let Ok(value) = serde_json::from_str::<Value>(data) else {
return false;
};
stream_json_has_client_visible_text(client_api_format, &value)
}
fn stream_json_has_client_visible_text(client_api_format: &str, value: &Value) -> bool {
if client_api_format.contains("openai:chat") || client_api_format.contains("openai_chat") {
return openai_chat_stream_json_has_visible_text(value);
}
if client_api_format.contains("openai:responses")
|| client_api_format.contains("openai_responses")
{
return openai_responses_stream_json_has_visible_text(value);
}
if client_api_format.contains("claude") {
return claude_stream_json_has_visible_text(value);
}
if client_api_format.contains("gemini") {
return gemini_stream_json_has_visible_text(value);
}
openai_chat_stream_json_has_visible_text(value)
|| openai_responses_stream_json_has_visible_text(value)
|| claude_stream_json_has_visible_text(value)
|| gemini_stream_json_has_visible_text(value)
}
fn visible_text_value(value: Option<&Value>) -> bool {
value
.and_then(Value::as_str)
.is_some_and(|text| text.chars().any(|ch| !ch.is_whitespace()))
}
fn openai_chat_stream_json_has_visible_text(value: &Value) -> bool {
value
.get("choices")
.and_then(Value::as_array)
.is_some_and(|choices| {
choices.iter().any(|choice| {
choice
.get("delta")
.and_then(Value::as_object)
.is_some_and(|delta| {
visible_text_value(delta.get("content"))
|| visible_text_value(delta.get("refusal"))
})
|| visible_text_value(choice.get("text"))
})
})
}
fn openai_responses_stream_json_has_visible_text(value: &Value) -> bool {
match value
.get("type")
.and_then(Value::as_str)
.unwrap_or_default()
{
"response.output_text.delta" => visible_text_value(value.get("delta")),
"response.output_text.done" => visible_text_value(value.get("text")),
"response.completed" => value
.get("response")
.is_some_and(openai_response_body_has_visible_text),
_ => false,
}
}
fn openai_response_body_has_visible_text(value: &Value) -> bool {
value
.get("output")
.and_then(Value::as_array)
.is_some_and(|items| {
items.iter().any(|item| {
item.get("content")
.and_then(Value::as_array)
.is_some_and(|content| {
content.iter().any(|part| {
part.get("type")
.and_then(Value::as_str)
.is_some_and(|kind| matches!(kind, "output_text" | "text"))
&& visible_text_value(part.get("text"))
})
})
})
})
}
fn claude_stream_json_has_visible_text(value: &Value) -> bool {
match value
.get("type")
.and_then(Value::as_str)
.unwrap_or_default()
{
"content_block_delta" => {
value
.get("delta")
.and_then(Value::as_object)
.is_some_and(|delta| {
delta.get("type").and_then(Value::as_str) == Some("text_delta")
&& visible_text_value(delta.get("text"))
})
}
"content_block_start" => value
.get("content_block")
.and_then(Value::as_object)
.is_some_and(|block| {
block.get("type").and_then(Value::as_str) == Some("text")
&& visible_text_value(block.get("text"))
}),
_ => false,
}
}
fn gemini_stream_json_has_visible_text(value: &Value) -> bool {
let event = value
.get("response")
.filter(|response| response.get("candidates").is_some())
.unwrap_or(value);
event
.get("candidates")
.and_then(Value::as_array)
.is_some_and(|candidates| {
candidates.iter().any(|candidate| {
candidate
.get("content")
.and_then(|content| content.get("parts"))
.and_then(Value::as_array)
.is_some_and(|parts| {
parts.iter().any(|part| {
!part
.get("thought")
.and_then(Value::as_bool)
.unwrap_or(false)
&& visible_text_value(part.get("text"))
})
})
})
})
}
fn observe_stream_usage_bytes(
observer: &mut StreamingStandardTerminalObserver,
report_context: &Value,
@@ -1736,19 +1939,72 @@ fn should_refresh_stream_usage_telemetry(
|| (next_elapsed.is_some() && next_elapsed != previous_elapsed)
}
fn stream_elapsed_ms_since(started_at: Instant) -> u64 {
started_at.elapsed().as_millis().min(u128::from(u64::MAX)) as u64
}
fn first_visible_text_telemetry(
stream_started_at: Instant,
upstream_telemetry: Option<&ExecutionTelemetry>,
) -> ExecutionTelemetry {
let elapsed_ms = stream_elapsed_ms_since(stream_started_at);
ExecutionTelemetry {
ttfb_ms: Some(elapsed_ms),
elapsed_ms: Some(elapsed_ms),
upstream_bytes: upstream_telemetry.and_then(|telemetry| telemetry.upstream_bytes),
}
}
fn usage_refresh_telemetry(
upstream_telemetry: &ExecutionTelemetry,
usage_stream_telemetry: Option<&ExecutionTelemetry>,
) -> ExecutionTelemetry {
ExecutionTelemetry {
ttfb_ms: usage_stream_telemetry.and_then(|telemetry| telemetry.ttfb_ms),
elapsed_ms: upstream_telemetry.elapsed_ms,
upstream_bytes: upstream_telemetry.upstream_bytes,
}
}
fn maybe_record_first_visible_text_started(
state: &AppState,
lifecycle_seed: &LifecycleUsageSeed,
status_code: u16,
tracker: &mut ClientVisibleTextTracker,
chunk: &[u8],
stream_started_at: Instant,
upstream_telemetry: Option<&ExecutionTelemetry>,
usage_stream_telemetry: &mut Option<ExecutionTelemetry>,
) {
if usage_stream_telemetry
.as_ref()
.and_then(|telemetry| telemetry.ttfb_ms)
.is_some()
{
return;
}
if !tracker.observe_chunk(chunk) {
return;
}
let telemetry = first_visible_text_telemetry(stream_started_at, upstream_telemetry);
state.usage_runtime.record_stream_started(
state.data.as_ref(),
lifecycle_seed,
status_code,
Some(&telemetry),
);
*usage_stream_telemetry = Some(telemetry);
}
fn build_terminal_stream_telemetry(
stream_started_at: Instant,
telemetry: Option<&ExecutionTelemetry>,
usage_stream_telemetry: Option<&ExecutionTelemetry>,
upstream_bytes: u64,
) -> ExecutionTelemetry {
let current_elapsed_ms = stream_started_at
.elapsed()
.as_millis()
.min(u128::from(u64::MAX)) as u64;
let ttfb_ms = telemetry
.and_then(|telemetry| telemetry.ttfb_ms)
.or_else(|| usage_stream_telemetry.and_then(|telemetry| telemetry.ttfb_ms));
let current_elapsed_ms = stream_elapsed_ms_since(stream_started_at);
let ttfb_ms = usage_stream_telemetry.and_then(|telemetry| telemetry.ttfb_ms);
let prior_elapsed_ms = telemetry
.and_then(|telemetry| telemetry.elapsed_ms)
.or_else(|| usage_stream_telemetry.and_then(|telemetry| telemetry.elapsed_ms))
@@ -2279,6 +2535,9 @@ async fn execute_stream_from_frame_stream(
let mut prefetched_body = Vec::new();
let mut prefetched_inspection_body = Vec::new();
let mut prefetched_telemetry: Option<ExecutionTelemetry> = None;
let mut prefetched_usage_telemetry: Option<ExecutionTelemetry> = None;
let mut prefetched_client_visible_text_tracker =
ClientVisibleTextTracker::new(plan.client_api_format.as_str());
let mut reached_eof = false;
let mut sync_json_stream_bridge_active = false;
if skip_direct_finalize_prefetch {
@@ -2356,7 +2615,7 @@ async fn execute_stream_from_frame_stream(
candidate_id,
report_kind,
headers,
prefetched_telemetry,
prefetched_usage_telemetry.clone(),
&provider_prefetched_body,
failure,
)
@@ -2389,7 +2648,7 @@ async fn execute_stream_from_frame_stream(
candidate_id,
report_kind,
headers,
prefetched_telemetry,
prefetched_usage_telemetry.clone(),
&prefetched_body,
failure,
)
@@ -2420,7 +2679,7 @@ async fn execute_stream_from_frame_stream(
candidate_id,
report_kind,
headers,
prefetched_telemetry,
prefetched_usage_telemetry.clone(),
&provider_prefetched_body,
error_status_code,
error_body_json,
@@ -2458,7 +2717,7 @@ async fn execute_stream_from_frame_stream(
headers,
Some(body_json),
None,
prefetched_telemetry,
prefetched_usage_telemetry.clone(),
);
record_sync_terminal_usage(
state,
@@ -2500,6 +2759,19 @@ async fn execute_stream_from_frame_stream(
"text/event-stream".to_string(),
);
stream_terminal_summary = outcome.terminal_summary;
if prefetched_usage_telemetry
.as_ref()
.and_then(|telemetry| telemetry.ttfb_ms)
.is_none()
&& prefetched_client_visible_text_tracker
.observe_chunk(&outcome.sse_body)
{
prefetched_usage_telemetry =
Some(first_visible_text_telemetry(
stream_started_at,
prefetched_telemetry.as_ref(),
));
}
prefetched_body.extend_from_slice(&outcome.sse_body);
prefetched_chunks.push(Bytes::from(outcome.sse_body));
sync_json_stream_bridge_active = true;
@@ -2524,7 +2796,7 @@ async fn execute_stream_from_frame_stream(
candidate_id,
report_kind,
headers,
prefetched_telemetry,
prefetched_usage_telemetry.clone(),
&provider_prefetched_body,
failure,
)
@@ -2557,7 +2829,7 @@ async fn execute_stream_from_frame_stream(
candidate_id,
report_kind,
headers,
prefetched_telemetry,
prefetched_usage_telemetry.clone(),
&provider_prefetched_body,
failure,
)
@@ -2588,7 +2860,7 @@ async fn execute_stream_from_frame_stream(
candidate_id,
report_kind,
headers,
prefetched_telemetry,
prefetched_usage_telemetry.clone(),
&provider_prefetched_body,
failure,
)
@@ -2599,6 +2871,18 @@ async fn execute_stream_from_frame_stream(
normalized_chunk
};
if !rewritten_chunk.is_empty() {
if prefetched_usage_telemetry
.as_ref()
.and_then(|telemetry| telemetry.ttfb_ms)
.is_none()
&& prefetched_client_visible_text_tracker
.observe_chunk(&rewritten_chunk)
{
prefetched_usage_telemetry = Some(first_visible_text_telemetry(
stream_started_at,
prefetched_telemetry.as_ref(),
));
}
prefetched_body.extend_from_slice(&rewritten_chunk);
prefetched_chunks.push(Bytes::from(rewritten_chunk));
}
@@ -2639,7 +2923,7 @@ async fn execute_stream_from_frame_stream(
candidate_id,
report_kind,
headers,
prefetched_telemetry,
prefetched_usage_telemetry.clone(),
&provider_prefetched_body,
build_stream_failure_from_execution_error(&error),
)
@@ -2652,11 +2936,16 @@ async fn execute_stream_from_frame_stream(
drop(private_stream_normalizer);
drop(local_stream_rewriter);
let initial_usage_telemetry = prefetched_usage_telemetry.clone().or_else(|| {
prefetched_telemetry
.as_ref()
.map(|telemetry| usage_refresh_telemetry(telemetry, None))
});
state.usage_runtime.record_stream_started(
state.data.as_ref(),
&lifecycle_seed,
status_code,
prefetched_telemetry.as_ref(),
initial_usage_telemetry.as_ref(),
);
if let Some(snapshot) = request_candidate_status_snapshot {
let state_bg = state.clone();
@@ -2698,6 +2987,7 @@ async fn execute_stream_from_frame_stream(
let prefetched_chunks_for_body = prefetched_chunks;
let sync_json_stream_bridge_active_for_report = sync_json_stream_bridge_active;
let initial_telemetry = prefetched_telemetry;
let initial_client_visible_text_tracker = prefetched_client_visible_text_tracker;
let initial_reached_eof = reached_eof;
let direct_stream_finalize_kind_owned = direct_stream_finalize_kind;
let candidate_started_unix_secs_for_report = candidate_started_unix_secs;
@@ -2787,7 +3077,8 @@ async fn execute_stream_from_frame_stream(
let mut client_stream_completion_tracker = ClientVisibleStreamCompletionTracker::default();
let mut client_visible_stream_completed =
client_stream_completion_tracker.observe_chunk(&prefetched_body_for_report);
let mut usage_stream_telemetry: Option<ExecutionTelemetry> = initial_telemetry.clone();
let mut client_visible_text_tracker = initial_client_visible_text_tracker;
let mut usage_stream_telemetry: Option<ExecutionTelemetry> = initial_usage_telemetry;
let mut telemetry: Option<ExecutionTelemetry> = initial_telemetry;
let reached_eof = initial_reached_eof;
let mut downstream_dropped = false;
@@ -3136,32 +3427,6 @@ async fn execute_stream_from_frame_stream(
continue;
}
if usage_stream_telemetry
.as_ref()
.and_then(|telemetry| telemetry.ttfb_ms)
.is_none()
{
let first_data_elapsed_ms = stream_started_at_for_report
.elapsed()
.as_millis()
.min(u128::from(u64::MAX))
as u64;
let first_data_telemetry = ExecutionTelemetry {
ttfb_ms: Some(first_data_elapsed_ms),
elapsed_ms: Some(first_data_elapsed_ms),
upstream_bytes: telemetry
.as_ref()
.and_then(|telemetry| telemetry.upstream_bytes),
};
state_for_report.usage_runtime.record_stream_started(
state_for_report.data.as_ref(),
&lifecycle_seed_for_report,
status_code,
Some(&first_data_telemetry),
);
usage_stream_telemetry = Some(first_data_telemetry);
}
append_stream_capture_bytes(
&mut buffered_body,
&rewritten_chunk,
@@ -3185,6 +3450,16 @@ async fn execute_stream_from_frame_stream(
);
downstream_dropped = true;
} else {
maybe_record_first_visible_text_started(
&state_for_report,
&lifecycle_seed_for_report,
status_code,
&mut client_visible_text_tracker,
rewritten_chunk.as_ref(),
stream_started_at_for_report,
telemetry.as_ref(),
&mut usage_stream_telemetry,
);
client_visible_stream_completed |= client_stream_completion_tracker
.observe_chunk(rewritten_chunk.as_ref());
client_stream_bytes.fetch_add(rewritten_chunk_len, Ordering::Relaxed);
@@ -3210,18 +3485,22 @@ async fn execute_stream_from_frame_stream(
StreamFramePayload::Telemetry {
telemetry: frame_telemetry,
} => {
let usage_frame_telemetry = usage_refresh_telemetry(
&frame_telemetry,
usage_stream_telemetry.as_ref(),
);
let should_refresh_stream_usage = should_refresh_stream_usage_telemetry(
usage_stream_telemetry.as_ref(),
&frame_telemetry,
&usage_frame_telemetry,
);
if should_refresh_stream_usage {
state_for_report.usage_runtime.record_stream_started(
state_for_report.data.as_ref(),
&lifecycle_seed_for_report,
status_code,
Some(&frame_telemetry),
Some(&usage_frame_telemetry),
);
usage_stream_telemetry = Some(frame_telemetry.clone());
usage_stream_telemetry = Some(usage_frame_telemetry);
}
telemetry = Some(frame_telemetry);
}
@@ -3327,6 +3606,16 @@ async fn execute_stream_from_frame_stream(
);
downstream_dropped = true;
} else {
maybe_record_first_visible_text_started(
&state_for_report,
&lifecycle_seed_for_report,
status_code,
&mut client_visible_text_tracker,
rewritten_chunk.as_ref(),
stream_started_at_for_report,
telemetry.as_ref(),
&mut usage_stream_telemetry,
);
client_visible_stream_completed |= client_stream_completion_tracker
.observe_chunk(rewritten_chunk.as_ref());
client_stream_bytes
@@ -3398,6 +3687,16 @@ async fn execute_stream_from_frame_stream(
);
downstream_dropped = true;
} else {
maybe_record_first_visible_text_started(
&state_for_report,
&lifecycle_seed_for_report,
status_code,
&mut client_visible_text_tracker,
flushed_chunk.as_ref(),
stream_started_at_for_report,
telemetry.as_ref(),
&mut usage_stream_telemetry,
);
client_visible_stream_completed |= client_stream_completion_tracker
.observe_chunk(flushed_chunk.as_ref());
client_stream_bytes.fetch_add(flushed_chunk_len, Ordering::Relaxed);
@@ -3831,7 +4130,7 @@ mod tests {
stream_requires_observed_terminal_event, stream_terminal_summary_missing_observed_finish,
stream_terminal_summary_missing_observed_finish_with_requirement,
stream_terminal_summary_represents_failure_with_requirement,
ClientVisibleStreamCompletionTracker,
ClientVisibleStreamCompletionTracker, ClientVisibleTextTracker,
};
use crate::control::GatewayControlDecision;
use crate::tunnel::{tunnel_protocol, TunnelProxyConn};
@@ -4726,6 +5025,35 @@ mod tests {
));
}
#[test]
fn client_visible_text_tracker_ignores_control_role_and_tool_events_until_text() {
let mut tracker = ClientVisibleTextTracker::new("openai:chat");
assert!(!tracker.observe_chunk(
b": keepalive\n\n\
data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n"
));
assert!(!tracker.observe_chunk(
b"data: {\"choices\":[{\"delta\":{\"tool_calls\":[{\"index\":0,\"function\":{\"arguments\":\"{}\"}}]}}]}\n\n"
));
assert!(
tracker.observe_chunk(b"data: {\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n")
);
assert!(!tracker
.observe_chunk(b"data: {\"choices\":[{\"delta\":{\"content\":\"again\"}}]}\n\n"));
}
#[test]
fn client_visible_text_tracker_detects_responses_output_text_across_chunks() {
let mut tracker = ClientVisibleTextTracker::new("openai:responses");
assert!(!tracker.observe_chunk(
b"event: response.created\n\
data: {\"type\":\"response.created\"}\n\n\
event: response.output_text.delta\n\
data: {\"type\":\"response.output_text.delta\",\"delta\":\""
));
assert!(tracker.observe_chunk(b"Hi\"}\n\n"));
}
#[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);
@@ -5315,7 +5643,7 @@ mod tests {
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
));
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"first\\\"}\\n\\n\"}}\n",
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"first\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{\\\"content\\\":\\\"hello\\\"}}]}\\n\\n\"}}\n",
));
release_terminal_for_stream.notified().await;
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(
@@ -5364,7 +5692,10 @@ mod tests {
})
.await
.expect("first business chunk should arrive");
assert_eq!(first.as_ref(), b"data: {\"id\":\"first\"}\n\n");
assert_eq!(
first.as_ref(),
b"data: {\"id\":\"first\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n"
);
tokio::time::sleep(Duration::from_millis(30)).await;
drop(body_stream);
release_terminal.notify_one();
@@ -5867,6 +6198,193 @@ mod tests {
server.abort();
}
#[tokio::test]
async fn execute_execution_runtime_stream_waits_for_visible_text_before_recording_first_byte() {
let listener = crate::test_support::bind_loopback_listener()
.await
.expect("listener should bind");
let addr = listener.local_addr().expect("local addr should resolve");
let role_only_seen = Arc::new(Notify::new());
let release_text = Arc::new(Notify::new());
let text_seen = Arc::new(Notify::new());
let release_terminal = Arc::new(Notify::new());
let role_only_seen_for_route = Arc::clone(&role_only_seen);
let release_text_for_route = Arc::clone(&release_text);
let text_seen_for_route = Arc::clone(&text_seen);
let release_terminal_for_route = Arc::clone(&release_terminal);
let server = tokio::spawn(async move {
let app = Router::new().route(
"/v1/execute/stream",
any(move |_request: Request| {
let role_only_seen = Arc::clone(&role_only_seen_for_route);
let release_text = Arc::clone(&release_text_for_route);
let text_seen = Arc::clone(&text_seen_for_route);
let release_terminal = Arc::clone(&release_terminal_for_route);
async move {
let frames = stream! {
yield Ok::<Bytes, Infallible>(Bytes::from_static(
b"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
));
yield Ok::<Bytes, Infallible>(Bytes::from_static(
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"choices\\\":[{\\\"delta\\\":{\\\"role\\\":\\\"assistant\\\"}}]}\\n\\n\"}}\n",
));
yield Ok::<Bytes, Infallible>(Bytes::from_static(
b"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"ttfb_ms\":11,\"elapsed_ms\":12}}}\n",
));
role_only_seen.notify_one();
release_text.notified().await;
yield Ok::<Bytes, Infallible>(Bytes::from_static(
b"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"choices\\\":[{\\\"delta\\\":{\\\"content\\\":\\\"hello\\\"}}]}\\n\\n\"}}\n",
));
text_seen.notify_one();
release_terminal.notified().await;
yield Ok::<Bytes, Infallible>(Bytes::from_static(
b"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":50}}}\n",
));
yield Ok::<Bytes, Infallible>(Bytes::from_static(
b"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n",
));
};
let mut response = axum::http::Response::new(Body::from_stream(frames));
response.headers_mut().insert(
header::CONTENT_TYPE,
HeaderValue::from_static("application/x-ndjson"),
);
response
}
}),
);
axum::serve(listener, app)
.await
.expect("server should start");
});
let usage_repository = Arc::new(InMemoryUsageReadRepository::default());
let request_candidate_repository = Arc::new(InMemoryRequestCandidateRepository::default());
let state = AppState::new()
.expect("app state should build")
.with_data_state_for_tests(
crate::data::GatewayDataState::with_request_candidate_and_usage_repository_for_tests(
Arc::clone(&request_candidate_repository),
Arc::clone(&usage_repository),
),
)
.with_usage_runtime_for_tests(UsageRuntimeConfig {
enabled: true,
..UsageRuntimeConfig::default()
})
.with_execution_runtime_override_base_url(format!("http://{addr}"));
let plan = ExecutionPlan {
request_id: "req-live-stream-visible-text".into(),
candidate_id: Some("cand-live-stream-visible-text".into()),
provider_name: Some("openai".into()),
provider_id: "prov-1".into(),
endpoint_id: "ep-1".into(),
key_id: "key-1".into(),
method: "POST".into(),
url: "https://api.openai.com/v1/chat/completions".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-5.4",
"messages": [{"role": "user", "content": "hello"}],
"stream": true
})),
stream: true,
client_api_format: "openai:chat".into(),
provider_api_format: "openai:chat".into(),
model_name: Some("gpt-5.4".into()),
proxy: None,
transport_profile: None,
timeouts: Some(ExecutionTimeouts {
connect_ms: Some(5_000),
total_ms: Some(5_000),
..ExecutionTimeouts::default()
}),
};
let decision = GatewayControlDecision::synthetic(
"/v1/chat/completions",
Some("ai_public".to_string()),
Some("openai".to_string()),
Some("chat".to_string()),
Some("openai:chat".to_string()),
)
.with_execution_runtime_candidate(true);
let response = execute_execution_runtime_stream(
&state,
plan,
"trace-live-stream-visible-text",
&decision,
"openai_chat_stream",
None,
Some(json!({
"provider_api_format": "openai:chat",
"client_api_format": "openai:chat",
})),
)
.await
.expect("execution should succeed")
.expect("execution should return a client response");
role_only_seen.notified().await;
let deadline = tokio::time::Instant::now() + Duration::from_secs(2);
let role_only_usage = loop {
let usage = usage_repository
.find_by_request_id("req-live-stream-visible-text")
.await
.expect("usage should read");
if usage.as_ref().is_some_and(|usage| {
usage.status == "streaming" && usage.response_time_ms == Some(12)
}) {
break usage.expect("streaming usage should exist");
}
assert!(
tokio::time::Instant::now() < deadline,
"usage should record streaming status"
);
tokio::time::sleep(Duration::from_millis(10)).await;
};
assert_eq!(role_only_usage.first_byte_time_ms, None);
assert_eq!(role_only_usage.response_time_ms, Some(12));
release_text.notify_one();
text_seen.notified().await;
let deadline = tokio::time::Instant::now() + Duration::from_secs(2);
let visible_text_usage = loop {
let usage = usage_repository
.find_by_request_id("req-live-stream-visible-text")
.await
.expect("usage should read");
if usage
.as_ref()
.is_some_and(|usage| usage.first_byte_time_ms.is_some())
{
break usage.expect("visible text usage should exist");
}
assert!(
tokio::time::Instant::now() < deadline,
"usage should record first byte only after visible text"
);
tokio::time::sleep(Duration::from_millis(10)).await;
};
assert!(visible_text_usage.first_byte_time_ms.unwrap_or(0) >= 12);
release_terminal.notify_one();
let body = to_bytes(response.into_body(), usize::MAX)
.await
.expect("response body should read");
let text = String::from_utf8(body.to_vec()).expect("response body should be utf8");
assert!(text.contains("\"role\":\"assistant\""));
assert!(text.contains("\"content\":\"hello\""));
server.abort();
}
#[tokio::test]
async fn execute_execution_runtime_stream_bridges_sync_json_body_from_remote_runtime_to_sse() {
let listener = crate::test_support::bind_loopback_listener()
apps/aether-gateway/src/execution_runtime/windsurf.rs
+4 -1
View File
@@ -4044,7 +4044,10 @@ fn language_server_env(proxy_url: Option<&str>) -> BTreeMap<String, String> {
.or_else(|| std::env::var_os("HOME"))
.filter(|value| !value.is_empty())
{
env.insert("USERPROFILE".to_string(), PathBuf::from(home).display().to_string());
env.insert(
"USERPROFILE".to_string(),
PathBuf::from(home).display().to_string(),
);
}
}
if let Some(proxy_url) = proxy_url.filter(|value| !value.trim().is_empty()) {
apps/aether-gateway/src/handlers/admin/system/shared/modules.rs
+10 -8
View File
@@ -317,14 +317,16 @@ pub(crate) fn build_admin_module_validation_result(
runtime: &AdminModuleRuntimeState,
) -> (bool, Option<String>) {
admin_system_kernel::build_admin_module_validation_result(
module.name,
&runtime.oauth_providers,
runtime.ldap_config.as_ref(),
runtime.gemini_files_has_capable_key,
runtime.important_notification_configured,
runtime.server_chan_push_configured,
runtime.bark_push_configured,
runtime.s3_backup_configured,
admin_system_kernel::AdminModuleValidationInput {
module_name: module.name,
oauth_providers: &runtime.oauth_providers,
ldap_config: runtime.ldap_config.as_ref(),
gemini_files_has_capable_key: runtime.gemini_files_has_capable_key,
important_notification_configured: runtime.important_notification_configured,
server_chan_push_configured: runtime.server_chan_push_configured,
bark_push_configured: runtime.bark_push_configured,
s3_backup_configured: runtime.s3_backup_configured,
},
)
}
apps/aether-gateway/src/tests/usage/direct.rs
+6 -4
View File
@@ -485,7 +485,7 @@ async fn gateway_records_usage_for_execution_runtime_stream_when_runtime_enabled
any(|_request: Request| async move {
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-usage-stream-123\\\",\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-usage-stream-123\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{\\\"content\\\":\\\"hello\\\"}}],\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: [DONE]\\n\\n\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":51,\"ttfb_ms\":19}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
@@ -539,7 +539,8 @@ async fn gateway_records_usage_for_execution_runtime_stream_when_runtime_enabled
assert_eq!(stored.status, "completed");
assert_eq!(stored.billing_status, "pending");
assert_eq!(stored.total_tokens, 6);
assert_eq!(stored.first_byte_time_ms, Some(19));
assert!(stored.first_byte_time_ms.is_some());
assert!(stored.response_time_ms >= stored.first_byte_time_ms);
assert_eq!(stored.is_stream, true);
gateway_handle.abort();
@@ -572,7 +573,7 @@ async fn gateway_records_pending_usage_before_execution_runtime_stream_headers_a
allow_execution_response.notified().await;
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-usage-stream-pending-123\\\",\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-usage-stream-pending-123\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{\\\"content\\\":\\\"hello\\\"}}],\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: [DONE]\\n\\n\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":51,\"ttfb_ms\":19}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
@@ -688,7 +689,8 @@ async fn gateway_records_pending_usage_before_execution_runtime_stream_headers_a
}
let stored = stored.expect("usage should be finalized");
assert_eq!(stored.status, "completed");
assert_eq!(stored.first_byte_time_ms, Some(19));
assert!(stored.first_byte_time_ms.is_some());
assert!(stored.response_time_ms >= stored.first_byte_time_ms);
gateway_handle.abort();
execution_runtime_handle.abort();
apps/aether-gateway/src/tests/usage/local.rs
+4 -3
View File
@@ -1213,7 +1213,7 @@ async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_wh
any(|_request: Request| async move {
let frames = concat!(
"{\"type\":\"headers\",\"payload\":{\"kind\":\"headers\",\"status_code\":200,\"headers\":{\"content-type\":\"text/event-stream\"}}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-local-report-stream-123\\\",\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: {\\\"id\\\":\\\"chatcmpl-local-report-stream-123\\\",\\\"choices\\\":[{\\\"index\\\":0,\\\"delta\\\":{\\\"content\\\":\\\"hello\\\"}}],\\\"usage\\\":{\\\"input_tokens\\\":2,\\\"output_tokens\\\":4,\\\"total_tokens\\\":6}}\\n\\n\"}}\n",
"{\"type\":\"data\",\"payload\":{\"kind\":\"data\",\"text\":\"data: [DONE]\\n\\n\"}}\n",
"{\"type\":\"telemetry\",\"payload\":{\"kind\":\"telemetry\",\"telemetry\":{\"elapsed_ms\":31,\"ttfb_ms\":11}}}\n",
"{\"type\":\"eof\",\"payload\":{\"kind\":\"eof\"}}\n"
@@ -1286,7 +1286,7 @@ async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_wh
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"
"data: {\"id\":\"chatcmpl-local-report-stream-123\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}],\"usage\":{\"input_tokens\":2,\"output_tokens\":4,\"total_tokens\":6}}\n\ndata: [DONE]\n\n"
);
let stored_usage = wait_for_usage_status(
@@ -1297,7 +1297,8 @@ async fn gateway_handles_local_openai_chat_stream_report_with_local_reporting_wh
.await;
assert_eq!(stored_usage.status, "completed");
assert_eq!(stored_usage.total_tokens, 6);
assert_eq!(stored_usage.first_byte_time_ms, Some(11));
assert!(stored_usage.first_byte_time_ms.is_some());
assert!(stored_usage.response_time_ms >= stored_usage.first_byte_time_ms);
assert!(stored_usage.is_stream);
let stored_candidates = request_candidate_repository
apps/aether-gateway/src/tests/usage/pricing.rs
+18 -4
View File
@@ -637,7 +637,20 @@ fn assert_usage_and_pricing(
Some(expected_response_time_ms)
);
}
assert_eq!(stored_usage.first_byte_time_ms, expected_ttfb_ms);
if expected_ttfb_ms.is_some() {
let first_byte_time_ms = stored_usage
.first_byte_time_ms
.expect("stream usage should record first visible text time");
assert!(
stored_usage
.response_time_ms
.is_some_and(|response_time_ms| response_time_ms >= first_byte_time_ms),
"stream first_byte_time_ms should not exceed response_time_ms: first_byte={first_byte_time_ms}, response={:?}",
stored_usage.response_time_ms
);
} else {
assert_eq!(stored_usage.first_byte_time_ms, None);
}
assert_eq!(
stored_usage.settlement_input_price_per_1m(),
Some(INPUT_PRICE_PER_1M)
@@ -893,7 +906,7 @@ async fn gateway_records_openai_stream_usage_and_pricing_with_cache_tokens_impl(
cache_read_tokens: 40,
};
let stream_body = [
"data: {\"id\":\"chatcmpl-openai-usage-pricing-stream-123\",\"choices\":[]}\n\n",
"data: {\"id\":\"chatcmpl-openai-usage-pricing-stream-123\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"}}]}\n\n",
"data: [DONE]\n\n",
];
let frames = build_stream_frames(
@@ -1084,7 +1097,7 @@ async fn gateway_records_claude_stream_usage_and_pricing_with_cache_breakdown_im
cache_read_tokens: 10,
};
let stream_body = [
"event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"id\":\"msg-claude-usage-pricing-stream-123\",\"type\":\"message\",\"model\":\"claude-sonnet-4-5-upstream\",\"role\":\"assistant\",\"content\":[]}}\n\n",
"event: content_block_delta\ndata: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n",
"event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n",
];
let frames = build_stream_frames(
@@ -1273,7 +1286,8 @@ async fn gateway_records_gemini_stream_usage_and_pricing_with_cache_read_tokens_
cache_creation_ephemeral_1h_tokens: 0,
cache_read_tokens: 30,
};
let stream_body = ["data: {\"candidates\":[]}\n\n"];
let stream_body =
["data: {\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"hello\"}]}}]}\n\n"];
let frames = build_stream_frames(
&stream_body,
standardized_usage_json(expected),
apps/aether-gateway/src/tests/video/registry_poller.rs
+4 -4
View File
@@ -177,7 +177,7 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
assert!(!background_tasks.is_empty(), "poller task should spawn");
let stored = {
let deadline = tokio::time::Instant::now() + std::time::Duration::from_millis(500);
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
let stored = repository
.find(VideoTaskLookupKey::Id("task-local-123"))
@@ -189,7 +189,7 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
}
assert!(
tokio::time::Instant::now() < deadline,
"poller did not refresh task within 500ms"
"poller did not refresh task within 2s"
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
@@ -283,7 +283,7 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
assert!(!background_tasks.is_empty(), "poller task should spawn");
let stored = {
let deadline = tokio::time::Instant::now() + std::time::Duration::from_millis(500);
let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2);
loop {
let stored = repository
.find(VideoTaskLookupKey::Id("task-local-123"))
@@ -295,7 +295,7 @@ async fn gateway_background_video_task_poller_refreshes_due_openai_task_from_rep
}
assert!(
tokio::time::Instant::now() < deadline,
"poller did not refresh task within 500ms"
"poller did not refresh task within 2s"
);
tokio::time::sleep(std::time::Duration::from_millis(10)).await;
}
crates/aether-admin/src/system.rs
+23 -8
View File
@@ -1284,16 +1284,31 @@ pub fn ldap_module_config_is_valid(config: Option<&StoredLdapModuleConfig>) -> b
.is_some()
}
pub struct AdminModuleValidationInput<'a> {
pub module_name: &'a str,
pub oauth_providers: &'a [StoredOAuthProviderModuleConfig],
pub ldap_config: Option<&'a StoredLdapModuleConfig>,
pub gemini_files_has_capable_key: bool,
pub important_notification_configured: bool,
pub server_chan_push_configured: bool,
pub bark_push_configured: bool,
pub s3_backup_configured: bool,
}
pub fn build_admin_module_validation_result(
module_name: &str,
oauth_providers: &[StoredOAuthProviderModuleConfig],
ldap_config: Option<&StoredLdapModuleConfig>,
gemini_files_has_capable_key: bool,
important_notification_configured: bool,
server_chan_push_configured: bool,
bark_push_configured: bool,
s3_backup_configured: bool,
input: AdminModuleValidationInput<'_>,
) -> (bool, Option<String>) {
let AdminModuleValidationInput {
module_name,
oauth_providers,
ldap_config,
gemini_files_has_capable_key,
important_notification_configured,
server_chan_push_configured,
bark_push_configured,
s3_backup_configured,
} = input;
match module_name {
"oauth" => {
if oauth_providers.is_empty() {
crates/aether-data/src/repository/usage/mysql.rs
+10 -2
View File
@@ -166,8 +166,16 @@ ON DUPLICATE KEY UPDATE
status_code = VALUES(status_code),
error_message = VALUES(error_message),
error_category = VALUES(error_category),
response_time_ms = VALUES(response_time_ms),
first_byte_time_ms = VALUES(first_byte_time_ms),
response_time_ms = CASE
WHEN VALUES(response_time_ms) IS NULL OR VALUES(response_time_ms) = 0
THEN COALESCE(response_time_ms, VALUES(response_time_ms))
ELSE VALUES(response_time_ms)
END,
first_byte_time_ms = CASE
WHEN VALUES(first_byte_time_ms) IS NULL OR VALUES(first_byte_time_ms) = 0
THEN COALESCE(first_byte_time_ms, VALUES(first_byte_time_ms))
ELSE VALUES(first_byte_time_ms)
END,
status = VALUES(status),
billing_status = VALUES(billing_status),
request_metadata = VALUES(request_metadata),
crates/aether-oauth/src/provider/providers/windsurf.rs
+9 -10
View File
@@ -353,16 +353,15 @@ enum RegisterUserBody {
LegacyJsonFirebaseToken,
}
type RegisterUserRequestParts = (
BTreeMap<String, String>,
&'static str,
Option<Value>,
Option<Vec<u8>>,
);
impl RegisterUserBody {
fn request_parts(
self,
token: &str,
) -> (
BTreeMap<String, String>,
&'static str,
Option<Value>,
Option<Vec<u8>>,
) {
fn request_parts(self, token: &str) -> RegisterUserRequestParts {
match self {
Self::ProtoOneTimeToken => (
register_user_proto_headers(),
@@ -1101,7 +1100,7 @@ mod tests {
);
let expected_body = proto_string_field_body(1, one_time_token);
assert_eq!(
requests[0].body_bytes.as_ref().map(|body| body.as_slice()),
requests[0].body_bytes.as_deref(),
Some(expected_body.as_slice())
);
}