mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-06 17:37:47 +08:00
fix(usage): record visible stream first byte timing
This commit is contained in:
@@ -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();
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user