fix(usage): record visible stream first byte timing

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
@@ -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();
+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
+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),
@@ -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;
}