use std::collections::BTreeMap; use std::io::{self, Write}; use aether_data_contracts::repository::usage::{ resolve_provider_cache_ttl_minutes, UsageBodyCaptureState, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY, }; use aether_data_contracts::DataLayerError; use serde::ser::{Impossible, SerializeMap, SerializeStruct}; use serde::{Serialize, Serializer}; use serde_json::Value; use super::{BorrowedUsageEventEnvelope, UsageEvent, UsageEventData, USAGE_EVENT_VERSION}; use crate::body_capture::mark_usage_event_capture_truncated; use crate::request_metadata::{ attach_client_request_body_metadata, attach_provider_request_body_metadata, attach_provider_response_body_metadata, attach_provider_response_model_metadata, clear_client_request_body_metadata, clear_provider_request_body_metadata, request_body_derived_facts_action, RequestBodyDerivedFactsAction, }; const DIAGNOSTIC_FIELDS: [&str; 8] = [ "request_body", "provider_request_body", "response_body", "client_response_body", "request_headers", "provider_request_headers", "response_headers", "client_response_headers", ]; const BODY_STATE_FIELDS: [&str; 4] = [ "request_body_state", "provider_request_body_state", "response_body_state", "client_response_body_state", ]; const BODY_METADATA_KEYS: [&str; 4] = ["request", "provider_request", "response", "client_response"]; #[derive(Debug)] pub(crate) struct EncodedUsageEvent { pub(crate) fields: BTreeMap, pub(crate) diagnostics_omitted: bool, } pub(super) fn encode( event: &UsageEvent, max_bytes: usize, ) -> Result { let mut writer = BoundedJsonWriter::new(max_bytes); if writer.serialize(&envelope(event, &event.data))? { return writer.into_event(false); } // Conservatively reject oversized original metadata before cloning it, even // when later fact normalization could make that metadata smaller. let core = ProjectedData { data: &event.data, overrides: None, }; if !writer.serialize(&envelope(event, &core))? { return Err(wire_limit_error(max_bytes)); } let overrides = WireOverrides::new(&event.data)?; let projected = ProjectedData { data: &event.data, overrides: Some(&overrides), }; if !writer.serialize(&envelope(event, &projected))? { return Err(wire_limit_error(max_bytes)); } writer.into_event(true) } fn envelope<'a, T: Serialize + ?Sized>( event: &'a UsageEvent, data: &'a T, ) -> BorrowedUsageEventEnvelope<'a, T> { BorrowedUsageEventEnvelope { v: USAGE_EVENT_VERSION, event_type: event.event_type, request_id: &event.request_id, timestamp_ms: event.timestamp_ms, data, } } fn wire_limit_error(max_bytes: usize) -> DataLayerError { DataLayerError::InvalidInput(format!( "usage event exceeds the {max_bytes}-byte wire limit after omitting diagnostic bodies and headers" )) } struct BoundedJsonWriter { bytes: Vec, max_bytes: usize, exceeded: bool, } impl BoundedJsonWriter { fn new(max_bytes: usize) -> Self { Self { bytes: Vec::new(), max_bytes, exceeded: false, } } fn serialize(&mut self, value: &T) -> Result { self.bytes.clear(); self.exceeded = false; match serde_json::to_writer(&mut *self, value) { Ok(()) => Ok(true), Err(_) if self.exceeded => Ok(false), Err(error) => Err(DataLayerError::UnexpectedValue(format!( "failed to serialize usage event payload: {error}" ))), } } fn into_event(self, diagnostics_omitted: bool) -> Result { let payload = String::from_utf8(self.bytes).map_err(|error| { DataLayerError::UnexpectedValue(format!("usage event JSON was not UTF-8: {error}")) })?; Ok(EncodedUsageEvent { fields: BTreeMap::from([("payload".to_string(), payload)]), diagnostics_omitted, }) } } impl Write for BoundedJsonWriter { fn write(&mut self, bytes: &[u8]) -> io::Result { if bytes.len() > self.max_bytes.saturating_sub(self.bytes.len()) { self.exceeded = true; return Err(io::Error::other("usage event wire limit exceeded")); } let required = self.bytes.len() + bytes.len(); if required > self.bytes.capacity() { let capacity = required .max(self.bytes.capacity().saturating_mul(2)) .min(self.max_bytes); self.bytes .try_reserve_exact(capacity - self.bytes.len()) .map_err(|error| { io::Error::other(format!("usage event wire allocation failed: {error}")) })?; } self.bytes.extend_from_slice(bytes); Ok(bytes.len()) } fn flush(&mut self) -> io::Result<()> { Ok(()) } } struct WireOverrides { truncated: [bool; 4], metadata: Option, } impl WireOverrides { fn new(data: &UsageEventData) -> Result { // The full v1 consumer decodes a JSON null Option as no body. let request_body = data.request_body.as_ref().filter(|body| !body.is_null()); let provider_request_body = data .provider_request_body .as_ref() .filter(|body| !body.is_null()); let body_cache_ttl = resolve_provider_cache_ttl_minutes( data.endpoint_api_format .as_deref() .or(data.api_format.as_deref()), data.target_model.as_deref().or(Some(data.model.as_str())), Some(data.model.as_str()), provider_request_body, ); if body_cache_ttl.is_some() && data.provider_request_body_state == Some(UsageBodyCaptureState::None) { return Err(DataLayerError::InvalidInput( "usage event cannot omit a provider request body whose cache TTL would be cleared by its explicit none capture state" .to_string(), )); } let mut metadata = data.request_metadata.clone(); match request_body_derived_facts_action(request_body, data.request_body_state) { RequestBodyDerivedFactsAction::Refresh => { if request_body.is_some_and(|body| !body.is_object()) { if let Some(Value::Object(object)) = metadata.as_mut() { object.remove(REQUESTED_REASONING_EFFORT_METADATA_KEY); } } else { metadata = attach_client_request_body_metadata(metadata, request_body); } } RequestBodyDerivedFactsAction::Clear if request_body.is_some() || data.request_body_state.is_some() => { metadata = clear_client_request_body_metadata(metadata); } RequestBodyDerivedFactsAction::Clear | RequestBodyDerivedFactsAction::Preserve => {} } match request_body_derived_facts_action( provider_request_body, data.provider_request_body_state, ) { RequestBodyDerivedFactsAction::Refresh => { if provider_request_body.is_some_and(|body| !body.is_object()) { // An authoritative scalar/array has no tier or reasoning, // but billing still falls back to the metadata's cache TTL. if let Some(Value::Object(object)) = metadata.as_mut() { object.remove(PROVIDER_REASONING_EFFORT_METADATA_KEY); object.remove(PROVIDER_SERVICE_TIER_METADATA_KEY); } } else { metadata = attach_provider_request_body_metadata( metadata, data.endpoint_api_format .as_deref() .or(data.api_format.as_deref()), data.target_model.as_deref().or(Some(data.model.as_str())), Some(data.model.as_str()), provider_request_body, ); } } RequestBodyDerivedFactsAction::Clear if provider_request_body.is_some() || data.provider_request_body_state.is_some() => { metadata = clear_provider_request_body_metadata(metadata); } RequestBodyDerivedFactsAction::Clear | RequestBodyDerivedFactsAction::Preserve => {} } metadata = attach_provider_response_body_metadata(metadata, data.response_body.as_ref()); metadata = attach_provider_response_model_metadata( metadata, request_body, data.request_body_state, data.api_format.as_deref(), data.response_body.as_ref(), data.response_body_state, data.endpoint_api_format.as_deref(), ); // Billing reads raw-body TTL before metadata regardless of capture state. // Preserve that precedence independently of reasoning and tier authority. if let Some(cache_ttl) = body_cache_ttl { let object = metadata .get_or_insert_with(|| Value::Object(serde_json::Map::new())) .as_object_mut() .ok_or_else(|| { DataLayerError::InvalidInput( "usage event cannot preserve provider cache TTL in non-object metadata after omitting diagnostic bodies" .to_string(), ) })?; object.insert( PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY.to_string(), Value::Number(cache_ttl.into()), ); } let truncated = [ (data.request_body.as_ref(), data.request_body_state), ( data.provider_request_body.as_ref(), data.provider_request_body_state, ), (data.response_body.as_ref(), data.response_body_state), ( data.client_response_body.as_ref(), data.client_response_body_state, ), ] .map(|(body, state)| { body.is_some_and(|body| !body.is_null()) && !matches!( state, Some( UsageBodyCaptureState::None | UsageBodyCaptureState::Disabled | UsageBodyCaptureState::Unavailable ) ) }); for (truncated, key) in truncated.into_iter().zip(BODY_METADATA_KEYS) { if !truncated { continue; } mark_usage_event_capture_truncated(&mut metadata, key); if let Some(entry) = metadata .as_mut() .and_then(|value| value.get_mut("body_capture")) .and_then(|value| value.get_mut(key)) .and_then(Value::as_object_mut) { entry.insert( "reason".to_string(), Value::String("wire_limit_exceeded".to_string()), ); } } Ok(Self { truncated, metadata, }) } } struct ProjectedData<'a> { data: &'a UsageEventData, overrides: Option<&'a WireOverrides>, } impl Serialize for ProjectedData<'_> { fn serialize(&self, serializer: S) -> Result { self.data.serialize(FieldProjectionSerializer { map: serializer.serialize_map(None)?, overrides: self.overrides, }) } } // Reuse UsageEventData's derived field traversal, including future fields and // skip_serializing_if rules. Only diagnostic fields and explicit overrides differ. struct FieldProjectionSerializer<'a, M> { map: M, overrides: Option<&'a WireOverrides>, } impl SerializeStruct for FieldProjectionSerializer<'_, M> { type Ok = M::Ok; type Error = M::Error; fn serialize_field( &mut self, key: &'static str, value: &T, ) -> Result<(), Self::Error> { if DIAGNOSTIC_FIELDS.contains(&key) { return Ok(()); } if let Some(overrides) = self.overrides { if key == "request_metadata" || BODY_STATE_FIELDS .iter() .zip(overrides.truncated) .any(|(state, truncated)| *state == key && truncated) { return Ok(()); } } self.map.serialize_entry(key, value) } fn end(mut self) -> Result { if let Some(overrides) = self.overrides { for (key, truncated) in BODY_STATE_FIELDS.into_iter().zip(overrides.truncated) { if truncated { self.map .serialize_entry(key, &UsageBodyCaptureState::Truncated)?; } } if let Some(metadata) = overrides.metadata.as_ref() { self.map.serialize_entry("request_metadata", metadata)?; } } self.map.end() } } fn expected_struct() -> Result { Err(E::custom("usage event data must serialize as a struct")) } macro_rules! reject_scalar_serialization { ($($name:ident($value:ident: $ty:ty)),* $(,)?) => { $(fn $name(self, $value: $ty) -> Result { let _ = $value; expected_struct() })* }; } impl Serializer for FieldProjectionSerializer<'_, M> { type Ok = M::Ok; type Error = M::Error; type SerializeSeq = Impossible; type SerializeTuple = Impossible; type SerializeTupleStruct = Impossible; type SerializeTupleVariant = Impossible; type SerializeMap = Impossible; type SerializeStruct = Self; type SerializeStructVariant = Impossible; reject_scalar_serialization! { serialize_bool(value: bool), serialize_i8(value: i8), serialize_i16(value: i16), serialize_i32(value: i32), serialize_i64(value: i64), serialize_i128(value: i128), serialize_u8(value: u8), serialize_u16(value: u16), serialize_u32(value: u32), serialize_u64(value: u64), serialize_u128(value: u128), serialize_f32(value: f32), serialize_f64(value: f64), serialize_char(value: char), serialize_str(value: &str), serialize_bytes(value: &[u8]), } fn serialize_none(self) -> Result { expected_struct() } fn serialize_some(self, _: &T) -> Result { expected_struct() } fn serialize_unit(self) -> Result { expected_struct() } fn serialize_unit_struct(self, _: &'static str) -> Result { expected_struct() } fn serialize_unit_variant( self, _: &'static str, _: u32, _: &'static str, ) -> Result { expected_struct() } fn serialize_newtype_struct( self, _: &'static str, _: &T, ) -> Result { expected_struct() } fn serialize_newtype_variant( self, _: &'static str, _: u32, _: &'static str, _: &T, ) -> Result { expected_struct() } fn serialize_seq(self, _: Option) -> Result { expected_struct() } fn serialize_tuple(self, _: usize) -> Result { expected_struct() } fn serialize_tuple_struct( self, _: &'static str, _: usize, ) -> Result { expected_struct() } fn serialize_tuple_variant( self, _: &'static str, _: u32, _: &'static str, _: usize, ) -> Result { expected_struct() } fn serialize_map(self, _: Option) -> Result { expected_struct() } fn serialize_struct( self, _: &'static str, _: usize, ) -> Result { Ok(self) } fn serialize_struct_variant( self, _: &'static str, _: u32, _: &'static str, _: usize, ) -> Result { expected_struct() } } #[cfg(test)] mod tests { use std::sync::Arc; use serde_json::json; use super::*; use crate::event::UsageEventType; use crate::event_capture_budget::EventCaptureMemoryBudget; fn event() -> UsageEvent { UsageEvent { event_type: UsageEventType::Completed, request_id: "wire-request".to_string(), timestamp_ms: 1_234_567, data: UsageEventData { provider_name: "provider".to_string(), model: "gpt-5.6-sol".to_string(), endpoint_api_format: Some("openai:responses".to_string()), user_id: Some("user-id".to_string()), api_key_id: Some("key-id".to_string()), provider_id: Some("provider-id".to_string()), provider_endpoint_id: Some("endpoint-id".to_string()), provider_api_key_id: Some("provider-key-id".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), cache_creation_ephemeral_5m_input_tokens: Some(0), cache_creation_ephemeral_1h_input_tokens: Some(25), cache_read_cost_usd: Some(0.0), total_cost_usd: Some(1.25), actual_total_cost_usd: Some(1.25), status_code: Some(200), is_stream: Some(false), candidate_index: Some(0), first_byte_time_ms: Some(0), error_message: Some(String::new()), request_metadata: Some(json!({ "plan_usage_reservation_token": "550e8400-e29b-41d4-a716-446655440000", "dimensions": {"image_count": 2, "size": "1024x1024", "quality": "high"}, "usage_available": true, "usage_pricing_available": true })), ..UsageEventData::default() }, } } fn wire_value(encoded: &EncodedUsageEvent) -> Value { serde_json::from_str(&encoded.fields["payload"]).expect("complete JSON wire payload") } #[test] fn event_wire_exact_limit_accounts_for_json_escaping_and_utf8() { let mut event = event(); event.request_id = "escaped\0\n\r\t\"\\\u{03bb}\u{1f600}".to_string(); let original = event.to_stream_fields().expect("original wire payload"); let length = original["payload"].len(); for limit in [length, length + 1] { let encoded = event.to_bounded_stream_fields(limit).expect("exact fit"); assert_eq!(encoded.fields, original); assert!(!encoded.diagnostics_omitted); } for limit in [0, 1, length - 1] { assert!(matches!( event.to_bounded_stream_fields(limit), Err(DataLayerError::InvalidInput(_)) )); } } #[test] fn event_wire_writer_does_not_append_past_limit_and_reuses_failed_buffer() { let mut writer = BoundedJsonWriter::new(5); writer.write_all(b"12345").expect("exact fit"); assert!(writer.write_all(b"6").is_err()); assert_eq!(writer.bytes, b"12345"); assert!(writer.exceeded); assert!(!writer.serialize(&"\u{0000}").expect("size rejection")); assert!(writer.bytes.len() <= 5); assert!(writer.serialize(&"\u{03bb}").expect("valid UTF-8 retry")); assert_eq!(writer.bytes, "\"\u{03bb}\"".as_bytes()); assert!(!writer.exceeded); } #[test] fn event_wire_projection_transparently_forwards_derived_fields() { #[derive(Serialize)] struct FutureFields<'a> { request_body: &'a Value, new_billing_field: &'a Value, zero_count: u64, enabled: bool, #[serde(skip_serializing_if = "Option::is_none")] missing_field: Option<&'a str>, } let diagnostic = json!({"large": "x".repeat(8_192)}); let billing = json!({"nested": [0, false, "unchanged"]}); let data = FutureFields { request_body: &diagnostic, new_billing_field: &billing, zero_count: 0, enabled: false, missing_field: None, }; let mut bytes = Vec::new(); let mut serializer = serde_json::Serializer::new(&mut bytes); data.serialize(FieldProjectionSerializer { map: (&mut serializer) .serialize_map(None) .expect("object serializer"), overrides: None, }) .expect("project derived fields"); assert_eq!( serde_json::from_slice::(&bytes).expect("projected JSON"), json!({"new_billing_field": billing, "zero_count": 0, "enabled": false}) ); } #[test] fn event_wire_omission_preserves_billing_facts_refs_and_source_ownership() { let mut event = event(); let padding = "x".repeat(16_384); event.data.request_body = Some(json!({"reasoning": {"effort": "high"}, "input": padding})); event.data.provider_request_body = Some( json!({"model": "gpt-5.6-sol", "reasoning": {"effort": "medium"}, "service_tier": "priority", "input": padding}), ); event.data.response_body = Some(json!({"service_tier": "Default", "output": padding})); event.data.client_response_body = Some(json!({"output": padding})); event.data.request_headers = Some(json!({"x-request": padding})); event.data.provider_request_headers = Some(json!({"x-provider": padding})); event.data.response_headers = Some(json!({"x-response": padding})); event.data.client_response_headers = Some(json!({"x-client": padding})); event.data.request_body_ref = Some("usage://wire-request/request_body".to_string()); event.data.request_body_state = Some(UsageBodyCaptureState::Inline); event.data.provider_request_body_state = Some(UsageBodyCaptureState::Inline); event.data.response_body_state = Some(UsageBodyCaptureState::Inline); event.data.request_metadata.as_mut().unwrap()["body_capture"] = json!({ "response": {"state": "inline", "source_bytes": 123_456} }); let weight = event.data.capture_heap_estimate(); let budget = Arc::new(EventCaptureMemoryBudget::new(weight)); event.data.apply_capture_memory_budget(Arc::clone(&budget)); let before = serde_json::to_value(&event).expect("source snapshot"); let original_wire: Value = serde_json::from_str(&event.to_stream_fields().unwrap()["payload"]).unwrap(); let encoded = event .to_bounded_stream_fields(8_192) .expect("diagnostic omission"); assert!(encoded.diagnostics_omitted); assert!(encoded.fields["payload"].len() <= 8_192); let value = wire_value(&encoded); for field in DIAGNOSTIC_FIELDS { assert!( value["data"].get(field).is_none(), "{field} must be omitted" ); } for (field, original) in original_wire["data"].as_object().unwrap() { if !DIAGNOSTIC_FIELDS.contains(&field.as_str()) && !BODY_STATE_FIELDS.contains(&field.as_str()) && field != "request_metadata" { assert_eq!(&value["data"][field], original, "core field {field}"); } } for (field, key) in BODY_STATE_FIELDS.into_iter().zip(BODY_METADATA_KEYS) { assert_eq!(value["data"][field], "truncated"); assert_eq!( value["data"]["request_metadata"]["body_capture"][key]["reason"], "wire_limit_exceeded" ); assert_eq!( value["data"]["request_metadata"]["body_capture"][key]["stored_bytes"], 0 ); } let metadata = &value["data"]["request_metadata"]; 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["body_capture"]["response"]["source_bytes"], 123_456 ); assert_eq!( metadata["dimensions"], before["data"]["request_metadata"]["dimensions"] ); assert_eq!(serde_json::to_value(&event).unwrap(), before); assert_eq!(budget.retained_bytes(), weight); assert_eq!(budget.downgraded_total(), 0); let decoded = UsageEvent::from_stream_fields(&encoded.fields).expect("wire decode"); let record = crate::build_upsert_usage_record_from_event(&decoded).expect("record mapping"); 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_ephemeral_5m_input_tokens, Some(0)); assert_eq!(record.cache_creation_ephemeral_1h_input_tokens, Some(25)); assert_eq!(record.error_message.as_deref(), Some("")); assert_eq!(record.request_body_ref, event.data.request_body_ref); assert_eq!( record.request_body_state, Some(UsageBodyCaptureState::Truncated) ); let metadata = record.request_metadata.expect("record billing metadata"); assert_eq!(metadata["provider_cache_ttl_minutes"], 30); assert_eq!(metadata["dimensions"]["image_count"], 2); drop(event); assert_eq!(budget.retained_bytes(), 0); } #[test] fn event_wire_preserves_explicit_capture_states_and_legacy_metadata() { for state in [ None, Some(UsageBodyCaptureState::None), Some(UsageBodyCaptureState::Disabled), Some(UsageBodyCaptureState::Unavailable), Some(UsageBodyCaptureState::Reference), ] { let mut event = event(); event.data.request_body_state = state; event.data.provider_request_body_state = state; event.data.response_body_state = state; event.data.client_response_body_state = state; event.data.request_body_ref = Some("usage://wire-request/request_body".to_string()); event.data.response_headers = Some(json!({"large-header": "x".repeat(16_384)})); event.data.request_metadata = Some(json!({ "requested_reasoning_effort": "high", "provider_reasoning_effort": "medium", "provider_service_tier": "priority", "provider_cache_ttl_minutes": 60, "provider_actual_service_tier": "flex" })); let before = serde_json::to_value(&event).unwrap(); let encoded = event .to_bounded_stream_fields(4_096) .expect("headers omitted"); let value = wire_value(&encoded); for field in BODY_STATE_FIELDS { assert_eq!(value["data"].get(field), before["data"].get(field)); } let metadata = &value["data"]["request_metadata"]; if state == Some(UsageBodyCaptureState::None) { assert!(metadata.get("requested_reasoning_effort").is_none()); assert!(metadata.get("provider_cache_ttl_minutes").is_none()); } else { assert_eq!(metadata["requested_reasoning_effort"], "high"); assert_eq!(metadata["provider_cache_ttl_minutes"], 60); } assert_eq!(metadata["provider_actual_service_tier"], "flex"); assert_eq!( value["data"]["request_body_ref"], before["data"]["request_body_ref"] ); assert_eq!(serde_json::to_value(&event).unwrap(), before); if matches!( state, Some( UsageBodyCaptureState::None | UsageBodyCaptureState::Disabled | UsageBodyCaptureState::Unavailable ) ) { event.data.endpoint_api_format = Some("claude:messages".to_string()); event.data.request_body = Some(json!({"reasoning": {"effort": "low"}})); event.data.provider_request_body = Some(json!({"reasoning": {"effort": "low"}, "service_tier": "default"})); event.data.response_body = Some(json!({"service_tier": "priority"})); event.data.client_response_body = Some(json!({"stale": true})); let baseline = UsageEvent::from_stream_fields_with_capture_budget( &event.to_stream_fields().expect("full v1 fields"), Arc::new(EventCaptureMemoryBudget::new(usize::MAX)), ) .expect("full v1 consumer baseline"); let with_stale_bodies = wire_value( &event .to_bounded_stream_fields(4_096) .expect("explicit capture states override stale bodies"), ); for field in BODY_STATE_FIELDS { assert_eq!( with_stale_bodies["data"].get(field), value["data"].get(field) ); } assert_eq!( with_stale_bodies["data"]["request_metadata"], serde_json::to_value(&baseline.data.request_metadata).unwrap(), "wire omission must preserve the full v1 consumer's billing facts" ); } } } #[test] fn event_wire_preserves_raw_body_cache_ttl_across_capture_states() { use aether_data_contracts::repository::usage::extract_provider_cache_ttl_minutes_from_metadata; for state in [ None, Some(UsageBodyCaptureState::Inline), Some(UsageBodyCaptureState::Reference), Some(UsageBodyCaptureState::Truncated), Some(UsageBodyCaptureState::Disabled), Some(UsageBodyCaptureState::Unavailable), ] { let mut event = event(); event.data.provider_request_body_state = state; event.data.provider_request_body = Some(json!({ "prompt_cache_options": {"ttl": "30m"}, "service_tier": "default", "reasoning": {"effort": "low"} })); event.data.response_headers = Some(json!({"large": "x".repeat(16_384)})); event.data.request_metadata = Some(json!({ "provider_cache_ttl_minutes": 60, "provider_service_tier": "priority", "provider_reasoning_effort": "high" })); let before = serde_json::to_value(&event).unwrap(); let baseline = UsageEvent::from_stream_fields_with_capture_budget( &event.to_stream_fields().expect("full v1 fields"), Arc::new(EventCaptureMemoryBudget::new(usize::MAX)), ) .expect("full v1 consumer baseline"); let encoded = event.to_bounded_stream_fields(4_096).expect("body omitted"); assert!(encoded.diagnostics_omitted); let decoded = UsageEvent::from_stream_fields_with_capture_budget( &encoded.fields, Arc::new(EventCaptureMemoryBudget::new(usize::MAX)), ) .expect("projected consumer event"); assert!(decoded.data.provider_request_body.is_none()); assert_eq!( extract_provider_cache_ttl_minutes_from_metadata( decoded.data.request_metadata.as_ref() ), Some(30), "raw body TTL must win over stale metadata for {state:?}" ); for field in ["provider_service_tier", "provider_reasoning_effort"] { assert_eq!( decoded.data.request_metadata.as_ref().unwrap().get(field), baseline.data.request_metadata.as_ref().unwrap().get(field), "TTL preservation must not change {field} authority for {state:?}" ); } assert_eq!(serde_json::to_value(&event).unwrap(), before); } } #[test] fn event_wire_non_object_bodies_match_full_consumer_facts() { use aether_data_contracts::repository::usage::{ extract_provider_cache_ttl_minutes_from_metadata, resolve_provider_service_tier_from_request_capture, }; for body in [ json!(null), json!("opaque"), json!([]), json!(42), json!(true), ] { for state in [ None, Some(UsageBodyCaptureState::None), Some(UsageBodyCaptureState::Inline), Some(UsageBodyCaptureState::Reference), Some(UsageBodyCaptureState::Truncated), Some(UsageBodyCaptureState::Disabled), Some(UsageBodyCaptureState::Unavailable), ] { let mut event = event(); event.data.request_body = Some(body.clone()); event.data.provider_request_body = Some(body.clone()); event.data.request_body_state = state; event.data.provider_request_body_state = state; event.data.response_headers = Some(json!({"large": "x".repeat(16_384)})); event.data.request_metadata = Some(json!({ "requested_reasoning_effort": "medium", "provider_reasoning_effort": "high", "provider_service_tier": "priority", "provider_cache_ttl_minutes": 60, "unrelated": {"preserved": true} })); let original = event.to_stream_fields().expect("full v1 fields"); let baseline = UsageEvent::from_stream_fields_with_capture_budget( &original, Arc::new(EventCaptureMemoryBudget::new(usize::MAX)), ) .expect("full v1 consumer baseline"); let encoded = event .to_bounded_stream_fields(4_096) .expect("omit diagnostics"); assert!(encoded.diagnostics_omitted); let decoded = UsageEvent::from_stream_fields_with_capture_budget( &encoded.fields, Arc::new(EventCaptureMemoryBudget::new(usize::MAX)), ) .expect("projected consumer event"); let tier = |data: &UsageEventData| { resolve_provider_service_tier_from_request_capture( data.provider_request_body.as_ref(), data.provider_request_body_state, data.request_metadata.as_ref(), ) }; assert_eq!( tier(&decoded.data), tier(&baseline.data), "provider tier for {body:?}, {state:?}" ); let metadata = decoded.data.request_metadata.as_ref().unwrap(); let baseline_metadata = baseline.data.request_metadata.as_ref().unwrap(); assert_eq!( extract_provider_cache_ttl_minutes_from_metadata(Some(metadata)), extract_provider_cache_ttl_minutes_from_metadata(Some(baseline_metadata)), "metadata TTL fallback for {body:?}, {state:?}" ); let authoritative = !body.is_null() && matches!( state, None | Some( UsageBodyCaptureState::Inline | UsageBodyCaptureState::Reference ) ); for field in [ REQUESTED_REASONING_EFFORT_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY, ] { assert_eq!( metadata.get(field), if authoritative { None } else { baseline_metadata.get(field) }, "reasoning authority for {body:?}, {state:?}, {field}" ); } assert_eq!(metadata["unrelated"], baseline_metadata["unrelated"]); if body.is_null() { assert_eq!( decoded.data.request_body_state, baseline.data.request_body_state ); assert_eq!( decoded.data.provider_request_body_state, baseline.data.provider_request_body_state ); } assert_eq!(event.to_stream_fields().unwrap(), original); } } } #[test] fn event_wire_rejects_cache_ttl_loss_from_explicit_none_capture_state() { let mut event = event(); event.data.provider_request_body_state = Some(UsageBodyCaptureState::None); event.data.provider_request_body = Some(json!({ "prompt_cache_options": {"ttl": "30m"}, "large": "x".repeat(16_384) })); let original = event.to_stream_fields().expect("full v1 fields"); let full = event .to_bounded_stream_fields(original["payload"].len()) .expect("complete diagnostics remain representable"); assert_eq!(full.fields, original); assert!(!full.diagnostics_omitted); assert!(matches!( event.to_bounded_stream_fields(4_096), Err(DataLayerError::InvalidInput(_)) )); assert_eq!(event.to_stream_fields().unwrap(), original); } #[test] fn event_wire_rejects_oversized_core_and_post_projection_metadata() { for field in ["request_metadata", "error_message"] { let mut event = event(); if field == "request_metadata" { event.data.request_metadata = Some(json!({"large": "x".repeat(32_768)})); } else { event.data.error_message = Some("x".repeat(32_768)); } let error = event .to_bounded_stream_fields(1_024) .expect_err("oversized core must fail"); assert!(matches!(error, DataLayerError::InvalidInput(_))); assert!( error.to_string().len() < 200, "errors must not include payload contents" ); } let mut event = event(); event.data.response_body = Some(json!({"large": "x".repeat(8_192)})); let core = ProjectedData { data: &event.data, overrides: None, }; let core_size = serde_json::to_vec(&envelope(&event, &core)).unwrap().len(); assert!( matches!( event.to_bounded_stream_fields(core_size), Err(DataLayerError::InvalidInput(_)) ), "added capture metadata must also fit the exact wire limit" ); } }