Add usage client family read model field

crates/aether-data/src/repository/usage/memory.rs

@@ -29,11 +29,12 @@ use serde_json::Value;
 use super::{
     api_key_usage_contribution, provider_api_key_usage_contribution,
     strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
-    ApiKeyUsageContribution, ApiKeyUsageDelta, ProviderApiKeyUsageContribution,
-    ProviderApiKeyUsageDelta, ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
-    StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
-    StoredRequestUsageAudit, StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery,
-    UsageDailyHeatmapQuery, UsageReadRepository, UsageWriteRepository,
+    usage_request_metadata_client_family, ApiKeyUsageContribution, ApiKeyUsageDelta,
+    ProviderApiKeyUsageContribution, ProviderApiKeyUsageDelta, ProviderApiKeyWindowUsageRequest,
+    StoredProviderApiKeyUsageSummary, StoredProviderApiKeyWindowUsageSummary,
+    StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
+    StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery, UsageDailyHeatmapQuery,
+    UsageReadRepository, UsageWriteRepository,
 };
 use crate::repository::auth::InMemoryAuthApiKeySnapshotRepository;
 use crate::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
@@ -56,6 +57,7 @@ impl InMemoryUsageReadRepository {
         let mut by_request_id = BTreeMap::new();
         for mut item in items {
             hydrate_legacy_body_refs(&mut item);
+            hydrate_client_family(&mut item);
             by_request_id.insert(item.request_id.clone(), item);
         }
         Self {
@@ -75,6 +77,7 @@ impl InMemoryUsageReadRepository {
         let mut detached_bodies = BTreeMap::new();
         for mut item in items {
             hydrate_legacy_body_refs(&mut item);
+            hydrate_client_family(&mut item);
             let request_id = item.request_id.clone();
             if let Some(body_ref) = detach_usage_body(
                 &request_id,
@@ -2546,6 +2549,13 @@ fn hydrate_legacy_body_refs(item: &mut StoredRequestUsageAudit) {
     }
 }
 
+fn hydrate_client_family(item: &mut StoredRequestUsageAudit) {
+    if item.client_family.is_none() {
+        item.client_family = usage_request_metadata_client_family(item.request_metadata.as_ref())
+            .map(ToOwned::to_owned);
+    }
+}
+
 fn persisted_usage_body_ref(
     incoming_ref: Option<&str>,
     incoming_body: Option<&Value>,
@@ -2851,6 +2861,13 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
                     })
                 },
             ),
+            client_family: usage_request_metadata_client_family(request_metadata.as_ref())
+                .map(ToOwned::to_owned)
+                .or_else(|| {
+                    existing
+                        .as_ref()
+                        .and_then(|existing| existing.client_family.clone())
+                }),
             request_metadata,
             created_at_unix_ms,
             updated_at_unix_secs: usage.updated_at_unix_secs,
@@ -4399,6 +4416,7 @@ mod tests {
             provider_endpoint_kind: Some("chat".to_string()),
             has_format_conversion: false,
             is_stream: false,
+            client_family: None,
             input_tokens: 10,
             output_tokens: 20,
             total_tokens: 30,