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

Preserve streaming usage timing on refresh

Browse Source
This commit is contained in:
fawney19
2026-05-28 16:33:10 +08:00
parent d076f64db3
commit 47886abd2b
1 changed files with 52 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/src/repository/usage/memory.rs
+52 -2
View File
@@ -161,6 +161,13 @@ fn usage_status_is_lifecycle(status: &str) -> bool {
matches!(status, "pending" | "streaming")
}
fn merge_usage_timing(existing: Option<u64>, incoming: Option<u64>) -> Option<u64> {
match incoming {
Some(0) | None => existing.or(incoming),
Some(value) => Some(value),
}
}
fn accumulate_provider_api_key_usage_contribution(
aggregates: &mut BTreeMap<String, ProviderApiKeyUsageContribution>,
contribution: ProviderApiKeyUsageContribution,
@@ -2781,8 +2788,18 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
status_code: usage.status_code,
error_message: usage.error_message,
error_category: usage.error_category,
response_time_ms: usage.response_time_ms,
first_byte_time_ms: usage.first_byte_time_ms,
response_time_ms: merge_usage_timing(
existing
.as_ref()
.and_then(|existing| existing.response_time_ms),
usage.response_time_ms,
),
first_byte_time_ms: merge_usage_timing(
existing
.as_ref()
.and_then(|existing| existing.first_byte_time_ms),
usage.first_byte_time_ms,
),
status: usage.status,
billing_status: usage.billing_status,
request_headers: usage.request_headers.or_else(|| {
@@ -3823,6 +3840,39 @@ mod tests {
assert_eq!(stored.total_tokens, 12);
}
#[tokio::test]
async fn streaming_refresh_without_timing_does_not_clear_stream_timing() {
let repository = InMemoryUsageReadRepository::default();
let mut first = sample_upsert_usage_record("req-streaming-refresh");
first.is_stream = Some(true);
first.status = "streaming".to_string();
first.status_code = Some(200);
first.response_time_ms = Some(45);
first.first_byte_time_ms = Some(12);
repository
.upsert(first)
.await
.expect("streaming usage should upsert");
let mut refresh = sample_upsert_usage_record("req-streaming-refresh");
refresh.is_stream = Some(true);
refresh.status = "streaming".to_string();
refresh.status_code = Some(200);
repository
.upsert(refresh)
.await
.expect("streaming refresh should upsert");
let stored = repository
.find_by_request_id("req-streaming-refresh")
.await
.expect("usage lookup should succeed")
.expect("usage should exist");
assert_eq!(stored.status, "streaming");
assert_eq!(stored.response_time_ms, Some(45));
assert_eq!(stored.first_byte_time_ms, Some(12));
}
#[tokio::test]
async fn seed_hydrates_legacy_body_ref_metadata_into_typed_fields() {
let repository = InMemoryUsageReadRepository::seed(vec![StoredRequestUsageAudit {