use std::collections::BTreeMap; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; use aether_data_contracts::repository::usage::UsageBodyCaptureState; use aether_data_contracts::DataLayerError; use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::body_capture::mark_usage_event_capture_truncated; pub use crate::event_capture_budget::UsageEventCaptureRetention; use crate::event_capture_budget::{ json_heap_estimate, shared_capture_memory_budget, EventCaptureMemoryBudget, }; pub const USAGE_EVENT_VERSION: u8 = 1; #[path = "event_wire.rs"] mod wire; pub(crate) use wire::EncodedUsageEvent; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum UsageEventType { Pending, Streaming, Completed, Failed, Cancelled, } #[derive(Debug, PartialEq, Serialize, Deserialize, Default)] pub struct UsageEventData { #[serde(default, skip_serializing_if = "Option::is_none")] pub user_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub api_key_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub username: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub api_key_name: Option, pub provider_name: String, pub model: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub target_model: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub model_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub global_model_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_endpoint_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_api_key_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub request_type: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub api_format: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub api_family: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub endpoint_kind: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub endpoint_api_format: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_api_family: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_endpoint_kind: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub has_format_conversion: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub is_stream: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub input_tokens: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub output_tokens: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub total_tokens: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub cache_creation_input_tokens: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub cache_creation_ephemeral_5m_input_tokens: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub cache_creation_ephemeral_1h_input_tokens: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub cache_read_input_tokens: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub cache_creation_cost_usd: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub cache_read_cost_usd: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub output_price_per_1m: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub total_cost_usd: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub actual_total_cost_usd: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub status_code: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub error_message: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub error_category: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub response_time_ms: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub first_byte_time_ms: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub request_headers: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub request_body: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub request_body_ref: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub request_body_state: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_request_headers: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_request_body: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_request_body_ref: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub provider_request_body_state: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub response_headers: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub response_body: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub response_body_ref: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub response_body_state: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub client_response_headers: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub client_response_body: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub client_response_body_ref: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub client_response_body_state: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub candidate_id: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub candidate_index: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub key_name: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub planner_kind: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub route_family: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub route_kind: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub execution_path: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub local_execution_runtime_miss_reason: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub request_metadata: Option, #[doc(hidden)] #[serde(skip)] pub capture_retention: UsageEventCaptureRetention, } impl UsageEventData { fn capture_heap_estimate(&self) -> usize { [ &self.request_body, &self.provider_request_body, &self.response_body, &self.client_response_body, ] .into_iter() .flatten() .fold(0usize, |bytes, body| { bytes .saturating_add(std::mem::size_of::()) .saturating_add(json_heap_estimate(body)) }) } fn captured_fields(&self) -> [bool; 4] { [ self.request_body.is_some(), self.provider_request_body.is_some(), self.response_body.is_some(), self.client_response_body.is_some(), ] } fn mark_capture_omitted(&mut self, captured: [bool; 4]) { for (present, key, state) in [ (captured[0], "request", &mut self.request_body_state), ( captured[1], "provider_request", &mut self.provider_request_body_state, ), (captured[2], "response", &mut self.response_body_state), ( captured[3], "client_response", &mut self.client_response_body_state, ), ] { if present && !matches!( *state, Some( UsageBodyCaptureState::None | UsageBodyCaptureState::Disabled | UsageBodyCaptureState::Unavailable ) ) { *state = Some(UsageBodyCaptureState::Truncated); mark_usage_event_capture_truncated(&mut self.request_metadata, key); } } } pub(crate) fn apply_capture_memory_budget( &mut self, budget: std::sync::Arc, ) { let bytes = self.capture_heap_estimate(); if self .capture_retention .reserve(std::sync::Arc::clone(&budget), bytes) { return; } let captured = self.captured_fields(); self.request_body = None; self.provider_request_body = None; self.response_body = None; self.client_response_body = None; self.mark_capture_omitted(captured); // The previous lease is released only after the owned JSON bodies are gone. self.capture_retention.clear(budget); } } impl Clone for UsageEventData { fn clone(&self) -> Self { let (capture_retention, retain_bodies) = self .capture_retention .clone_for_bodies(|| self.capture_heap_estimate()); // Enumerate every field so additions require an explicit ownership decision. let mut cloned = Self { user_id: self.user_id.clone(), api_key_id: self.api_key_id.clone(), username: self.username.clone(), api_key_name: self.api_key_name.clone(), provider_name: self.provider_name.clone(), model: self.model.clone(), target_model: self.target_model.clone(), model_id: self.model_id.clone(), global_model_id: self.global_model_id.clone(), provider_id: self.provider_id.clone(), provider_endpoint_id: self.provider_endpoint_id.clone(), provider_api_key_id: self.provider_api_key_id.clone(), request_type: self.request_type.clone(), api_format: self.api_format.clone(), api_family: self.api_family.clone(), endpoint_kind: self.endpoint_kind.clone(), endpoint_api_format: self.endpoint_api_format.clone(), provider_api_family: self.provider_api_family.clone(), provider_endpoint_kind: self.provider_endpoint_kind.clone(), has_format_conversion: self.has_format_conversion, is_stream: self.is_stream, input_tokens: self.input_tokens, output_tokens: self.output_tokens, total_tokens: self.total_tokens, cache_creation_input_tokens: self.cache_creation_input_tokens, cache_creation_ephemeral_5m_input_tokens: self.cache_creation_ephemeral_5m_input_tokens, cache_creation_ephemeral_1h_input_tokens: self.cache_creation_ephemeral_1h_input_tokens, cache_read_input_tokens: self.cache_read_input_tokens, cache_creation_cost_usd: self.cache_creation_cost_usd, cache_read_cost_usd: self.cache_read_cost_usd, output_price_per_1m: self.output_price_per_1m, total_cost_usd: self.total_cost_usd, actual_total_cost_usd: self.actual_total_cost_usd, status_code: self.status_code, error_message: self.error_message.clone(), error_category: self.error_category.clone(), response_time_ms: self.response_time_ms, first_byte_time_ms: self.first_byte_time_ms, request_headers: self.request_headers.clone(), request_body: retain_bodies.then(|| self.request_body.clone()).flatten(), request_body_ref: self.request_body_ref.clone(), request_body_state: self.request_body_state, provider_request_headers: self.provider_request_headers.clone(), provider_request_body: retain_bodies .then(|| self.provider_request_body.clone()) .flatten(), provider_request_body_ref: self.provider_request_body_ref.clone(), provider_request_body_state: self.provider_request_body_state, response_headers: self.response_headers.clone(), response_body: retain_bodies.then(|| self.response_body.clone()).flatten(), response_body_ref: self.response_body_ref.clone(), response_body_state: self.response_body_state, client_response_headers: self.client_response_headers.clone(), client_response_body: retain_bodies .then(|| self.client_response_body.clone()) .flatten(), client_response_body_ref: self.client_response_body_ref.clone(), client_response_body_state: self.client_response_body_state, candidate_id: self.candidate_id.clone(), candidate_index: self.candidate_index, key_name: self.key_name.clone(), planner_kind: self.planner_kind.clone(), route_family: self.route_family.clone(), route_kind: self.route_kind.clone(), execution_path: self.execution_path.clone(), local_execution_runtime_miss_reason: self.local_execution_runtime_miss_reason.clone(), request_metadata: self.request_metadata.clone(), capture_retention, }; if !retain_bodies { cloned.mark_capture_omitted(self.captured_fields()); } cloned } } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub struct UsageEvent { pub event_type: UsageEventType, pub request_id: String, pub timestamp_ms: u64, pub data: UsageEventData, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] struct UsageEventEnvelope { v: u8, #[serde(rename = "type")] event_type: UsageEventType, request_id: String, timestamp_ms: u64, data: UsageEventData, } #[derive(Serialize)] struct BorrowedUsageEventEnvelope<'a, T: ?Sized> { v: u8, #[serde(rename = "type")] event_type: UsageEventType, request_id: &'a str, timestamp_ms: u64, data: &'a T, } impl UsageEvent { pub fn new( event_type: UsageEventType, request_id: impl Into, data: UsageEventData, ) -> Self { Self { event_type, request_id: request_id.into(), timestamp_ms: now_ms(), data, } } pub fn to_stream_fields(&self) -> Result, DataLayerError> { let payload = BorrowedUsageEventEnvelope { v: USAGE_EVENT_VERSION, event_type: self.event_type, request_id: &self.request_id, timestamp_ms: self.timestamp_ms, data: &self.data, }; let payload = serde_json::to_string(&payload).map_err(|err| { DataLayerError::UnexpectedValue(format!( "failed to serialize usage event payload: {err}" )) })?; Ok(BTreeMap::from([("payload".to_string(), payload)])) } pub(crate) fn to_bounded_stream_fields( &self, max_bytes: usize, ) -> Result { wire::encode(self, max_bytes) } pub fn from_stream_fields(fields: &BTreeMap) -> Result { Self::from_stream_fields_with_capture_budget(fields, shared_capture_memory_budget()) } pub(crate) fn from_stream_fields_with_capture_budget( fields: &BTreeMap, budget: Arc, ) -> Result { let payload = fields.get("payload").ok_or_else(|| { DataLayerError::UnexpectedValue( "usage event stream entry missing payload field".to_string(), ) })?; let envelope: UsageEventEnvelope = serde_json::from_str(payload).map_err(|err| { DataLayerError::UnexpectedValue(format!( "failed to deserialize usage event payload: {err}" )) })?; if envelope.v != USAGE_EVENT_VERSION { return Err(DataLayerError::UnexpectedValue(format!( "unsupported usage event version: {}", envelope.v ))); } let mut event = Self { event_type: envelope.event_type, request_id: envelope.request_id, timestamp_ms: envelope.timestamp_ms, data: envelope.data, }; // The wire format has no ownership lease. Preserve billing facts before a decoded // body can be omitted; the raw Redis response and serde allocation are not budgeted here. crate::runtime::prepare_decoded_event_capture_memory(&mut event, budget); Ok(event) } } pub fn now_ms() -> u64 { SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap_or_default() .as_millis() as u64 } #[cfg(test)] mod tests { use std::collections::BTreeMap; use std::sync::Arc; use aether_data_contracts::repository::usage::UsageBodyCaptureState; use aether_data_contracts::DataLayerError; use serde_json::json; use crate::event_capture_budget::EventCaptureMemoryBudget; use crate::{ apply_usage_body_capture_policy_to_event, build_upsert_usage_record_from_event, UsageBodyCapturePolicy, }; use super::{UsageEvent, UsageEventData, UsageEventType}; fn captured_event() -> UsageEvent { UsageEvent { event_type: UsageEventType::Failed, request_id: "capture-budget-request".to_string(), timestamp_ms: 123_456, data: UsageEventData { provider_name: "provider".to_string(), model: "model".to_string(), input_tokens: Some(100), output_tokens: Some(500), total_tokens: Some(600), cache_read_input_tokens: Some(0), cache_creation_input_tokens: Some(25), actual_total_cost_usd: Some(1.25), status_code: Some(502), error_category: Some("upstream_error".to_string()), error_message: Some("upstream failed".to_string()), request_body: Some(json!({"messages": [{"content": "request"}]})), provider_request_body: Some(json!({"input": "upstream request"})), response_body: Some(json!({"usage": {"input_tokens": 100, "output_tokens": 500}})), client_response_body: Some(json!({"error": "client response"})), request_body_state: Some(UsageBodyCaptureState::Inline), provider_request_body_state: Some(UsageBodyCaptureState::Inline), response_body_state: Some(UsageBodyCaptureState::Inline), client_response_body_state: Some(UsageBodyCaptureState::Inline), request_metadata: Some(json!({ "requested_reasoning_effort": "high", "provider_reasoning_effort": "medium", "provider_service_tier": "priority", "provider_actual_service_tier": "default", "provider_cache_ttl_minutes": 60, "plan_usage_reservation_token": "550e8400-e29b-41d4-a716-446655440000", "body_capture": {"response": {"state": "inline", "source_bytes": 1000}} })), ..UsageEventData::default() }, } } #[test] fn event_capture_budget_zero_preserves_billing_refs_and_database_truncation() { let budget = Arc::new(EventCaptureMemoryBudget::new(0)); let mut event = captured_event(); event.data.request_body_ref = Some("usage://capture-budget-request/request_body".to_string()); event.data.response_body_ref = Some("usage://capture-budget-request/response_body".to_string()); event.data.apply_capture_memory_budget(Arc::clone(&budget)); assert!(event.data.request_body.is_none()); assert!(event.data.provider_request_body.is_none()); assert!(event.data.response_body.is_none()); assert!(event.data.client_response_body.is_none()); let capture_metadata = event .data .request_metadata .as_ref() .expect("capture metadata"); assert_eq!( capture_metadata["body_capture"]["response"]["source_bytes"], 1000 ); assert_eq!( capture_metadata["body_capture"]["response"]["stored_bytes"], 0 ); assert_eq!( capture_metadata["body_capture"]["response"]["reason"], "usage_event_memory_budget_exceeded" ); let record = build_upsert_usage_record_from_event(&event).expect("record mapping"); assert_eq!(record.status, "failed"); assert_eq!(record.input_tokens, Some(100)); assert_eq!(record.output_tokens, Some(500)); assert_eq!(record.cache_read_input_tokens, Some(0)); assert_eq!(record.cache_creation_input_tokens, Some(25)); assert_eq!(record.actual_total_cost_usd, Some(1.25)); assert_eq!(record.error_category.as_deref(), Some("upstream_error")); assert_eq!(record.request_body_ref, event.data.request_body_ref); assert_eq!(record.response_body_ref, event.data.response_body_ref); for state in [ record.request_body_state, record.provider_request_body_state, record.response_body_state, record.client_response_body_state, ] { assert_eq!(state, Some(UsageBodyCaptureState::Truncated)); } let metadata = record.request_metadata.expect("preserved metadata"); assert_eq!(metadata["provider_service_tier"], "priority"); assert_eq!(metadata["provider_actual_service_tier"], "default"); assert_eq!(metadata["provider_cache_ttl_minutes"], 60); assert_eq!( metadata["plan_usage_reservation_token"], "550e8400-e29b-41d4-a716-446655440000" ); // Persistence projects billing metadata; capture state remains in typed columns. assert!(metadata.get("body_capture").is_none()); assert_eq!(budget.retained_bytes(), 0); assert_eq!(budget.downgraded_total(), 1); } #[test] fn event_capture_budget_clone_reserves_each_copy_before_cloning_bodies() { let mut event = captured_event(); let weight = event.data.capture_heap_estimate(); let budget = Arc::new(EventCaptureMemoryBudget::new(weight * 2)); event.data.apply_capture_memory_budget(Arc::clone(&budget)); let copy = event.clone(); assert_eq!(copy, event); assert_eq!(budget.retained_bytes(), weight * 2); let downgraded = event.clone(); assert!(event.data.response_body.is_some()); assert!(copy.data.response_body.is_some()); assert!(downgraded.data.response_body.is_none()); assert_eq!( downgraded.data.response_body_state, Some(UsageBodyCaptureState::Truncated) ); assert_eq!(downgraded.data.total_tokens, event.data.total_tokens); assert_eq!(downgraded.timestamp_ms, event.timestamp_ms); assert_eq!(downgraded.event_type, event.event_type); assert_eq!(budget.retained_bytes(), weight * 2); drop(copy); assert_eq!(budget.retained_bytes(), weight); drop((event, downgraded)); assert_eq!(budget.retained_bytes(), 0); } #[test] fn event_capture_budget_serialization_borrows_bodies_without_charging_a_clone() { let mut event = captured_event(); let weight = event.data.capture_heap_estimate(); let budget = Arc::new(EventCaptureMemoryBudget::new(weight)); crate::runtime::prepare_event_capture_memory(&mut event, Arc::clone(&budget)); let decoded_budget = Arc::new(EventCaptureMemoryBudget::new(usize::MAX)); for _ in 0..3 { let fields = event.to_stream_fields().expect("wire serialization"); assert!(!fields["payload"].contains("capture_retention")); let decoded = UsageEvent::from_stream_fields_with_capture_budget( &fields, Arc::clone(&decoded_budget), ) .expect("wire decode"); assert_eq!(decoded, event); assert_eq!(budget.retained_bytes(), weight); assert!(decoded_budget.retained_bytes() > 0); drop(decoded); assert_eq!(decoded_budget.retained_bytes(), 0); } assert_eq!(budget.downgraded_total(), 0); drop(event); assert_eq!(budget.retained_bytes(), 0); } #[test] fn from_stream_fields_legacy_body_budget_preserves_billing_and_request_facts() { let mut event = captured_event(); event.data.model = "gpt-5.6-sol".to_string(); event.data.endpoint_api_format = Some("openai:responses".to_string()); event.data.request_body = Some(json!({"reasoning": {"effort": "high"}})); event.data.provider_request_body = Some(json!({ "model": "gpt-5.6-sol", "reasoning": {"effort": "medium"}, "service_tier": "priority" })); event.data.response_body = Some(json!({"service_tier": "Default"})); event.data.request_body_state = None; event.data.provider_request_body_state = None; event.data.response_body_state = None; event.data.client_response_body_state = None; event.data.cache_creation_ephemeral_5m_input_tokens = Some(0); event.data.cache_creation_ephemeral_1h_input_tokens = Some(25); event.data.cache_read_cost_usd = Some(0.0); event.data.request_body_ref = Some("usage://legacy/request".to_string()); event.data.request_metadata = Some(json!({ "plan_usage_reservation_token": "550e8400-e29b-41d4-a716-446655440000" })); let fields = event.to_stream_fields().expect("legacy wire serialization"); assert!(!fields["payload"].contains("request_body_state")); let budget = Arc::new(EventCaptureMemoryBudget::new(0)); let decoded = UsageEvent::from_stream_fields_with_capture_budget(&fields, Arc::clone(&budget)) .expect("legacy wire decode"); assert_eq!(decoded.event_type, UsageEventType::Failed); assert_eq!(decoded.request_id, event.request_id); assert_eq!(decoded.timestamp_ms, event.timestamp_ms); assert_eq!(decoded.data.input_tokens, Some(100)); assert_eq!(decoded.data.output_tokens, Some(500)); assert_eq!(decoded.data.total_tokens, Some(600)); assert_eq!(decoded.data.cache_creation_input_tokens, Some(25)); assert_eq!( decoded.data.cache_creation_ephemeral_5m_input_tokens, Some(0) ); assert_eq!( decoded.data.cache_creation_ephemeral_1h_input_tokens, Some(25) ); assert_eq!(decoded.data.cache_read_input_tokens, Some(0)); assert_eq!(decoded.data.cache_read_cost_usd, Some(0.0)); assert_eq!(decoded.data.actual_total_cost_usd, Some(1.25)); assert_eq!(decoded.data.status_code, Some(502)); assert_eq!( decoded.data.error_category.as_deref(), Some("upstream_error") ); assert_eq!(decoded.data.request_body_ref, event.data.request_body_ref); assert!(decoded.data.request_body.is_none()); assert!(decoded.data.provider_request_body.is_none()); assert!(decoded.data.response_body.is_none()); assert!(decoded.data.client_response_body.is_none()); for state in [ decoded.data.request_body_state, decoded.data.provider_request_body_state, decoded.data.response_body_state, decoded.data.client_response_body_state, ] { assert_eq!(state, Some(UsageBodyCaptureState::Truncated)); } let metadata = decoded .data .request_metadata .as_ref() .expect("preserved facts"); assert_eq!(metadata["requested_reasoning_effort"], "high"); assert_eq!(metadata["provider_reasoning_effort"], "medium"); assert_eq!(metadata["provider_service_tier"], "priority"); assert_eq!(metadata["provider_actual_service_tier"], "default"); assert_eq!(metadata["provider_cache_ttl_minutes"], 30); assert_eq!( metadata["plan_usage_reservation_token"], "550e8400-e29b-41d4-a716-446655440000" ); assert_eq!(budget.retained_bytes(), 0); assert_eq!(budget.downgraded_total(), 1); } #[test] fn from_stream_fields_reconstructed_lease_also_bounds_recorder_clones() { let fields = captured_event() .to_stream_fields() .expect("wire serialization"); let probe_budget = Arc::new(EventCaptureMemoryBudget::new(usize::MAX)); let probe = UsageEvent::from_stream_fields_with_capture_budget(&fields, Arc::clone(&probe_budget)) .expect("estimate decoded allocation"); let weight = probe_budget.retained_bytes(); assert!(weight > 0); drop(probe); assert_eq!(probe_budget.retained_bytes(), 0); let budget = Arc::new(EventCaptureMemoryBudget::new(weight * 2)); let event = UsageEvent::from_stream_fields_with_capture_budget(&fields, Arc::clone(&budget)) .expect("wire decode"); let recorder_copy = event.clone(); assert!(recorder_copy.data.response_body.is_some()); assert_eq!(budget.retained_bytes(), weight * 2); let omitted_copy = event.clone(); assert!(omitted_copy.data.response_body.is_none()); assert_eq!(omitted_copy.data.total_tokens, Some(600)); assert_eq!(omitted_copy.data.cache_read_input_tokens, Some(0)); assert_eq!(budget.retained_bytes(), weight * 2); drop(event); assert_eq!(budget.retained_bytes(), weight); drop((recorder_copy, omitted_copy)); assert_eq!(budget.retained_bytes(), 0); } fn assert_typed_body_clear_is_preserved(event: &UsageEvent, state: UsageBodyCaptureState) { for (body, reference, actual_state) in [ ( &event.data.request_body, &event.data.request_body_ref, event.data.request_body_state, ), ( &event.data.provider_request_body, &event.data.provider_request_body_ref, event.data.provider_request_body_state, ), ( &event.data.response_body, &event.data.response_body_ref, event.data.response_body_state, ), ( &event.data.client_response_body, &event.data.client_response_body_ref, event.data.client_response_body_state, ), ] { assert!(body.is_none()); assert_eq!(actual_state, Some(state)); assert_eq!(reference.as_deref(), Some("usage://stale/reference")); } assert!(event .data .request_metadata .as_ref() .and_then(|value| value.get("body_capture")) .is_none()); } #[test] fn from_stream_fields_and_clone_budget_preserve_typed_clear_with_residual_bodies() { for state in [ UsageBodyCaptureState::None, UsageBodyCaptureState::Disabled, UsageBodyCaptureState::Unavailable, ] { let mut source = captured_event(); source.data.request_metadata = None; source.data.request_body_state = Some(state); source.data.provider_request_body_state = Some(state); source.data.response_body_state = Some(state); source.data.client_response_body_state = Some(state); source.data.request_body_ref = Some("usage://stale/reference".to_string()); source.data.provider_request_body_ref = Some("usage://stale/reference".to_string()); source.data.response_body_ref = Some("usage://stale/reference".to_string()); source.data.client_response_body_ref = Some("usage://stale/reference".to_string()); let fields = source.to_stream_fields().expect("wire serialization"); let decoded_budget = Arc::new(EventCaptureMemoryBudget::new(0)); let decoded = UsageEvent::from_stream_fields_with_capture_budget( &fields, Arc::clone(&decoded_budget), ) .expect("wire decode"); assert_typed_body_clear_is_preserved(&decoded, state); assert_eq!(decoded_budget.retained_bytes(), 0); let clone_budget = Arc::new(EventCaptureMemoryBudget::new( source.data.capture_heap_estimate(), )); crate::runtime::prepare_event_capture_memory(&mut source, Arc::clone(&clone_budget)); let cloned = source.clone(); assert_typed_body_clear_is_preserved(&cloned, state); assert!(source.data.request_body.is_some()); assert_eq!(clone_budget.downgraded_total(), 1); drop(source); assert_eq!(clone_budget.retained_bytes(), 0); } } #[test] fn from_stream_fields_legacy_metadata_only_facts_survive_before_billing() { for limit in [0, 8192] { for include_response in [false, true] { let mut source = legacy_metadata_only_event(); if include_response { source.data.response_body = Some(json!({"result": "response capture"})); } let fields = source .to_stream_fields() .expect("legacy wire serialization"); let budget = Arc::new(EventCaptureMemoryBudget::new(limit)); let decoded = UsageEvent::from_stream_fields_with_capture_budget( &fields, Arc::clone(&budget), ) .expect("legacy wire decode"); // The worker enriches this clone before DTO conversion. Missing legacy bodies // must not erase a previously derived TTL or turn an explicit zero into unknown. let billing_event = decoded.clone(); let metadata = billing_event .data .request_metadata .as_ref() .expect("legacy facts"); assert_eq!(metadata["requested_reasoning_effort"], "high"); assert_eq!(metadata["provider_reasoning_effort"], "medium"); assert_eq!(metadata["provider_service_tier"], "priority"); assert_eq!(metadata["provider_actual_service_tier"], "default"); assert_eq!(metadata["provider_cache_ttl_minutes"], 60); assert_eq!(billing_event.data.input_tokens, Some(0)); assert_eq!(billing_event.data.output_tokens, Some(0)); assert_eq!(billing_event.data.total_tokens, Some(0)); assert_eq!(billing_event.data.cache_read_input_tokens, Some(0)); assert_eq!(billing_event.data.cache_creation_input_tokens, Some(0)); assert_eq!(billing_event.data.actual_total_cost_usd, Some(0.0)); assert_eq!(billing_event.data.request_body_state, None); assert_eq!(billing_event.data.provider_request_body_state, None); drop((billing_event, decoded)); assert_eq!(budget.retained_bytes(), 0); } } } #[test] fn from_stream_fields_typed_none_still_clears_metadata_only_request_facts() { for limit in [0, 8192] { let mut source = legacy_metadata_only_event(); source.data.request_body_state = Some(UsageBodyCaptureState::None); source.data.provider_request_body_state = Some(UsageBodyCaptureState::None); let fields = source.to_stream_fields().expect("wire serialization"); let budget = Arc::new(EventCaptureMemoryBudget::new(limit)); let decoded = UsageEvent::from_stream_fields_with_capture_budget(&fields, Arc::clone(&budget)) .expect("wire decode"); let metadata = decoded .data .request_metadata .as_ref() .expect("response facts remain"); for key in [ "requested_reasoning_effort", "provider_reasoning_effort", "provider_service_tier", "provider_cache_ttl_minutes", ] { assert!(metadata.get(key).is_none(), "typed none must clear {key}"); } assert_eq!(metadata["provider_actual_service_tier"], "default"); assert_eq!( decoded.data.request_body_state, Some(UsageBodyCaptureState::None) ); assert_eq!( decoded.data.provider_request_body_state, Some(UsageBodyCaptureState::None) ); assert_eq!(decoded.data.cache_read_input_tokens, Some(0)); assert_eq!(budget.retained_bytes(), 0); assert_eq!(budget.downgraded_total(), 0); } } fn legacy_metadata_only_event() -> UsageEvent { UsageEvent::new( UsageEventType::Completed, "legacy-metadata-only", UsageEventData { provider_name: "openai".to_string(), model: "gpt-5.6-sol".to_string(), endpoint_api_format: Some("openai:responses".to_string()), input_tokens: Some(0), output_tokens: Some(0), total_tokens: Some(0), cache_read_input_tokens: Some(0), cache_creation_input_tokens: Some(0), actual_total_cost_usd: Some(0.0), request_metadata: Some(json!({ "requested_reasoning_effort": "high", "provider_reasoning_effort": "medium", "provider_service_tier": "priority", "provider_actual_service_tier": "default", "provider_cache_ttl_minutes": 60 })), ..UsageEventData::default() }, ) } #[test] fn from_stream_fields_body_omission_keeps_unknown_usage_unknown() { let budget = Arc::new(EventCaptureMemoryBudget::new(0)); for event_type in [ UsageEventType::Completed, UsageEventType::Failed, UsageEventType::Cancelled, ] { let event = UsageEvent::new( event_type, "usage-unavailable", UsageEventData { provider_name: "openai".to_string(), model: "gpt-5".to_string(), response_body: Some(json!({"error": "usage unavailable"})), request_metadata: Some(json!({ "usage_available": false, "usage_pricing_available": false })), ..UsageEventData::default() }, ); let fields = event.to_stream_fields().expect("wire serialization"); let decoded = UsageEvent::from_stream_fields_with_capture_budget(&fields, Arc::clone(&budget)) .expect("wire decode"); assert_eq!(decoded.event_type, event_type); assert_eq!(decoded.data.input_tokens, None); assert_eq!(decoded.data.output_tokens, None); assert_eq!(decoded.data.total_tokens, None); assert_eq!(decoded.data.cache_read_input_tokens, None); assert_eq!(decoded.data.cache_creation_input_tokens, None); assert_eq!(decoded.data.actual_total_cost_usd, None); assert_eq!( decoded.data.request_metadata.as_ref().expect("metadata")["usage_available"], false ); assert_eq!( decoded.data.request_metadata.as_ref().expect("metadata") ["usage_pricing_available"], false ); assert_eq!( decoded.data.response_body_state, Some(UsageBodyCaptureState::Truncated) ); } assert_eq!(budget.retained_bytes(), 0); assert_eq!(budget.downgraded_total(), 3); } #[test] fn from_stream_fields_invalid_envelopes_do_not_reserve_capture_memory() { let budget = Arc::new(EventCaptureMemoryBudget::new(1024)); let mut unsupported = captured_event() .to_stream_fields() .expect("wire serialization"); let mut payload: serde_json::Value = serde_json::from_str(&unsupported["payload"]).expect("json"); payload["v"] = json!(99); unsupported.insert("payload".to_string(), payload.to_string()); for fields in [ BTreeMap::new(), BTreeMap::from([("payload".to_string(), "not json".to_string())]), unsupported, ] { assert!(matches!( UsageEvent::from_stream_fields_with_capture_budget(&fields, Arc::clone(&budget)), Err(DataLayerError::UnexpectedValue(_)) )); assert_eq!(budget.retained_bytes(), 0); assert_eq!(budget.downgraded_total(), 0); } } #[test] fn event_capture_budget_basic_policy_needs_no_diagnostic_allocation() { let mut event = captured_event(); let budget = Arc::new(EventCaptureMemoryBudget::new(0)); apply_usage_body_capture_policy_to_event(UsageBodyCapturePolicy::default(), &mut event); event.data.apply_capture_memory_budget(Arc::clone(&budget)); assert_eq!( event.data.response_body_state, Some(UsageBodyCaptureState::Disabled) ); assert_eq!(event.data.total_tokens, Some(600)); assert_eq!(budget.retained_bytes(), 0); assert_eq!(budget.downgraded_total(), 0); } #[test] fn usage_event_round_trips_through_stream_fields() { let event = UsageEvent::new( UsageEventType::Completed, "req-1", UsageEventData { provider_name: "OpenAI".to_string(), model: "gpt-5".to_string(), input_tokens: Some(10), output_tokens: Some(20), ..UsageEventData::default() }, ); let fields = event.to_stream_fields().expect("event should serialize"); let parsed = UsageEvent::from_stream_fields(&fields).expect("event should parse"); assert_eq!(parsed.request_id, "req-1"); assert_eq!(parsed.event_type, UsageEventType::Completed); assert_eq!(parsed.data.total_tokens, None); assert_eq!(parsed.data.output_tokens, Some(20)); } }