mirror of
https://github.com/fawney19/Aether.git
synced 2026-10-07 01:47:47 +08:00
fix: harden concurrency limits and high-RPM runtime paths
Bound request, stream, queue, and shutdown resource lifetimes. Reduce scheduler and Redis hot-path work and isolate database maintenance. Include regression coverage, load probes, and concurrency audit results.
This commit is contained in:
elky
@@ -1,4 +1,5 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
|
||||
@@ -6,8 +7,18 @@ 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 {
|
||||
@@ -18,7 +29,7 @@ pub enum UsageEventType {
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
|
||||
#[derive(Debug, PartialEq, Serialize, Deserialize, Default)]
|
||||
pub struct UsageEventData {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub user_id: Option<String>,
|
||||
@@ -144,6 +155,171 @@ pub struct UsageEventData {
|
||||
pub local_execution_runtime_miss_reason: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub request_metadata: Option<Value>,
|
||||
#[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::<Value>())
|
||||
.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<EventCaptureMemoryBudget>,
|
||||
) {
|
||||
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)]
|
||||
@@ -164,6 +340,16 @@ struct UsageEventEnvelope {
|
||||
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,
|
||||
@@ -179,12 +365,12 @@ impl UsageEvent {
|
||||
}
|
||||
|
||||
pub fn to_stream_fields(&self) -> Result<BTreeMap<String, String>, DataLayerError> {
|
||||
let payload = UsageEventEnvelope {
|
||||
let payload = BorrowedUsageEventEnvelope {
|
||||
v: USAGE_EVENT_VERSION,
|
||||
event_type: self.event_type,
|
||||
request_id: self.request_id.clone(),
|
||||
request_id: &self.request_id,
|
||||
timestamp_ms: self.timestamp_ms,
|
||||
data: self.data.clone(),
|
||||
data: &self.data,
|
||||
};
|
||||
let payload = serde_json::to_string(&payload).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
@@ -194,7 +380,21 @@ impl UsageEvent {
|
||||
Ok(BTreeMap::from([("payload".to_string(), payload)]))
|
||||
}
|
||||
|
||||
pub(crate) fn to_bounded_stream_fields(
|
||||
&self,
|
||||
max_bytes: usize,
|
||||
) -> Result<EncodedUsageEvent, DataLayerError> {
|
||||
wire::encode(self, max_bytes)
|
||||
}
|
||||
|
||||
pub fn from_stream_fields(fields: &BTreeMap<String, String>) -> Result<Self, DataLayerError> {
|
||||
Self::from_stream_fields_with_capture_budget(fields, shared_capture_memory_budget())
|
||||
}
|
||||
|
||||
pub(crate) fn from_stream_fields_with_capture_budget(
|
||||
fields: &BTreeMap<String, String>,
|
||||
budget: Arc<EventCaptureMemoryBudget>,
|
||||
) -> Result<Self, DataLayerError> {
|
||||
let payload = fields.get("payload").ok_or_else(|| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"usage event stream entry missing payload field".to_string(),
|
||||
@@ -212,12 +412,16 @@ impl UsageEvent {
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -230,8 +434,564 @@ pub fn now_ms() -> 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(
|
||||
|
||||
Reference in New Issue
Block a user