i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor usage body capture and stream terminal reporting

Browse Source
This commit is contained in:
fawney19
2026-04-18 17:48:21 +08:00
parent 569242d72f
commit 3363592751
36 changed files with 2673 additions and 512 deletions
Download Patch File Download Diff File Expand all files Collapse all files

659
crates/aether-usage-runtime/src/body_capture.rs Normal file
View File

@@ -0,0 +1,659 @@
use aether_data_contracts::repository::usage::{
UpsertUsageRecord, UsageBodyCaptureState, UsageBodyField,
};
use serde_json::{json, Map, Value};
use crate::event::UsageEvent;
use crate::runtime::{UsageBodyCapturePolicy, UsageRequestRecordLevel};
const TRUNCATED_BODY_STRING_SUFFIX: &str = "...[truncated]";
#[derive(Debug)]
struct LimitedUsageBodyCapture {
value: Value,
source_bytes: Option<u64>,
stored_bytes: Option<u64>,
truncated: bool,
reason: Option<&'static str>,
}
struct UsageBodyCapturePayloadMut<'a> {
request_body: &'a mut Option<Value>,
request_body_ref: &'a mut Option<String>,
request_body_state: &'a mut Option<UsageBodyCaptureState>,
provider_request_body: &'a mut Option<Value>,
provider_request_body_ref: &'a mut Option<String>,
provider_request_body_state: &'a mut Option<UsageBodyCaptureState>,
response_body: &'a mut Option<Value>,
response_body_ref: &'a mut Option<String>,
response_body_state: &'a mut Option<UsageBodyCaptureState>,
client_response_body: &'a mut Option<Value>,
client_response_body_ref: &'a mut Option<String>,
client_response_body_state: &'a mut Option<UsageBodyCaptureState>,
request_metadata: &'a mut Option<Value>,
}
impl<'a> UsageBodyCapturePayloadMut<'a> {
fn from_event(event: &'a mut UsageEvent) -> Self {
Self {
request_body: &mut event.data.request_body,
request_body_ref: &mut event.data.request_body_ref,
request_body_state: &mut event.data.request_body_state,
provider_request_body: &mut event.data.provider_request_body,
provider_request_body_ref: &mut event.data.provider_request_body_ref,
provider_request_body_state: &mut event.data.provider_request_body_state,
response_body: &mut event.data.response_body,
response_body_ref: &mut event.data.response_body_ref,
response_body_state: &mut event.data.response_body_state,
client_response_body: &mut event.data.client_response_body,
client_response_body_ref: &mut event.data.client_response_body_ref,
client_response_body_state: &mut event.data.client_response_body_state,
request_metadata: &mut event.data.request_metadata,
}
}
fn from_record(record: &'a mut UpsertUsageRecord) -> Self {
Self {
request_body: &mut record.request_body,
request_body_ref: &mut record.request_body_ref,
request_body_state: &mut record.request_body_state,
provider_request_body: &mut record.provider_request_body,
provider_request_body_ref: &mut record.provider_request_body_ref,
provider_request_body_state: &mut record.provider_request_body_state,
response_body: &mut record.response_body,
response_body_ref: &mut record.response_body_ref,
response_body_state: &mut record.response_body_state,
client_response_body: &mut record.client_response_body,
client_response_body_ref: &mut record.client_response_body_ref,
client_response_body_state: &mut record.client_response_body_state,
request_metadata: &mut record.request_metadata,
}
}
}
#[derive(Debug, Clone, Copy)]
pub struct UsageBodyCaptureEngine {
policy: UsageBodyCapturePolicy,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct RuntimeBodyCaptureStates {
pub request: UsageBodyCaptureState,
pub provider_request: UsageBodyCaptureState,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct RuntimeBodyCaptureMetadataInput<'a> {
pub request_has_inline_body: bool,
pub request_body_ref: Option<&'a str>,
pub provider_request_has_inline_body: bool,
pub provider_request_body_ref: Option<&'a str>,
pub provider_request_source_bytes: Option<u64>,
pub provider_request_unavailable: bool,
pub provider_request_unavailable_reason: Option<&'a str>,
}
impl UsageBodyCaptureEngine {
pub fn new(policy: UsageBodyCapturePolicy) -> Self {
Self { policy }
}
pub fn apply_to_event(self, event: &mut UsageEvent) {
self.apply_to_payload(UsageBodyCapturePayloadMut::from_event(event));
}
pub fn apply_to_record(self, record: &mut UpsertUsageRecord) {
self.apply_to_payload(UsageBodyCapturePayloadMut::from_record(record));
}
fn apply_to_payload(self, payload: UsageBodyCapturePayloadMut<'_>) {
if matches!(self.policy.record_level, UsageRequestRecordLevel::Basic) {
disable_usage_body_capture_field(
UsageBodyField::RequestBody,
"request",
payload.request_body,
payload.request_body_ref,
payload.request_body_state,
payload.request_metadata,
);
disable_usage_body_capture_field(
UsageBodyField::ProviderRequestBody,
"provider_request",
payload.provider_request_body,
payload.provider_request_body_ref,
payload.provider_request_body_state,
payload.request_metadata,
);
disable_usage_body_capture_field(
UsageBodyField::ResponseBody,
"response",
payload.response_body,
payload.response_body_ref,
payload.response_body_state,
payload.request_metadata,
);
disable_usage_body_capture_field(
UsageBodyField::ClientResponseBody,
"client_response",
payload.client_response_body,
payload.client_response_body_ref,
payload.client_response_body_state,
payload.request_metadata,
);
return;
}
apply_usage_body_capture_limit(
UsageBodyField::RequestBody,
"request",
self.policy.max_request_body_bytes,
payload.request_body,
payload.request_body_ref,
payload.request_body_state,
payload.request_metadata,
);
apply_usage_body_capture_limit(
UsageBodyField::ProviderRequestBody,
"provider_request",
self.policy.max_request_body_bytes,
payload.provider_request_body,
payload.provider_request_body_ref,
payload.provider_request_body_state,
payload.request_metadata,
);
apply_usage_body_capture_limit(
UsageBodyField::ResponseBody,
"response",
self.policy.max_response_body_bytes,
payload.response_body,
payload.response_body_ref,
payload.response_body_state,
payload.request_metadata,
);
apply_usage_body_capture_limit(
UsageBodyField::ClientResponseBody,
"client_response",
self.policy.max_response_body_bytes,
payload.client_response_body,
payload.client_response_body_ref,
payload.client_response_body_state,
payload.request_metadata,
);
}
}
pub fn apply_usage_body_capture_policy_to_event(
policy: UsageBodyCapturePolicy,
event: &mut UsageEvent,
) {
UsageBodyCaptureEngine::new(policy).apply_to_event(event);
}
pub fn apply_usage_body_capture_policy_to_record(
policy: UsageBodyCapturePolicy,
record: &mut UpsertUsageRecord,
) {
UsageBodyCaptureEngine::new(policy).apply_to_record(record);
}
fn disable_usage_body_capture_field(
field: UsageBodyField,
metadata_key: &str,
body: &mut Option<Value>,
body_ref: &mut Option<String>,
state: &mut Option<UsageBodyCaptureState>,
request_metadata: &mut Option<Value>,
) {
*body = None;
*body_ref = None;
*state = Some(UsageBodyCaptureState::Disabled);
sync_usage_body_ref_metadata(request_metadata, field, None);
upsert_body_capture_metadata_value_entry(
request_metadata,
metadata_key,
Some(UsageBodyCaptureState::Disabled),
None,
None,
Some("request_record_level_basic"),
);
}
fn apply_usage_body_capture_limit(
field: UsageBodyField,
metadata_key: &str,
max_bytes: Option<usize>,
body: &mut Option<Value>,
body_ref: &mut Option<String>,
state: &mut Option<UsageBodyCaptureState>,
request_metadata: &mut Option<Value>,
) {
*body_ref = sanitize_usage_body_ref(body_ref.take());
if body.is_some() && body_ref.is_some() {
*body = None;
}
if let Some(body_ref_value) = body_ref.as_ref() {
*state = Some(UsageBodyCaptureState::Reference);
sync_usage_body_ref_metadata(request_metadata, field, Some(body_ref_value));
upsert_body_capture_metadata_value_entry(
request_metadata,
metadata_key,
Some(UsageBodyCaptureState::Reference),
None,
None,
None,
);
return;
}
let Some(value) = body.take() else {
if matches!(state, Some(UsageBodyCaptureState::Unavailable)) {
upsert_body_capture_metadata_value_entry(
request_metadata,
metadata_key,
*state,
None,
None,
None,
);
} else if state.is_none() {
*state = Some(UsageBodyCaptureState::None);
}
sync_usage_body_ref_metadata(request_metadata, field, None);
return;
};
let limited = limit_usage_body_capture_value(value, max_bytes);
let next_state = if limited.truncated {
UsageBodyCaptureState::Truncated
} else {
UsageBodyCaptureState::Inline
};
*state = Some(next_state);
*body = Some(limited.value);
sync_usage_body_ref_metadata(request_metadata, field, None);
upsert_body_capture_metadata_value_entry(
request_metadata,
metadata_key,
Some(next_state),
limited.stored_bytes,
limited.source_bytes,
limited.reason,
);
}
fn limit_usage_body_capture_value(
value: Value,
max_bytes: Option<usize>,
) -> LimitedUsageBodyCapture {
let source_bytes = serde_json::to_vec(&value)
.ok()
.map(|bytes| bytes.len() as u64);
let Some(limit) = max_bytes.filter(|value| *value > 0) else {
return LimitedUsageBodyCapture {
stored_bytes: source_bytes,
source_bytes,
value,
truncated: false,
reason: None,
};
};
let Some(source_len) = source_bytes else {
return LimitedUsageBodyCapture {
stored_bytes: None,
source_bytes: None,
value,
truncated: false,
reason: None,
};
};
if source_len <= limit as u64 {
return LimitedUsageBodyCapture {
stored_bytes: Some(source_len),
source_bytes: Some(source_len),
value,
truncated: false,
reason: None,
};
}
let truncated_value = match value {
Value::String(text) => Value::String(truncate_usage_body_string(&text, limit)),
other => json!({
"truncated": true,
"reason": "body_capture_limit_exceeded",
"max_bytes": limit,
"source_bytes": source_len,
"value_kind": usage_value_kind(&other),
}),
};
let stored_bytes = serde_json::to_vec(&truncated_value)
.ok()
.map(|bytes| bytes.len() as u64);
LimitedUsageBodyCapture {
value: truncated_value,
source_bytes: Some(source_len),
stored_bytes,
truncated: true,
reason: Some("body_capture_limit_exceeded"),
}
}
fn truncate_usage_body_string(value: &str, max_bytes: usize) -> String {
if serde_json::to_vec(&Value::String(value.to_string()))
.ok()
.is_some_and(|bytes| bytes.len() <= max_bytes)
{
return value.to_string();
}
let mut end = value.len();
while end > 0 {
while end > 0 && !value.is_char_boundary(end) {
end -= 1;
}
let mut candidate = value[..end].to_string();
candidate.push_str(TRUNCATED_BODY_STRING_SUFFIX);
if serde_json::to_vec(&Value::String(candidate.clone()))
.ok()
.is_some_and(|bytes| bytes.len() <= max_bytes)
{
return candidate;
}
end = value[..end]
.char_indices()
.last()
.map(|(index, _)| index)
.unwrap_or(0);
if end == 0 {
break;
}
}
json!({
"truncated": true,
"reason": "body_capture_limit_exceeded",
"max_bytes": max_bytes,
"value_kind": "string",
})
.to_string()
}
pub(crate) fn sync_usage_body_ref_metadata(
metadata: &mut Option<Value>,
field: UsageBodyField,
body_ref: Option<&str>,
) {
let Some(body_ref) = body_ref.map(str::trim).filter(|value| !value.is_empty()) else {
if let Some(object) = metadata.as_mut().and_then(Value::as_object_mut) {
object.remove(field.as_ref_key());
}
return;
};
let object = metadata
.get_or_insert_with(|| Value::Object(Map::new()))
.as_object_mut();
let Some(object) = object else {
return;
};
object.insert(
field.as_ref_key().to_string(),
Value::String(body_ref.to_string()),
);
}
pub(crate) fn build_payload_body_capture_metadata(
provider_body_base64: Option<&str>,
client_body_base64: Option<&str>,
provider_body_state: Option<UsageBodyCaptureState>,
client_body_state: Option<UsageBodyCaptureState>,
) -> Option<Value> {
let mut metadata = Map::new();
if let Some(decoded_len) = provider_body_base64.and_then(decoded_base64_len_hint) {
metadata.insert(
"provider_response_body_base64_bytes".to_string(),
Value::Number(decoded_len.into()),
);
}
if let Some(decoded_len) = client_body_base64.and_then(decoded_base64_len_hint) {
metadata.insert(
"client_response_body_base64_bytes".to_string(),
Value::Number(decoded_len.into()),
);
}
let mut body_capture = Map::new();
append_body_capture_metadata_entry(
&mut body_capture,
"response",
provider_body_state,
provider_body_base64.and_then(decoded_base64_len_hint),
provider_body_base64.and_then(decoded_base64_len_hint),
);
append_body_capture_metadata_entry(
&mut body_capture,
"client_response",
client_body_state,
client_body_base64.and_then(decoded_base64_len_hint),
client_body_base64.and_then(decoded_base64_len_hint),
);
if !body_capture.is_empty() {
metadata.insert("body_capture".to_string(), Value::Object(body_capture));
}
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
pub(crate) fn build_runtime_body_capture_states(
request_has_inline_body: bool,
request_body_ref: Option<&str>,
provider_request_has_inline_body: bool,
provider_request_body_ref: Option<&str>,
provider_request_unavailable: bool,
) -> RuntimeBodyCaptureStates {
RuntimeBodyCaptureStates {
request: UsageBodyCaptureState::from_capture_parts(
request_has_inline_body,
request_body_ref.is_some(),
false,
),
provider_request: UsageBodyCaptureState::from_capture_parts(
provider_request_has_inline_body,
provider_request_body_ref.is_some(),
provider_request_unavailable,
),
}
}
pub(crate) fn append_runtime_body_capture_metadata(
metadata: &mut Map<String, Value>,
input: RuntimeBodyCaptureMetadataInput<'_>,
) {
let states = build_runtime_body_capture_states(
input.request_has_inline_body,
input.request_body_ref,
input.provider_request_has_inline_body,
input.provider_request_body_ref,
input.provider_request_unavailable,
);
upsert_body_capture_metadata_entry(metadata, "request", Some(states.request), None, None, None);
upsert_body_capture_metadata_entry(
metadata,
"provider_request",
Some(states.provider_request),
input.provider_request_source_bytes,
input.provider_request_source_bytes,
input.provider_request_unavailable_reason,
);
}
pub(crate) fn build_plan_body_capture_metadata(
provider_request_body_base64: Option<&str>,
) -> Option<Value> {
let mut metadata = Map::new();
append_plan_body_capture_metadata(&mut metadata, provider_request_body_base64);
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
pub(crate) fn append_plan_body_capture_metadata(
metadata: &mut Map<String, Value>,
provider_request_body_base64: Option<&str>,
) {
if let Some(body_bytes_b64) = provider_request_body_base64 {
let decoded_len = decoded_base64_len_hint(body_bytes_b64);
if let Some(decoded_len) = decoded_len {
metadata.insert(
"provider_request_body_base64_bytes".to_string(),
Value::Number(decoded_len.into()),
);
}
upsert_body_capture_metadata_entry(
metadata,
"provider_request",
Some(UsageBodyCaptureState::Unavailable),
decoded_len,
decoded_len,
Some("body_bytes_base64_only"),
);
}
}
fn append_body_capture_metadata_entry(
target: &mut Map<String, Value>,
key: &str,
state: Option<UsageBodyCaptureState>,
stored_bytes: Option<u64>,
source_bytes: Option<u64>,
) {
let Some(state) = state else {
return;
};
let mut entry = Map::new();
entry.insert(
"state".to_string(),
Value::String(state.as_str().to_string()),
);
if let Some(stored_bytes) = stored_bytes {
entry.insert("stored_bytes".to_string(), json!(stored_bytes));
}
if let Some(source_bytes) = source_bytes {
entry.insert("source_bytes".to_string(), json!(source_bytes));
}
if matches!(state, UsageBodyCaptureState::Truncated) {
entry.insert(
"reason".to_string(),
Value::String("body_capture_limit_exceeded".to_string()),
);
}
target.insert(key.to_string(), Value::Object(entry));
}
pub(crate) fn upsert_body_capture_metadata_entry(
metadata: &mut Map<String, Value>,
key: &str,
state: Option<UsageBodyCaptureState>,
stored_bytes: Option<u64>,
source_bytes: Option<u64>,
reason: Option<&str>,
) {
let body_capture = metadata
.entry("body_capture".to_string())
.or_insert_with(|| Value::Object(Map::new()));
let Some(body_capture_object) = body_capture.as_object_mut() else {
return;
};
let Some(state) = state else {
return;
};
let mut entry = Map::new();
entry.insert(
"state".to_string(),
Value::String(state.as_str().to_string()),
);
if let Some(bytes) = stored_bytes {
entry.insert("stored_bytes".to_string(), json!(bytes));
}
if let Some(bytes) = source_bytes {
entry.insert("source_bytes".to_string(), json!(bytes));
}
if let Some(reason) = reason {
entry.insert("reason".to_string(), Value::String(reason.to_string()));
}
body_capture_object.insert(key.to_string(), Value::Object(entry));
}
fn upsert_body_capture_metadata_value_entry(
metadata: &mut Option<Value>,
key: &str,
state: Option<UsageBodyCaptureState>,
stored_bytes: Option<u64>,
source_bytes: Option<u64>,
reason: Option<&str>,
) {
let Some(state) = state else {
return;
};
let metadata_object = metadata
.get_or_insert_with(|| Value::Object(Map::new()))
.as_object_mut();
let Some(metadata_object) = metadata_object else {
return;
};
upsert_body_capture_metadata_entry(
metadata_object,
key,
Some(state),
stored_bytes,
source_bytes,
reason,
);
}
pub(crate) fn decoded_base64_len_hint(body_base64: &str) -> Option<u64> {
let body_base64 = body_base64.trim();
if body_base64.is_empty() {
return None;
}
let usable_len = body_base64.len();
if usable_len % 4 == 1 {
return None;
}
let padding = body_base64
.chars()
.rev()
.take_while(|char| *char == '=')
.count();
let full_quads = usable_len / 4;
let remainder = usable_len % 4;
let base_len = full_quads.saturating_mul(3);
let remainder_len = match remainder {
0 => 0,
2 => 1,
3 => 2,
_ => return None,
};
let decoded_len = base_len
.saturating_add(remainder_len)
.saturating_sub(padding.min(2));
Some(decoded_len as u64)
}
fn sanitize_usage_body_ref(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
}
fn usage_value_kind(value: &Value) -> &'static str {
match value {
Value::Null => "null",
Value::Bool(_) => "bool",
Value::Number(_) => "number",
Value::String(_) => "string",
Value::Array(_) => "array",
Value::Object(_) => "object",
}
}

9
crates/aether-usage-runtime/src/event.rs
View File

@@ -1,6 +1,7 @@
use std::collections::BTreeMap;
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;
@@ -96,24 +97,32 @@ pub struct UsageEventData {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_body_ref: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_body_state: Option<UsageBodyCaptureState>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_request_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_request_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_request_body_ref: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_request_body_state: Option<UsageBodyCaptureState>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_body_ref: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_body_state: Option<UsageBodyCaptureState>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_response_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_response_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_response_body_ref: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_response_body_state: Option<UsageBodyCaptureState>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub candidate_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub candidate_index: Option<u64>,

9
crates/aether-usage-runtime/src/lib.rs
View File

@@ -1,3 +1,4 @@
mod body_capture;
pub mod config;
pub mod event;
mod executor;
@@ -13,6 +14,10 @@ pub mod usage_mapper;
pub mod worker;
pub mod write;
pub use body_capture::{
apply_usage_body_capture_policy_to_event, apply_usage_body_capture_policy_to_record,
UsageBodyCaptureEngine,
};
pub use config::UsageRuntimeConfig;
pub use event::{now_ms, UsageEvent, UsageEventData, UsageEventType, USAGE_EVENT_VERSION};
pub use queue::UsageQueue;
@@ -31,7 +36,9 @@ pub use report_context::{
build_locally_actionable_report_context_from_video_task, report_context_is_locally_actionable,
};
pub use runtime::{
UsageBillingEventEnricher, UsageRequestRecordLevel, UsageRuntime, UsageRuntimeAccess,
UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageRequestRecordLevel, UsageRuntime,
UsageRuntimeAccess, DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES,
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
};
pub use settlement::{settle_usage_if_needed, UsageSettlementWriter};
pub use standardized_usage::StandardizedUsage;

4
crates/aether-usage-runtime/src/record.rs
View File

@@ -76,20 +76,24 @@ pub fn build_upsert_usage_record_from_event(
request_body: data.request_body,
request_body_ref: empty_to_none(data.request_body_ref)
.or_else(|| metadata_string(data.request_metadata.as_ref(), "request_body_ref")),
request_body_state: data.request_body_state,
provider_request_headers: data.provider_request_headers,
provider_request_body: data.provider_request_body,
provider_request_body_ref: empty_to_none(data.provider_request_body_ref).or_else(|| {
metadata_string(data.request_metadata.as_ref(), "provider_request_body_ref")
}),
provider_request_body_state: data.provider_request_body_state,
response_headers: data.response_headers,
response_body: data.response_body,
response_body_ref: empty_to_none(data.response_body_ref)
.or_else(|| metadata_string(data.request_metadata.as_ref(), "response_body_ref")),
response_body_state: data.response_body_state,
client_response_headers: data.client_response_headers,
client_response_body: data.client_response_body,
client_response_body_ref: empty_to_none(data.client_response_body_ref).or_else(|| {
metadata_string(data.request_metadata.as_ref(), "client_response_body_ref")
}),
client_response_body_state: data.client_response_body_state,
candidate_id: data
.candidate_id
.or_else(|| metadata_string(data.request_metadata.as_ref(), "candidate_id")),

9
crates/aether-usage-runtime/src/report.rs
View File

@@ -1,6 +1,7 @@
use std::collections::BTreeMap;
use aether_contracts::ExecutionTelemetry;
use aether_contracts::{ExecutionStreamTerminalSummary, ExecutionTelemetry};
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
use base64::Engine as _;
use serde::{Deserialize, Serialize};
use serde_json::Value;
@@ -37,8 +38,14 @@ pub struct GatewayStreamReportRequest {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_body_base64: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_body_state: Option<UsageBodyCaptureState>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_body_base64: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_body_state: Option<UsageBodyCaptureState>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub terminal_summary: Option<ExecutionStreamTerminalSummary>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub telemetry: Option<ExecutionTelemetry>,
}

92
crates/aether-usage-runtime/src/runtime.rs
View File

@@ -11,6 +11,7 @@ use tracing::warn;
use crate::executor::spawn_on_usage_background_runtime;
use crate::{
apply_usage_body_capture_policy_to_event, apply_usage_body_capture_policy_to_record,
build_pending_usage_record_from_seed, build_stream_terminal_usage_seed,
build_streaming_usage_record_from_seed, build_sync_terminal_usage_seed,
build_terminal_usage_event_from_seed, build_upsert_usage_record_from_event,
@@ -31,6 +32,26 @@ pub enum UsageRequestRecordLevel {
Full,
}
pub const DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES: usize = 5 * 1024 * 1024;
pub const DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES: usize = 5 * 1024 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct UsageBodyCapturePolicy {
pub record_level: UsageRequestRecordLevel,
pub max_request_body_bytes: Option<usize>,
pub max_response_body_bytes: Option<usize>,
}
impl Default for UsageBodyCapturePolicy {
fn default() -> Self {
Self {
record_level: UsageRequestRecordLevel::Full,
max_request_body_bytes: Some(DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES),
max_response_body_bytes: Some(DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES),
}
}
}
#[async_trait]
pub trait UsageRuntimeAccess:
UsageRecordWriter + UsageSettlementWriter + UsageBillingEventEnricher + Send + Sync
@@ -39,8 +60,12 @@ pub trait UsageRuntimeAccess:
fn has_usage_worker_runner(&self) -> bool;
fn usage_worker_runner(&self) -> Option<RedisStreamRunner>;
async fn body_capture_policy(&self) -> Result<UsageBodyCapturePolicy, DataLayerError> {
Ok(UsageBodyCapturePolicy::default())
}
async fn request_record_level(&self) -> Result<UsageRequestRecordLevel, DataLayerError> {
Ok(UsageRequestRecordLevel::Full)
Ok(self.body_capture_policy().await?.record_level)
}
}
@@ -114,7 +139,8 @@ impl UsageRuntime {
spawn_on_usage_background_runtime(boxed_usage_task(async move {
let now_unix_secs = now_unix_secs();
match build_pending_usage_record_offthread(&seed, now_unix_secs).await {
Ok(record) => {
Ok(mut record) => {
apply_body_capture_policy_to_record_from_data(&data, &mut record).await;
if let Err(err) = data.upsert_usage_record(record).await {
warn!(
event_name = "usage_pending_record_failed",
@@ -164,7 +190,8 @@ impl UsageRuntime {
)
.await
{
Ok(record) => {
Ok(mut record) => {
apply_body_capture_policy_to_record_from_data(&data, &mut record).await;
if let Err(err) = data.upsert_usage_record(record).await {
warn!(
event_name = "usage_stream_record_failed",
@@ -209,7 +236,7 @@ impl UsageRuntime {
spawn_on_usage_background_runtime(boxed_usage_task(async move {
match build_sync_terminal_usage_event_offthread(input).await {
Ok(mut event) => {
apply_request_record_level_from_data(&data, &mut event).await;
apply_body_capture_policy_from_data(&data, &mut event).await;
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_sync_terminal_billing_enrichment_failed",
@@ -257,7 +284,7 @@ impl UsageRuntime {
spawn_on_usage_background_runtime(boxed_usage_task(async move {
match build_stream_terminal_usage_event_offthread(input).await {
Ok(mut event) => {
apply_request_record_level_from_data(&data, &mut event).await;
apply_body_capture_policy_from_data(&data, &mut event).await;
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_stream_terminal_billing_enrichment_failed",
@@ -303,7 +330,7 @@ impl UsageRuntime {
if !self.is_enabled() {
return;
}
apply_request_record_level_from_data(data, &mut event).await;
apply_body_capture_policy_from_data(data, &mut event).await;
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_terminal_billing_enrichment_failed",
@@ -446,38 +473,44 @@ fn join_error_to_data_layer(err: tokio::task::JoinError) -> DataLayerError {
DataLayerError::UnexpectedValue(format!("usage builder task join failed: {err}"))
}
async fn apply_request_record_level_from_data<T>(data: &T, event: &mut UsageEvent)
async fn apply_body_capture_policy_from_data<T>(data: &T, event: &mut UsageEvent)
where
T: UsageRuntimeAccess,
{
match data.request_record_level().await {
Ok(level) => apply_request_record_level(level, event),
match data.body_capture_policy().await {
Ok(policy) => apply_usage_body_capture_policy_to_event(policy, event),
Err(err) => {
warn!(
event_name = "usage_request_record_level_read_failed",
event_name = "usage_body_capture_policy_read_failed",
log_type = "event",
request_id = %event.request_id,
fallback = "full",
fallback = "default",
error = %err,
"usage runtime failed to read request record level; keeping full capture"
"usage runtime failed to read body capture policy; keeping default capture"
);
apply_usage_body_capture_policy_to_event(UsageBodyCapturePolicy::default(), event);
}
}
}
fn apply_request_record_level(level: UsageRequestRecordLevel, event: &mut UsageEvent) {
if !matches!(level, UsageRequestRecordLevel::Basic) {
return;
async fn apply_body_capture_policy_to_record_from_data<T>(data: &T, record: &mut UpsertUsageRecord)
where
T: UsageRuntimeAccess,
{
match data.body_capture_policy().await {
Ok(policy) => apply_usage_body_capture_policy_to_record(policy, record),
Err(err) => {
warn!(
event_name = "usage_body_capture_policy_read_failed",
log_type = "event",
request_id = %record.request_id,
fallback = "default",
error = %err,
"usage runtime failed to read body capture policy; keeping default capture"
);
apply_usage_body_capture_policy_to_record(UsageBodyCapturePolicy::default(), record);
}
}
event.data.request_body = None;
event.data.request_body_ref = None;
event.data.provider_request_body = None;
event.data.provider_request_body_ref = None;
event.data.response_body = None;
event.data.response_body_ref = None;
event.data.client_response_body = None;
event.data.client_response_body_ref = None;
}
fn boxed_usage_task<F>(task: F) -> Pin<Box<dyn Future<Output = ()> + Send>>
@@ -498,7 +531,8 @@ fn now_unix_secs() -> u64 {
mod tests {
use serde_json::json;
use super::{apply_request_record_level, UsageRequestRecordLevel};
use super::{UsageBodyCapturePolicy, UsageRequestRecordLevel};
use crate::apply_usage_body_capture_policy_to_event;
use crate::{UsageEvent, UsageEventData, UsageEventType};
#[test]
@@ -527,7 +561,13 @@ mod tests {
},
);
apply_request_record_level(UsageRequestRecordLevel::Basic, &mut event);
apply_usage_body_capture_policy_to_event(
UsageBodyCapturePolicy {
record_level: UsageRequestRecordLevel::Basic,
..UsageBodyCapturePolicy::default()
},
&mut event,
);
assert_eq!(event.data.total_tokens, Some(42));
assert_eq!(event.data.error_message.as_deref(), Some("upstream failed"));

114
crates/aether-usage-runtime/src/standardized_usage.rs
View File

@@ -1,113 +1 @@
use std::collections::BTreeMap;
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, Default)]
pub struct StandardizedUsage {
pub input_tokens: i64,
pub output_tokens: i64,
pub cache_creation_tokens: i64,
pub cache_creation_ephemeral_5m_tokens: i64,
pub cache_creation_ephemeral_1h_tokens: i64,
pub cache_read_tokens: i64,
pub reasoning_tokens: i64,
pub cache_storage_token_hours: f64,
pub request_count: i64,
pub dimensions: BTreeMap<String, serde_json::Value>,
}
impl StandardizedUsage {
pub fn new() -> Self {
Self {
request_count: 1,
..Self::default()
}
}
pub fn get(&self, field_name: &str) -> Option<serde_json::Value> {
match field_name {
"input_tokens" => Some(serde_json::json!(self.input_tokens)),
"output_tokens" => Some(serde_json::json!(self.output_tokens)),
"cache_creation_tokens" => Some(serde_json::json!(self.cache_creation_tokens)),
"cache_creation_ephemeral_5m_tokens" => {
Some(serde_json::json!(self.cache_creation_ephemeral_5m_tokens))
}
"cache_creation_ephemeral_1h_tokens" => {
Some(serde_json::json!(self.cache_creation_ephemeral_1h_tokens))
}
"cache_read_tokens" => Some(serde_json::json!(self.cache_read_tokens)),
"reasoning_tokens" => Some(serde_json::json!(self.reasoning_tokens)),
"cache_storage_token_hours" => Some(serde_json::json!(self.cache_storage_token_hours)),
"request_count" => Some(serde_json::json!(self.request_count)),
"extra" | "dimensions" => Some(serde_json::json!(self.dimensions)),
_ => self.dimensions.get(field_name).cloned(),
}
}
pub fn set(&mut self, field_name: &str, value: impl Into<serde_json::Value>) {
let value = value.into();
match field_name {
"input_tokens" => self.input_tokens = as_i64(&value, 0),
"output_tokens" => self.output_tokens = as_i64(&value, 0),
"cache_creation_tokens" => self.cache_creation_tokens = as_i64(&value, 0),
"cache_creation_ephemeral_5m_tokens" => {
self.cache_creation_ephemeral_5m_tokens = as_i64(&value, 0)
}
"cache_creation_ephemeral_1h_tokens" => {
self.cache_creation_ephemeral_1h_tokens = as_i64(&value, 0)
}
"cache_read_tokens" => self.cache_read_tokens = as_i64(&value, 0),
"reasoning_tokens" => self.reasoning_tokens = as_i64(&value, 0),
"cache_storage_token_hours" => self.cache_storage_token_hours = as_f64(&value, 0.0),
"request_count" => self.request_count = as_i64(&value, 0),
"extra" | "dimensions" => {
self.dimensions = match value {
serde_json::Value::Object(map) => map.into_iter().collect(),
_ => BTreeMap::new(),
}
}
_ => {
self.dimensions.insert(field_name.to_string(), value);
}
}
}
pub fn normalize_cache_creation_breakdown(mut self) -> Self {
if self.cache_creation_tokens <= 0 {
let derived = self
.cache_creation_ephemeral_5m_tokens
.saturating_add(self.cache_creation_ephemeral_1h_tokens);
if derived > 0 {
self.cache_creation_tokens = derived;
}
}
self
}
}
fn as_i64(value: &serde_json::Value, default: i64) -> i64 {
value
.as_i64()
.or_else(|| value.as_u64().and_then(|v| i64::try_from(v).ok()))
.unwrap_or(default)
}
fn as_f64(value: &serde_json::Value, default: f64) -> f64 {
value.as_f64().unwrap_or(default)
}
#[cfg(test)]
mod tests {
use super::StandardizedUsage;
#[test]
fn standardized_usage_reads_and_writes_known_and_extra_fields() {
let mut usage = StandardizedUsage::new();
usage.set("input_tokens", 10);
usage.set("custom_dimension", "value");
assert_eq!(usage.get("input_tokens"), Some(serde_json::json!(10)));
assert_eq!(
usage.get("custom_dimension"),
Some(serde_json::json!("value"))
);
}
}
pub use aether_contracts::StandardizedUsage;

274
crates/aether-usage-runtime/src/write.rs
View File

@@ -1,11 +1,16 @@
use std::collections::BTreeMap;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
use aether_data_contracts::repository::usage::UpsertUsageRecord;
use aether_contracts::{ExecutionPlan, ExecutionStreamTerminalSummary, ExecutionTelemetry};
use aether_data_contracts::repository::usage::{UpsertUsageRecord, UsageBodyCaptureState};
use aether_data_contracts::DataLayerError;
use base64::Engine as _;
use serde_json::{json, Map, Value};
use crate::body_capture::{
append_runtime_body_capture_metadata, build_payload_body_capture_metadata,
build_plan_body_capture_metadata, build_runtime_body_capture_states, decoded_base64_len_hint,
RuntimeBodyCaptureMetadataInput,
};
use crate::request_metadata::{
build_usage_request_metadata_seed, merge_usage_request_metadata,
sanitize_usage_request_metadata,
@@ -41,6 +46,14 @@ struct UsageBodyRefsSeed {
client_response_body_ref: Option<String>,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct UsageBodyStatesSeed {
request_body_state: Option<UsageBodyCaptureState>,
provider_request_body_state: Option<UsageBodyCaptureState>,
response_body_state: Option<UsageBodyCaptureState>,
client_response_body_state: Option<UsageBodyCaptureState>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct LifecycleUsageSeed {
pub request_id: String,
@@ -64,6 +77,7 @@ pub struct LifecycleUsageSeed {
pub has_format_conversion: Option<bool>,
pub is_stream: bool,
routing: UsageRoutingSeed,
body_states: UsageBodyStatesSeed,
pub request_metadata: Option<Value>,
}
@@ -97,6 +111,7 @@ pub struct TerminalUsageContextSeed {
pub provider_request_headers: Option<Value>,
pub provider_request: Option<Value>,
body_refs: UsageBodyRefsSeed,
body_states: UsageBodyStatesSeed,
routing: UsageRoutingSeed,
pub request_metadata: Option<Value>,
}
@@ -109,8 +124,10 @@ pub struct SyncTerminalUsagePayloadSeed {
pub first_byte_time_ms: Option<u64>,
pub provider_response_headers: Option<Value>,
pub provider_response_full: Option<Value>,
pub provider_response_body_state: Option<UsageBodyCaptureState>,
pub client_response_headers: Option<Value>,
pub client_response: Option<Value>,
pub client_response_body_state: Option<UsageBodyCaptureState>,
pub capture_metadata: Option<Value>,
}
@@ -122,8 +139,11 @@ pub struct StreamTerminalUsagePayloadSeed {
pub first_byte_time_ms: Option<u64>,
pub provider_response_headers: Option<Value>,
pub provider_response_full: Option<Value>,
pub provider_response_body_state: Option<UsageBodyCaptureState>,
pub client_response_headers: Option<Value>,
pub client_response: Option<Value>,
pub client_response_body_state: Option<UsageBodyCaptureState>,
pub terminal_summary: Option<ExecutionStreamTerminalSummary>,
pub capture_metadata: Option<Value>,
}
@@ -154,6 +174,7 @@ pub struct TerminalUsageSeed {
pub provider_request_headers: Option<Value>,
pub provider_request: Option<Value>,
body_refs: UsageBodyRefsSeed,
body_states: UsageBodyStatesSeed,
pub provider_response_headers: Option<Value>,
pub provider_response: Option<Value>,
pub client_response_headers: Option<Value>,
@@ -162,6 +183,7 @@ pub struct TerminalUsageSeed {
pub request_metadata: Option<Value>,
pub audit_payload: Option<Value>,
pub standardized_usage: Option<StandardizedUsage>,
pub terminal_summary: Option<ExecutionStreamTerminalSummary>,
}
pub type TerminalUsageOutcome = TerminalUsageSeed;
@@ -241,6 +263,7 @@ pub fn build_lifecycle_usage_seed(
has_format_conversion: context_bool(context, "needs_conversion"),
is_stream: plan.stream,
routing: build_runtime_routing_seed(plan, context),
body_states: build_runtime_body_states_seed(plan, context),
request_metadata: build_runtime_request_metadata_seed(plan, context),
}
}
@@ -391,15 +414,19 @@ pub fn build_terminal_usage_event_from_seed(
request_headers: seed.request_headers,
request_body: seed.request_body,
request_body_ref: body_refs.request_body_ref,
request_body_state: seed.body_states.request_body_state,
provider_request_headers: seed.provider_request_headers,
provider_request_body: seed.provider_request,
provider_request_body_ref: body_refs.provider_request_body_ref,
provider_request_body_state: seed.body_states.provider_request_body_state,
response_headers: seed.provider_response_headers,
response_body: provider_response.clone(),
response_body_ref: body_refs.response_body_ref,
response_body_state: seed.body_states.response_body_state,
client_response_headers: seed.client_response_headers,
client_response_body: client_response.clone(),
client_response_body_ref: body_refs.client_response_body_ref,
client_response_body_state: seed.body_states.client_response_body_state,
candidate_id: routing.candidate_id,
key_name: routing.key_name,
planner_kind: routing.planner_kind,
@@ -415,12 +442,14 @@ pub fn build_terminal_usage_event_from_seed(
apply_standardized_usage_seed(usage, &mut data);
}
if let Some(response_body) = provider_response.as_ref() {
apply_standardized_usage(
Some(seed.provider_contract.clone()),
response_body,
&mut data,
);
if seed.standardized_usage.is_none() {
if let Some(response_body) = provider_response.as_ref() {
apply_standardized_usage(
Some(seed.provider_contract.clone()),
response_body,
&mut data,
);
}
}
if data.total_tokens.is_none() {
if let Some(tokens) = provider_response
@@ -494,9 +523,10 @@ pub fn build_terminal_usage_context_seed(
provider_request: context_body_value(context, "provider_request_body")
.or_else(|| plan_json_body_capture_for_usage(plan)),
body_refs: build_runtime_body_refs_seed(plan, context),
body_states: build_runtime_body_states_seed(plan, context),
request_metadata: merge_usage_request_metadata(
build_usage_request_metadata_seed(plan, context),
build_plan_body_capture_metadata(plan),
build_plan_body_capture_metadata(plan.body.body_bytes_b64.as_deref()),
),
}
}
@@ -504,6 +534,12 @@ pub fn build_terminal_usage_context_seed(
pub fn build_sync_terminal_usage_payload_seed(
payload: &GatewaySyncReportRequest,
) -> SyncTerminalUsagePayloadSeed {
let provider_response_full = payload
.body_json
.as_ref()
.cloned()
.or_else(|| decode_body_for_storage(payload.body_base64.as_deref()));
let client_response = payload.client_body_json.as_ref().cloned();
SyncTerminalUsagePayloadSeed {
report_kind: payload.report_kind.clone(),
status_code: payload.status_code,
@@ -513,14 +549,33 @@ pub fn build_sync_terminal_usage_payload_seed(
.and_then(|value| value.elapsed_ms),
first_byte_time_ms: payload.telemetry.as_ref().and_then(|value| value.ttfb_ms),
provider_response_headers: Some(headers_to_json(&payload.headers)),
provider_response_full: payload
.body_json
.as_ref()
.cloned()
.or_else(|| decode_body_for_storage(payload.body_base64.as_deref())),
provider_response_full: provider_response_full.clone(),
provider_response_body_state: Some(UsageBodyCaptureState::from_capture_parts(
provider_response_full.is_some(),
false,
false,
)),
client_response_headers: Some(headers_to_json(&payload.headers)),
client_response: payload.client_body_json.as_ref().cloned(),
capture_metadata: build_payload_body_capture_metadata(payload.body_base64.as_deref(), None),
client_response: client_response.clone(),
client_response_body_state: Some(UsageBodyCaptureState::from_capture_parts(
client_response.is_some(),
false,
false,
)),
capture_metadata: build_payload_body_capture_metadata(
payload.body_base64.as_deref(),
None,
Some(UsageBodyCaptureState::from_capture_parts(
provider_response_full.is_some(),
false,
false,
)),
Some(UsageBodyCaptureState::from_capture_parts(
client_response.is_some(),
false,
false,
)),
),
}
}
@@ -537,11 +592,16 @@ pub fn build_stream_terminal_usage_payload_seed(
first_byte_time_ms: payload.telemetry.as_ref().and_then(|value| value.ttfb_ms),
provider_response_headers: Some(headers_to_json(&payload.headers)),
provider_response_full: decode_body_for_storage(payload.provider_body_base64.as_deref()),
provider_response_body_state: payload.provider_body_state,
client_response_headers: Some(headers_to_json(&payload.headers)),
client_response: decode_body_for_storage(payload.client_body_base64.as_deref()),
client_response_body_state: payload.client_body_state,
terminal_summary: payload.terminal_summary.clone(),
capture_metadata: build_payload_body_capture_metadata(
payload.provider_body_base64.as_deref(),
payload.client_body_base64.as_deref(),
payload.provider_body_state,
payload.client_body_state,
),
}
}
@@ -586,6 +646,12 @@ pub fn build_sync_terminal_usage_seed(
provider_request_headers: context_seed.provider_request_headers,
provider_request: context_seed.provider_request,
body_refs: context_seed.body_refs,
body_states: UsageBodyStatesSeed {
request_body_state: context_seed.body_states.request_body_state,
provider_request_body_state: context_seed.body_states.provider_request_body_state,
response_body_state: payload_seed.provider_response_body_state,
client_response_body_state: payload_seed.client_response_body_state,
},
routing: context_seed.routing,
provider_response_headers: payload_seed.provider_response_headers,
provider_response: payload_seed.provider_response_full,
@@ -594,6 +660,7 @@ pub fn build_sync_terminal_usage_seed(
request_metadata: context_seed.request_metadata,
audit_payload: payload_seed.capture_metadata,
standardized_usage,
terminal_summary: None,
}
}
@@ -603,9 +670,17 @@ pub fn build_stream_terminal_usage_seed(
cancelled: bool,
) -> TerminalUsageSeed {
let standardized_usage = payload_seed
.provider_response_full
.terminal_summary
.as_ref()
.map(|response| map_usage_from_response(response, context_seed.provider_contract.as_str()));
.and_then(|summary| summary.standardized_usage.clone())
.or_else(|| {
payload_seed
.provider_response_full
.as_ref()
.map(|response| {
map_usage_from_response(response, context_seed.provider_contract.as_str())
})
});
let terminal_state = infer_stream_terminal_state(
payload_seed.report_kind.as_str(),
payload_seed.status_code,
@@ -638,6 +713,12 @@ pub fn build_stream_terminal_usage_seed(
provider_request_headers: context_seed.provider_request_headers,
provider_request: context_seed.provider_request,
body_refs: context_seed.body_refs,
body_states: UsageBodyStatesSeed {
request_body_state: context_seed.body_states.request_body_state,
provider_request_body_state: context_seed.body_states.provider_request_body_state,
response_body_state: payload_seed.provider_response_body_state,
client_response_body_state: payload_seed.client_response_body_state,
},
routing: context_seed.routing,
provider_response_headers: payload_seed.provider_response_headers,
provider_response: payload_seed.provider_response_full,
@@ -646,6 +727,7 @@ pub fn build_stream_terminal_usage_seed(
request_metadata: context_seed.request_metadata,
audit_payload: payload_seed.capture_metadata,
standardized_usage,
terminal_summary: payload_seed.terminal_summary,
}
}
@@ -758,15 +840,19 @@ fn build_lifecycle_usage_record(
request_headers: None,
request_body: None,
request_body_ref: body_refs.request_body_ref,
request_body_state: seed.body_states.request_body_state,
provider_request_headers: None,
provider_request_body: None,
provider_request_body_ref: body_refs.provider_request_body_ref,
provider_request_body_state: seed.body_states.provider_request_body_state,
response_headers: sanitize_usage_header_capture(response_headers),
response_body: None,
response_body_ref: body_refs.response_body_ref,
response_body_state: Some(UsageBodyCaptureState::None),
client_response_headers: sanitize_usage_header_capture(client_response_headers),
client_response_body: None,
client_response_body_ref: body_refs.client_response_body_ref,
client_response_body_state: Some(UsageBodyCaptureState::None),
candidate_id: routing.candidate_id,
candidate_index: routing.candidate_index,
key_name: routing.key_name,
@@ -796,6 +882,7 @@ fn build_usage_event_data_seed_with_detail(
let context = report_context.and_then(Value::as_object);
let routing = build_runtime_routing_seed(plan, context);
let body_refs = build_runtime_body_refs_seed(plan, context);
let body_states = build_runtime_body_states_seed(plan, context);
let api_format = context_string(context, "client_api_format")
.or_else(|| non_empty_string(Some(plan.client_api_format.clone())));
let endpoint_api_format = context_string(context, "provider_api_format")
@@ -845,13 +932,17 @@ fn build_usage_event_data_seed_with_detail(
request_headers: context_usage_value(context, "original_headers"),
request_body: context_body_value(context, "original_request_body"),
request_body_ref: body_refs.request_body_ref,
request_body_state: body_states.request_body_state,
provider_request_headers: context_usage_value(context, "provider_request_headers")
.or_else(|| Some(headers_to_json(&plan.headers))),
provider_request_body: context_body_value(context, "provider_request_body")
.or_else(|| plan_json_body_capture_for_usage(plan)),
provider_request_body_ref: body_refs.provider_request_body_ref,
provider_request_body_state: body_states.provider_request_body_state,
response_body_ref: body_refs.response_body_ref,
response_body_state: body_states.response_body_state,
client_response_body_ref: body_refs.client_response_body_ref,
client_response_body_state: body_states.client_response_body_state,
candidate_id: routing.candidate_id,
candidate_index: routing.candidate_index,
key_name: routing.key_name,
@@ -1007,7 +1098,37 @@ fn build_runtime_request_metadata_seed(
if let Some(trace_id) = context_string(context, "trace_id") {
metadata.insert("trace_id".to_string(), Value::String(trace_id));
}
append_plan_body_capture_metadata(plan, &mut metadata);
let request_body = context_body_value(context, "original_request_body");
let request_body_ref = context_string(context, "request_body_ref");
let provider_request_body = context_body_value(context, "provider_request_body")
.or_else(|| plan_json_body_capture_for_usage(plan));
let provider_request_body_ref = context_string(context, "provider_request_body_ref")
.or_else(|| non_empty_string(plan.body.body_ref.clone()));
let provider_source_bytes = plan
.body
.body_bytes_b64
.as_deref()
.and_then(decoded_base64_len_hint);
append_runtime_body_capture_metadata(
&mut metadata,
RuntimeBodyCaptureMetadataInput {
request_has_inline_body: request_body.is_some(),
request_body_ref: request_body_ref.as_deref(),
provider_request_has_inline_body: provider_request_body.is_some(),
provider_request_body_ref: provider_request_body_ref.as_deref(),
provider_request_source_bytes: provider_source_bytes,
provider_request_unavailable: plan.body.body_bytes_b64.is_some(),
provider_request_unavailable_reason: plan
.body
.body_bytes_b64
.as_ref()
.map(|_| "body_bytes_base64_only"),
},
);
crate::body_capture::append_plan_body_capture_metadata(
&mut metadata,
plan.body.body_bytes_b64.as_deref(),
);
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
@@ -1026,23 +1147,6 @@ fn capture_usage_storage_value(value: Value) -> Value {
})
}
fn build_plan_body_capture_metadata(plan: &ExecutionPlan) -> Option<Value> {
let mut metadata = Map::new();
append_plan_body_capture_metadata(plan, &mut metadata);
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
fn append_plan_body_capture_metadata(plan: &ExecutionPlan, metadata: &mut Map<String, Value>) {
if let Some(body_bytes_b64) = plan.body.body_bytes_b64.as_deref() {
if let Some(decoded_len) = decoded_base64_len_hint(body_bytes_b64) {
metadata.insert(
"provider_request_body_base64_bytes".to_string(),
Value::Number(decoded_len.into()),
);
}
}
}
fn build_runtime_body_refs_seed(
plan: &ExecutionPlan,
context: Option<&Map<String, Value>>,
@@ -1056,6 +1160,32 @@ fn build_runtime_body_refs_seed(
}
}
fn build_runtime_body_states_seed(
plan: &ExecutionPlan,
context: Option<&Map<String, Value>>,
) -> UsageBodyStatesSeed {
let request_body = context_body_value(context, "original_request_body");
let request_body_ref = context_string(context, "request_body_ref");
let provider_request_body = context_body_value(context, "provider_request_body")
.or_else(|| plan_json_body_capture_for_usage(plan));
let provider_request_body_ref = context_string(context, "provider_request_body_ref")
.or_else(|| non_empty_string(plan.body.body_ref.clone()));
let states = build_runtime_body_capture_states(
request_body.is_some(),
request_body_ref.as_deref(),
provider_request_body.is_some(),
provider_request_body_ref.as_deref(),
plan.body.body_bytes_b64.is_some(),
);
UsageBodyStatesSeed {
request_body_state: Some(states.request),
provider_request_body_state: Some(states.provider_request),
response_body_state: Some(UsageBodyCaptureState::None),
client_response_body_state: Some(UsageBodyCaptureState::None),
}
}
fn merge_body_refs_seed_with_metadata(
seed: &UsageBodyRefsSeed,
metadata: Option<&Value>,
@@ -1081,26 +1211,6 @@ fn merge_body_refs_seed_with_metadata(
}
}
fn build_payload_body_capture_metadata(
provider_body_base64: Option<&str>,
client_body_base64: Option<&str>,
) -> Option<Value> {
let mut metadata = Map::new();
if let Some(decoded_len) = provider_body_base64.and_then(decoded_base64_len_hint) {
metadata.insert(
"provider_response_body_base64_bytes".to_string(),
Value::Number(decoded_len.into()),
);
}
if let Some(decoded_len) = client_body_base64.and_then(decoded_base64_len_hint) {
metadata.insert(
"client_response_body_base64_bytes".to_string(),
Value::Number(decoded_len.into()),
);
}
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
fn plan_json_body_capture_for_usage(plan: &ExecutionPlan) -> Option<Value> {
if plan.body.body_ref.is_some() || plan.body.body_bytes_b64.is_some() {
return None;
@@ -1438,38 +1548,6 @@ fn decode_body_for_storage(body_base64: Option<&str>) -> Option<Value> {
Some(Value::String(body_base64.to_string()))
}
fn decoded_base64_len_hint(body_base64: &str) -> Option<u64> {
let body_base64 = body_base64.trim();
if body_base64.is_empty() {
return None;
}
let usable_len = body_base64.len();
if usable_len % 4 == 1 {
return None;
}
let padding = body_base64
.chars()
.rev()
.take_while(|char| *char == '=')
.count();
let full_quads = usable_len / 4;
let remainder = usable_len % 4;
let base_len = full_quads.saturating_mul(3);
let remainder_len = match remainder {
0 => 0,
2 => 1,
3 => 2,
_ => return None,
};
let decoded_len = base_len
.saturating_add(remainder_len)
.saturating_sub(padding.min(2));
Some(decoded_len as u64)
}
fn parse_sse_body_for_storage(text: &str) -> Option<Value> {
if !text.contains("data:") {
return None;
@@ -1641,14 +1719,15 @@ mod tests {
build_terminal_usage_event_from_seed, build_usage_event_data_seed,
extract_token_counts_from_json, headers_to_json, mask_header_value,
mask_sensitive_headers_in_json_value, LifecycleUsageSeed, TerminalUsageSeed,
UsageBodyRefsSeed, UsageRoutingSeed, UsageTerminalState, MAX_USAGE_CAPTURE_BYTES,
MAX_USAGE_CAPTURE_DEPTH,
UsageBodyRefsSeed, UsageBodyStatesSeed, UsageRoutingSeed, UsageTerminalState,
MAX_USAGE_CAPTURE_BYTES, MAX_USAGE_CAPTURE_DEPTH,
};
use crate::{
build_upsert_usage_record_from_event, GatewayStreamReportRequest, GatewaySyncReportRequest,
UsageEvent, UsageEventData, UsageEventType,
};
use aether_contracts::{ExecutionPlan, RequestBody};
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
use base64::Engine as _;
use serde_json::{json, Value};
use std::collections::BTreeMap;
@@ -1933,10 +2012,13 @@ mod tests {
.expect("provider body should encode"),
),
),
provider_body_state: Some(UsageBodyCaptureState::Inline),
client_body_base64: Some(
base64::engine::general_purpose::STANDARD
.encode("data: {\"id\":\"chatcmpl_123\"}\n\ndata: [DONE]\n"),
),
client_body_state: Some(UsageBodyCaptureState::Inline),
terminal_summary: None,
telemetry: None,
};
@@ -2022,7 +2104,10 @@ mod tests {
status_code: 200,
headers: BTreeMap::new(),
provider_body_base64: Some(base64::engine::general_purpose::STANDARD.encode(sse_body)),
provider_body_state: Some(UsageBodyCaptureState::Inline),
client_body_base64: None,
client_body_state: Some(UsageBodyCaptureState::None),
terminal_summary: None,
telemetry: None,
};
@@ -2350,9 +2435,12 @@ mod tests {
provider_body_base64: Some(
base64::engine::general_purpose::STANDARD.encode(provider_bytes),
),
provider_body_state: Some(UsageBodyCaptureState::Inline),
client_body_base64: Some(
base64::engine::general_purpose::STANDARD.encode(client_bytes),
),
client_body_state: Some(UsageBodyCaptureState::Inline),
terminal_summary: None,
telemetry: None,
};
@@ -2430,7 +2518,10 @@ mod tests {
status_code: 200,
headers: BTreeMap::new(),
provider_body_base64: Some(base64::engine::general_purpose::STANDARD.encode(&sse_body)),
provider_body_state: Some(UsageBodyCaptureState::Truncated),
client_body_base64: None,
client_body_state: Some(UsageBodyCaptureState::None),
terminal_summary: None,
telemetry: None,
};
@@ -2599,6 +2690,7 @@ mod tests {
has_format_conversion: false,
is_stream: false,
body_refs: UsageBodyRefsSeed::default(),
body_states: UsageBodyStatesSeed::default(),
routing: UsageRoutingSeed {
candidate_id: Some("cand-1".to_string()),
..UsageRoutingSeed::default()
@@ -2639,6 +2731,7 @@ mod tests {
})),
audit_payload: None,
standardized_usage: None,
terminal_summary: None,
})
.expect("usage event should build");
@@ -2726,6 +2819,7 @@ mod tests {
provider_endpoint_kind: Some("chat".to_string()),
has_format_conversion: Some(false),
is_stream: false,
body_states: UsageBodyStatesSeed::default(),
routing: UsageRoutingSeed {
candidate_id: Some("cand-1".to_string()),
..UsageRoutingSeed::default()