i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 09:20:22 +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

218
crates/aether-admin/src/observability/usage.rs
View File

@@ -265,84 +265,7 @@ fn admin_usage_has_body_value(
body: Option<&Value>,
field: UsageBodyField,
) -> bool {
body.is_some() || item.body_ref(field).is_some()
}
fn admin_usage_body_capture_storage(
item: &StoredRequestUsageAudit,
body: Option<&Value>,
field: UsageBodyField,
) -> &'static str {
if item.body_ref(field).is_some() {
"reference"
} else if body.is_some() {
"inline"
} else {
"missing"
}
}
fn admin_usage_body_capture_entry(
item: &StoredRequestUsageAudit,
body: Option<&Value>,
field: UsageBodyField,
) -> serde_json::Map<String, Value> {
let mut entry = serde_json::Map::new();
entry.insert(
"available".to_string(),
json!(admin_usage_has_body_value(item, body, field)),
);
entry.insert(
"storage".to_string(),
json!(admin_usage_body_capture_storage(item, body, field)),
);
if let Some(body_ref) = item.body_ref(field) {
entry.insert("body_ref".to_string(), json!(body_ref));
}
entry
}
fn admin_usage_request_preview_source(item: &StoredRequestUsageAudit) -> &'static str {
if item.body_ref(UsageBodyField::RequestBody).is_some() {
"stored_reference"
} else if item.request_body.is_some() {
"stored_original"
} else {
"local_reconstruction"
}
}
fn admin_usage_request_capture_entry(
item: &StoredRequestUsageAudit,
) -> serde_json::Map<String, Value> {
let mut request = admin_usage_body_capture_entry(
item,
item.request_body.as_ref(),
UsageBodyField::RequestBody,
);
request.insert(
"preview_source".to_string(),
json!(admin_usage_request_preview_source(item)),
);
request
}
fn admin_usage_curl_body_source(item: &StoredRequestUsageAudit) -> &'static str {
if admin_usage_has_body_value(
item,
item.provider_request_body.as_ref(),
UsageBodyField::ProviderRequestBody,
) {
"provider_request"
} else if admin_usage_has_body_value(
item,
item.request_body.as_ref(),
UsageBodyField::RequestBody,
) {
"request"
} else {
"local_reconstruction"
}
item.body_capture_result(field, body).available
}
fn admin_usage_strip_body_ref_metadata(metadata: &mut serde_json::Map<String, Value>) {
@@ -411,28 +334,12 @@ fn maybe_insert_bool_field(
}
fn admin_usage_body_capture_json(item: &StoredRequestUsageAudit) -> Value {
let request = admin_usage_request_capture_entry(item);
let provider_request = admin_usage_body_capture_entry(
item,
item.provider_request_body.as_ref(),
item.body_capture_json_for_fields(&[
UsageBodyField::RequestBody,
UsageBodyField::ProviderRequestBody,
);
let response = admin_usage_body_capture_entry(
item,
item.response_body.as_ref(),
UsageBodyField::ResponseBody,
);
let client_response = admin_usage_body_capture_entry(
item,
item.client_response_body.as_ref(),
UsageBodyField::ClientResponseBody,
);
json!({
"request": request,
"provider_request": provider_request,
"response": response,
"client_response": client_response,
})
])
}
fn admin_usage_settlement_json(item: &StoredRequestUsageAudit) -> Value {
@@ -511,21 +418,19 @@ fn admin_usage_trace_json(item: &StoredRequestUsageAudit) -> Value {
}
fn admin_usage_replay_body_capture_json(item: &StoredRequestUsageAudit) -> Value {
json!({
"request": admin_usage_request_capture_entry(item),
})
item.body_capture_json_for_fields(&[UsageBodyField::RequestBody])
}
fn admin_usage_curl_body_capture_json(item: &StoredRequestUsageAudit) -> Value {
json!({
"body_source": admin_usage_curl_body_source(item),
"request": admin_usage_request_capture_entry(item),
"provider_request": admin_usage_body_capture_entry(
item,
item.provider_request_body.as_ref(),
UsageBodyField::ProviderRequestBody,
),
})
let mut object = item.body_capture_json_object_for_fields(&[
UsageBodyField::RequestBody,
UsageBodyField::ProviderRequestBody,
]);
object.insert(
"body_source".to_string(),
Value::String(item.curl_body_source().to_string()),
);
Value::Object(object)
}
pub fn admin_usage_provider_key_name(
@@ -1339,19 +1244,12 @@ fn admin_usage_resolve_replay_mode(same_provider: bool, same_endpoint: bool) ->
}
}
pub fn admin_usage_resolve_request_preview_body(
pub fn admin_usage_resolve_request_capture_body(
item: &StoredRequestUsageAudit,
body_override: Option<Value>,
) -> Value {
) -> Option<Value> {
let resolved_model = item.model.clone();
let mut request_body = body_override
.or_else(|| item.request_body.clone())
.unwrap_or_else(|| {
json!({
"model": resolved_model,
"stream": item.is_stream,
})
});
let mut request_body = body_override.or_else(|| item.request_body.clone())?;
if let Some(body) = request_body.as_object_mut() {
body.entry("model".to_string())
.or_insert_with(|| json!(resolved_model));
@@ -1383,7 +1281,7 @@ pub fn admin_usage_resolve_request_preview_body(
.or_insert_with(|| json!(api_format));
}
}
request_body
Some(request_body)
}
pub fn admin_usage_headers_from_value(value: &Value) -> Option<BTreeMap<String, String>> {
@@ -1597,17 +1495,17 @@ pub fn build_admin_usage_curl_response(
url: Option<String>,
headers_json: Option<Value>,
headers: &BTreeMap<String, String>,
body: &Value,
body: Option<&Value>,
) -> Response<Body> {
let curl = admin_usage_build_curl_command(url.as_deref(), headers, Some(body));
let curl = admin_usage_build_curl_command(url.as_deref(), headers, body);
Json(json!({
"url": url,
"method": "POST",
"headers": headers_json.unwrap_or_else(|| json!(headers.clone())),
"body": body,
"body": body.cloned().unwrap_or(Value::Null),
"curl": curl,
"body_capture": admin_usage_curl_body_capture_json(item),
"original_request_body_available": admin_usage_has_body_value(
"captured_request_body_available": admin_usage_has_body_value(
item,
item.request_body.as_ref(),
UsageBodyField::RequestBody
@@ -1630,7 +1528,7 @@ pub fn build_admin_usage_detail_payload(
auth_api_key_reader_available: bool,
provider_key_name: Option<&str>,
include_bodies: bool,
request_body: Value,
request_body: Option<Value>,
default_headers: &BTreeMap<String, String>,
) -> Value {
let mut payload = admin_usage_record_json(
@@ -1651,9 +1549,6 @@ pub fn build_admin_usage_detail_payload(
admin_usage_strip_routing_metadata(object);
admin_usage_strip_settlement_metadata(object);
admin_usage_strip_trace_metadata(object);
object.remove("request_preview_source");
object.remove("original_request_body_available");
object.remove("original_response_body_available");
}
payload["user"] = match item.user_id.as_ref() {
Some(user_id) => json!({
@@ -1689,7 +1584,11 @@ pub fn build_admin_usage_detail_payload(
payload["body_capture"] = admin_usage_body_capture_json(item);
payload["settlement"] = admin_usage_settlement_json(item);
payload["trace"] = admin_usage_trace_json(item);
payload["has_request_body"] = json!(true);
payload["has_request_body"] = json!(admin_usage_has_body_value(
item,
item.request_body.as_ref(),
UsageBodyField::RequestBody
));
payload["has_provider_request_body"] = json!(admin_usage_has_body_value(
item,
item.provider_request_body.as_ref(),
@@ -1707,7 +1606,7 @@ pub fn build_admin_usage_detail_payload(
));
payload["tiered_pricing"] = Value::Null;
if include_bodies {
payload["request_body"] = request_body;
payload["request_body"] = request_body.unwrap_or(Value::Null);
payload["provider_request_body"] =
item.provider_request_body.clone().unwrap_or(Value::Null);
payload["response_body"] = item.response_body.clone().unwrap_or(Value::Null);
@@ -1727,7 +1626,7 @@ pub fn build_admin_usage_replay_plan_response(
target_provider: &StoredProviderCatalogProvider,
target_endpoint: &StoredProviderCatalogEndpoint,
target_api_key_id: Option<String>,
request_body: Value,
request_body: Option<Value>,
url: &str,
headers: &BTreeMap<String, String>,
same_provider: bool,
@@ -1735,7 +1634,7 @@ pub fn build_admin_usage_replay_plan_response(
) -> Response<Body> {
let resolved_model = item.model.clone();
let mapping_source = "none";
let curl = admin_usage_build_curl_command(Some(url), headers, Some(&request_body));
let curl = admin_usage_build_curl_command(Some(url), headers, request_body.as_ref());
Json(json!({
"dry_run": true,
@@ -1752,13 +1651,14 @@ pub fn build_admin_usage_replay_plan_response(
"method": "POST",
"url": url,
"request_headers": headers,
"request_body": request_body,
"request_body": request_body.clone().unwrap_or(Value::Null),
"body_capture": admin_usage_replay_body_capture_json(item),
"original_request_body_available": admin_usage_has_body_value(
"captured_request_body_available": admin_usage_has_body_value(
item,
item.request_body.as_ref(),
UsageBodyField::RequestBody
),
"request_body_available": request_body.is_some(),
"note": "Rust local replay currently exposes a dry-run plan and does not dispatch upstream",
"curl": curl,
}))
@@ -1911,13 +1811,10 @@ mod tests {
false,
None,
false,
json!({"model": "gpt-5"}),
Some(json!({"model": "gpt-5"})),
&BTreeMap::new(),
);
assert!(payload["metadata"]["request_preview_source"].is_null());
assert!(payload["metadata"]["original_request_body_available"].is_null());
assert!(payload["metadata"]["original_response_body_available"].is_null());
assert!(payload["metadata"]["request_body_ref"].is_null());
assert!(payload["metadata"]["provider_request_body_ref"].is_null());
assert!(payload["metadata"]["response_body_ref"].is_null());
@@ -1929,7 +1826,7 @@ mod tests {
"usage://request/req-1/request_body"
);
assert_eq!(
payload["body_capture"]["request"]["preview_source"],
payload["body_capture"]["request"]["capture_source"],
"stored_reference"
);
assert_eq!(
@@ -1954,6 +1851,42 @@ mod tests {
assert_eq!(payload["has_client_response_body"], true);
}
#[test]
fn detail_payload_preserves_legacy_body_capture_metadata_keys() {
let item = StoredRequestUsageAudit {
request_metadata: Some(json!({
"request_preview_source": "stored_original",
"original_request_body_available": true,
"original_response_body_available": true,
"custom_key": "keep-me"
})),
..sample_usage("completed", Some(200), None)
};
let payload = build_admin_usage_detail_payload(
&item,
&BTreeMap::new(),
&BTreeMap::new(),
false,
false,
None,
false,
Some(json!({"model": "gpt-5"})),
&BTreeMap::new(),
);
assert_eq!(
payload["metadata"]["request_preview_source"],
"stored_original"
);
assert_eq!(payload["metadata"]["original_request_body_available"], true);
assert_eq!(
payload["metadata"]["original_response_body_available"],
true
);
assert_eq!(payload["metadata"]["custom_key"], "keep-me");
}
#[test]
fn admin_usage_record_prefers_typed_pricing_metadata_with_billing_snapshot_fallback() {
let item = StoredRequestUsageAudit {
@@ -2022,7 +1955,7 @@ mod tests {
false,
None,
false,
json!({"model": "gpt-5"}),
Some(json!({"model": "gpt-5"})),
&BTreeMap::new(),
);
@@ -2096,7 +2029,7 @@ mod tests {
false,
Some("resolved-primary"),
false,
json!({"model": "gpt-5"}),
Some(json!({"model": "gpt-5"})),
&BTreeMap::new(),
);
@@ -2119,9 +2052,6 @@ mod tests {
assert!(payload["metadata"]["key_name"].is_null());
assert!(payload["metadata"]["planner_kind"].is_null());
assert!(payload["metadata"]["trace_id"].is_null());
assert!(payload["metadata"]["request_preview_source"].is_null());
assert!(payload["metadata"]["original_request_body_available"].is_null());
assert!(payload["metadata"]["original_response_body_available"].is_null());
assert!(payload["metadata"]["route_family"].is_null());
assert!(payload["metadata"]["route_kind"].is_null());
assert!(payload["metadata"]["execution_path"].is_null());

2
crates/aether-ai-pipeline/src/api.rs
View File

@@ -91,7 +91,7 @@ pub use crate::finalize::standard::openai::stream::{
};
pub use crate::finalize::standard::stream_core::common::*;
pub use crate::finalize::standard::stream_core::{
CanonicalStreamFrame, StreamingStandardFormatMatrix,
CanonicalStreamFrame, StreamingStandardFormatMatrix, StreamingStandardTerminalObserver,
};
pub use crate::finalize::sync_products::{
aggregate_claude_stream_sync_response, aggregate_gemini_stream_sync_response,

142
crates/aether-ai-pipeline/src/finalize/standard/stream_core/format_matrix.rs
View File

@@ -1,3 +1,4 @@
use aether_contracts::{ExecutionStreamTerminalSummary, StandardizedUsage};
use serde_json::Value;
use crate::conversion::{build_core_error_body_for_client_format, LocalCoreSyncErrorKind};
@@ -8,7 +9,9 @@ use crate::finalize::standard::openai::stream::{
OpenAIChatClientEmitter, OpenAIChatProviderState, OpenAICliClientEmitter,
OpenAICliProviderState,
};
use crate::finalize::standard::stream_core::common::{decode_json_data_line, CanonicalStreamFrame};
use crate::finalize::standard::stream_core::common::{
decode_json_data_line, CanonicalStreamEvent, CanonicalStreamFrame, CanonicalUsage,
};
use crate::finalize::PipelineFinalizeError;
#[derive(Default)]
@@ -60,18 +63,8 @@ impl StreamingStandardFormatMatrix {
return;
}
let provider_api_format = report_context
.get("provider_api_format")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
let client_api_format = report_context
.get("client_api_format")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase();
let provider_api_format = provider_api_format_for_context(report_context);
let client_api_format = client_api_format_for_context(report_context);
self.provider = ProviderStreamParser::for_api_format(provider_api_format.as_str());
self.client = ClientStreamEmitter::for_api_format(client_api_format.as_str());
@@ -99,6 +92,100 @@ impl StreamingStandardFormatMatrix {
}
}
#[derive(Default)]
pub struct StreamingStandardTerminalObserver {
provider: Option<ProviderStreamParser>,
latest_summary: Option<ExecutionStreamTerminalSummary>,
}
impl StreamingStandardTerminalObserver {
pub fn push_line(
&mut self,
report_context: &Value,
line: Vec<u8>,
) -> Result<(), PipelineFinalizeError> {
self.ensure_initialized(report_context);
let Some(provider) = self.provider.as_mut() else {
return Ok(());
};
let frames = provider.push_line(report_context, line)?;
self.observe_frames(frames);
Ok(())
}
pub fn finish(
&mut self,
report_context: &Value,
) -> Result<Option<ExecutionStreamTerminalSummary>, PipelineFinalizeError> {
self.ensure_initialized(report_context);
let Some(provider) = self.provider.as_mut() else {
return Ok(self.latest_summary.clone());
};
let frames = provider.finish(report_context)?;
self.observe_frames(frames);
Ok(self.latest_summary.clone())
}
pub fn disable_with_error(&mut self, parser_error: impl Into<String>) {
let parser_error = parser_error.into();
if let Some(summary) = self.latest_summary.as_mut() {
if summary.parser_error.is_none() {
summary.parser_error = Some(parser_error);
}
} else {
self.latest_summary = Some(ExecutionStreamTerminalSummary {
parser_error: Some(parser_error),
..ExecutionStreamTerminalSummary::default()
});
}
self.provider = None;
}
pub fn latest_summary(&self) -> Option<&ExecutionStreamTerminalSummary> {
self.latest_summary.as_ref()
}
fn ensure_initialized(&mut self, report_context: &Value) {
if self.provider.is_some() || self.latest_summary.is_some() {
return;
}
let provider_api_format = provider_api_format_for_context(report_context);
self.provider = ProviderStreamParser::for_api_format(provider_api_format.as_str());
}
fn observe_frames(&mut self, frames: Vec<CanonicalStreamFrame>) {
for frame in frames {
self.observe_frame(frame);
}
}
fn observe_frame(&mut self, frame: CanonicalStreamFrame) {
let CanonicalStreamFrame { id, model, event } = frame;
let summary = self
.latest_summary
.get_or_insert_with(|| ExecutionStreamTerminalSummary {
response_id: Some(id.clone()),
model: Some(model.clone()),
..ExecutionStreamTerminalSummary::default()
});
if summary.response_id.is_none() {
summary.response_id = Some(id);
}
if summary.model.is_none() {
summary.model = Some(model);
}
if let CanonicalStreamEvent::Finish {
finish_reason,
usage,
} = event
{
summary.finish_reason = finish_reason;
summary.standardized_usage = usage.map(standardized_usage_from_canonical);
summary.observed_finish = true;
}
}
}
enum ProviderStreamParser {
OpenAIChat(OpenAIChatProviderState),
OpenAICli(OpenAICliProviderState),
@@ -150,6 +237,35 @@ enum ClientStreamEmitter {
Gemini(GeminiClientEmitter),
}
fn provider_api_format_for_context(report_context: &Value) -> String {
report_context
.get("provider_api_format")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase()
}
fn client_api_format_for_context(report_context: &Value) -> String {
report_context
.get("client_api_format")
.and_then(Value::as_str)
.unwrap_or_default()
.trim()
.to_ascii_lowercase()
}
fn standardized_usage_from_canonical(usage: CanonicalUsage) -> StandardizedUsage {
let mut standardized = StandardizedUsage::new();
standardized.input_tokens = usage.input_tokens as i64;
standardized.output_tokens = usage.output_tokens as i64;
standardized.dimensions.insert(
"total_tokens".to_string(),
serde_json::json!(usage.total_tokens),
);
standardized
}
impl ClientStreamEmitter {
fn for_api_format(client_api_format: &str) -> Option<Self> {
Some(match client_api_format {

2
crates/aether-ai-pipeline/src/finalize/standard/stream_core/mod.rs
View File

@@ -2,4 +2,4 @@ pub mod common;
pub mod format_matrix;
pub use common::{CanonicalStreamEvent, CanonicalStreamFrame};
pub use format_matrix::StreamingStandardFormatMatrix;
pub use format_matrix::{StreamingStandardFormatMatrix, StreamingStandardTerminalObserver};

11
crates/aether-contracts/src/frame.rs
View File

@@ -1,9 +1,8 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use crate::{ExecutionError, ExecutionTelemetry};
use crate::{ExecutionError, ExecutionStreamTerminalSummary, ExecutionTelemetry};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
@@ -37,7 +36,7 @@ pub enum StreamFramePayload {
},
Eof {
#[serde(default, skip_serializing_if = "Option::is_none")]
summary: Option<Value>,
summary: Option<ExecutionStreamTerminalSummary>,
},
}
@@ -50,9 +49,13 @@ pub struct StreamFrame {
impl StreamFrame {
pub fn eof() -> Self {
Self::eof_with_summary(None)
}
pub fn eof_with_summary(summary: Option<ExecutionStreamTerminalSummary>) -> Self {
Self {
frame_type: StreamFrameType::Eof,
payload: StreamFramePayload::Eof { summary: None },
payload: StreamFramePayload::Eof { summary },
}
}
}

2
crates/aether-contracts/src/lib.rs
View File

@@ -3,6 +3,7 @@ mod frame;
mod plan;
mod result;
pub mod tunnel;
mod usage;
pub use error::{ExecutionError, ExecutionErrorKind, ExecutionPhase};
pub use frame::{StreamFrame, StreamFramePayload, StreamFrameType};
@@ -11,3 +12,4 @@ pub use plan::{
EXECUTION_REQUEST_FOLLOW_REDIRECTS_HEADER, EXECUTION_REQUEST_HTTP1_ONLY_HEADER,
};
pub use result::{ExecutionResult, ExecutionTelemetry, ResponseBody};
pub use usage::{ExecutionStreamTerminalSummary, StandardizedUsage};

131
crates/aether-contracts/src/usage.rs Normal file
View File

@@ -0,0 +1,131 @@
use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, PartialEq, Serialize, 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
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub struct ExecutionStreamTerminalSummary {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub standardized_usage: Option<StandardizedUsage>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub finish_reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model: Option<String>,
#[serde(default)]
pub observed_finish: bool,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub parser_error: Option<String>,
}
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"))
);
}
}

3
crates/aether-data-contracts/src/repository/usage/mod.rs
View File

@@ -11,7 +11,8 @@ pub use types::{
StoredUsagePerformancePercentilesRow, StoredUsageSettledCostSummary,
StoredUsageTimeSeriesBucket, UpsertUsageRecord, UsageAuditAggregationGroupBy,
UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditListQuery,
UsageAuditSummaryQuery, UsageBodyField, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery,
UsageAuditSummaryQuery, UsageBodyCaptureResult, UsageBodyCaptureState, UsageBodyCaptureStorage,
UsageBodyField, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery,
UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy,
UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCostSavingsSummaryQuery,
UsageDailyHeatmapQuery, UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery,

383
crates/aether-data-contracts/src/repository/usage/types.rs
View File

@@ -1,12 +1,21 @@
use async_trait::async_trait;
use serde_json::Value;
/// Joined usage read model assembled from the accounting row plus the newer audit/snapshot
/// satellite tables.
///
/// The canonical owners are split across `public.usage`, `public.usage_http_audits`,
/// `public.usage_body_blobs`, `public.usage_routing_snapshots`, and
/// `public.usage_settlement_snapshots`. Fallback reads from deprecated `public.usage.*` mirror
/// columns still exist for historical rows, but those columns are compatibility-only.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredRequestUsageAudit {
pub id: String,
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
// Legacy display-cache mirrors from `public.usage`. New writes intentionally avoid treating
// them as authoritative fields.
pub username: Option<String>,
pub api_key_name: Option<String>,
pub provider_name: String,
@@ -42,7 +51,12 @@ pub struct StoredRequestUsageAudit {
pub response_time_ms: Option<u64>,
pub first_byte_time_ms: Option<u64>,
pub status: String,
// Settlement state prefers `public.usage_settlement_snapshots` and only falls back to the
// deprecated `public.usage.billing_status` mirror for older rows.
pub billing_status: String,
// HTTP capture read model. Canonical owners are `public.usage_http_audits` plus
// `public.usage_body_blobs`; deprecated `public.usage.*headers/*body*` columns are historical
// fallback only.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
@@ -50,24 +64,32 @@ pub struct StoredRequestUsageAudit {
#[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>,
@@ -210,15 +232,19 @@ impl StoredRequestUsageAudit {
request_headers: None,
request_body: None,
request_body_ref: None,
request_body_state: None,
provider_request_headers: None,
provider_request_body: None,
provider_request_body_ref: None,
provider_request_body_state: None,
response_headers: None,
response_body: None,
response_body_ref: None,
response_body_state: None,
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -350,6 +376,108 @@ impl StoredRequestUsageAudit {
}
}
pub fn body_value(&self, field: UsageBodyField) -> Option<&Value> {
match field {
UsageBodyField::RequestBody => self.request_body.as_ref(),
UsageBodyField::ProviderRequestBody => self.provider_request_body.as_ref(),
UsageBodyField::ResponseBody => self.response_body.as_ref(),
UsageBodyField::ClientResponseBody => self.client_response_body.as_ref(),
}
}
pub fn body_state(&self, field: UsageBodyField) -> Option<UsageBodyCaptureState> {
match field {
UsageBodyField::RequestBody => self.request_body_state,
UsageBodyField::ProviderRequestBody => self.provider_request_body_state,
UsageBodyField::ResponseBody => self.response_body_state,
UsageBodyField::ClientResponseBody => self.client_response_body_state,
}
}
pub fn body_capture_result(
&self,
field: UsageBodyField,
body: Option<&Value>,
) -> UsageBodyCaptureResult {
resolve_usage_body_capture_result(
self.body_state(field),
body.is_some(),
self.body_ref(field).is_some(),
)
}
pub fn body_capture_json_entry(
&self,
field: UsageBodyField,
body: Option<&Value>,
) -> serde_json::Map<String, Value> {
self.body_capture_result(field, body)
.as_json_entry(self.body_ref(field))
}
pub fn request_body_capture_json_entry(&self) -> serde_json::Map<String, Value> {
let mut entry =
self.body_capture_json_entry(UsageBodyField::RequestBody, self.request_body.as_ref());
entry.insert(
"capture_source".to_string(),
Value::String(
self.body_capture_result(UsageBodyField::RequestBody, self.request_body.as_ref())
.request_capture_source()
.to_string(),
),
);
entry
}
pub fn body_capture_json_object_for_fields(
&self,
fields: &[UsageBodyField],
) -> serde_json::Map<String, Value> {
let mut object = serde_json::Map::new();
for field in fields {
let entry = match field {
UsageBodyField::RequestBody => self.request_body_capture_json_entry(),
other => self.body_capture_json_entry(*other, self.body_value(*other)),
};
object.insert(field.as_capture_key().to_string(), Value::Object(entry));
}
object
}
pub fn body_capture_json_for_fields(&self, fields: &[UsageBodyField]) -> Value {
Value::Object(self.body_capture_json_object_for_fields(fields))
}
pub fn preferred_request_body_source_field(&self) -> Option<UsageBodyField> {
if self
.body_capture_result(
UsageBodyField::ProviderRequestBody,
self.body_value(UsageBodyField::ProviderRequestBody),
)
.available
{
Some(UsageBodyField::ProviderRequestBody)
} else if self
.body_capture_result(
UsageBodyField::RequestBody,
self.body_value(UsageBodyField::RequestBody),
)
.available
{
Some(UsageBodyField::RequestBody)
} else {
None
}
}
pub fn curl_body_source(&self) -> &'static str {
match self.preferred_request_body_source_field() {
Some(UsageBodyField::ProviderRequestBody) => "provider_request",
Some(UsageBodyField::RequestBody) => "request",
_ => "unavailable",
}
}
pub fn routing_candidate_id(&self) -> Option<&str> {
self.candidate_id
.as_deref()
@@ -895,6 +1023,143 @@ pub struct StoredUsageDailySummary {
pub actual_total_cost_usd: f64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UsageBodyCaptureState {
None,
Inline,
Reference,
Truncated,
Disabled,
Unavailable,
}
impl UsageBodyCaptureState {
pub fn as_str(self) -> &'static str {
match self {
Self::None => "none",
Self::Inline => "inline",
Self::Reference => "reference",
Self::Truncated => "truncated",
Self::Disabled => "disabled",
Self::Unavailable => "unavailable",
}
}
pub fn from_capture_parts(
has_inline_body: bool,
has_reference: bool,
unavailable: bool,
) -> Self {
if has_inline_body {
Self::Inline
} else if has_reference {
Self::Reference
} else if unavailable {
Self::Unavailable
} else {
Self::None
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum UsageBodyCaptureStorage {
Inline,
Reference,
Truncated,
Disabled,
Unavailable,
None,
Missing,
}
impl UsageBodyCaptureStorage {
pub fn as_str(self) -> &'static str {
match self {
Self::Inline => "inline",
Self::Reference => "reference",
Self::Truncated => "truncated",
Self::Disabled => "disabled",
Self::Unavailable => "unavailable",
Self::None => "none",
Self::Missing => "missing",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct UsageBodyCaptureResult {
pub available: bool,
pub storage: UsageBodyCaptureStorage,
pub state: Option<UsageBodyCaptureState>,
}
impl UsageBodyCaptureResult {
pub fn state_label(self) -> &'static str {
self.state
.map(UsageBodyCaptureState::as_str)
.unwrap_or("legacy_unknown")
}
pub fn as_json_entry(self, body_ref: Option<&str>) -> serde_json::Map<String, Value> {
let mut entry = serde_json::Map::new();
entry.insert("available".to_string(), Value::Bool(self.available));
entry.insert(
"storage".to_string(),
Value::String(self.storage.as_str().to_string()),
);
entry.insert(
"state".to_string(),
Value::String(self.state_label().to_string()),
);
if let Some(body_ref) = body_ref {
entry.insert("body_ref".to_string(), Value::String(body_ref.to_string()));
}
entry
}
pub fn request_capture_source(self) -> &'static str {
match self.state {
Some(UsageBodyCaptureState::Reference) => "stored_reference",
Some(UsageBodyCaptureState::Inline) => "stored_original",
Some(UsageBodyCaptureState::Truncated) => "stored_truncated",
Some(UsageBodyCaptureState::Disabled) => "disabled",
Some(UsageBodyCaptureState::Unavailable) => "unavailable",
Some(UsageBodyCaptureState::None) => "not_captured",
None => match self.storage {
UsageBodyCaptureStorage::Reference => "stored_reference",
UsageBodyCaptureStorage::Inline => "stored_original",
_ => "legacy_unknown",
},
}
}
}
pub fn resolve_usage_body_capture_result(
state: Option<UsageBodyCaptureState>,
has_inline_body: bool,
has_reference: bool,
) -> UsageBodyCaptureResult {
let (available, storage) = match state {
Some(UsageBodyCaptureState::Inline) => (true, UsageBodyCaptureStorage::Inline),
Some(UsageBodyCaptureState::Reference) => (true, UsageBodyCaptureStorage::Reference),
Some(UsageBodyCaptureState::Truncated) => (true, UsageBodyCaptureStorage::Truncated),
Some(UsageBodyCaptureState::Disabled) => (false, UsageBodyCaptureStorage::Disabled),
Some(UsageBodyCaptureState::Unavailable) => (false, UsageBodyCaptureStorage::Unavailable),
Some(UsageBodyCaptureState::None) => (false, UsageBodyCaptureStorage::None),
None if has_reference => (true, UsageBodyCaptureStorage::Reference),
None if has_inline_body => (true, UsageBodyCaptureStorage::Inline),
None => (false, UsageBodyCaptureStorage::Missing),
};
UsageBodyCaptureResult {
available,
storage,
state,
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum UsageBodyField {
@@ -905,6 +1170,15 @@ pub enum UsageBodyField {
}
impl UsageBodyField {
pub fn as_capture_key(&self) -> &'static str {
match self {
Self::RequestBody => "request",
Self::ProviderRequestBody => "provider_request",
Self::ResponseBody => "response",
Self::ClientResponseBody => "client_response",
}
}
pub fn as_ref_key(&self) -> &'static str {
match self {
Self::RequestBody => "request_body_ref",
@@ -1111,11 +1385,18 @@ pub trait UsageReadRepository: Send + Sync {
) -> Result<Vec<StoredUsageDailySummary>, crate::DataLayerError>;
}
/// Repository write model for a single usage aggregate.
///
/// Request/response headers and bodies here are capture inputs that the repository persists into
/// the dedicated HTTP audit/body stores. Deprecated mirror columns on `public.usage` remain in the
/// schema for compatibility only and are not the intended long-term destination for new writes.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct UpsertUsageRecord {
pub request_id: String,
pub user_id: Option<String>,
pub api_key_id: Option<String>,
// Legacy display-cache mirrors on `public.usage`. Repository write paths strip these so new
// values do not keep populating the deprecated columns.
pub username: Option<String>,
pub api_key_name: Option<String>,
pub provider_name: String,
@@ -1151,19 +1432,27 @@ pub struct UpsertUsageRecord {
pub response_time_ms: Option<u64>,
pub first_byte_time_ms: Option<u64>,
pub status: String,
// Settlement state is also projected into `public.usage_settlement_snapshots`; the base usage
// row mirror exists for compatibility and indexing until a later schema cleanup.
pub billing_status: String,
// HTTP capture payload. Canonical persistence goes through `usage_http_audits` and
// `usage_body_blobs`; any remaining `public.usage` body/header columns are compatibility-only.
pub request_headers: Option<Value>,
pub request_body: Option<Value>,
pub request_body_ref: Option<String>,
pub request_body_state: Option<UsageBodyCaptureState>,
pub provider_request_headers: Option<Value>,
pub provider_request_body: Option<Value>,
pub provider_request_body_ref: Option<String>,
pub provider_request_body_state: Option<UsageBodyCaptureState>,
pub response_headers: Option<Value>,
pub response_body: Option<Value>,
pub response_body_ref: Option<String>,
pub response_body_state: Option<UsageBodyCaptureState>,
pub client_response_headers: Option<Value>,
pub client_response_body: Option<Value>,
pub client_response_body_ref: Option<String>,
pub client_response_body_state: Option<UsageBodyCaptureState>,
pub candidate_id: Option<String>,
pub candidate_index: Option<u64>,
pub key_name: Option<String>,
@@ -1293,8 +1582,11 @@ fn parse_timestamp(value: i64, field_name: &str) -> Result<u64, crate::DataLayer
#[cfg(test)]
mod tests {
use super::{StoredRequestUsageAudit, UpsertUsageRecord, UsageBodyField};
use serde_json::json;
use super::{
StoredRequestUsageAudit, UpsertUsageRecord, UsageBodyCaptureState, UsageBodyCaptureStorage,
UsageBodyField,
};
use serde_json::{json, Value};
fn sample_usage() -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
@@ -1478,6 +1770,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -1623,4 +1919,87 @@ mod tests {
assert_eq!(usage.body_ref(UsageBodyField::ResponseBody), None);
assert_eq!(usage.body_ref(UsageBodyField::ClientResponseBody), None);
}
#[test]
fn body_capture_result_prefers_typed_state_over_inline_body_presence() {
let mut usage = sample_usage();
usage.request_body = Some(json!({"model": "gpt-5"}));
usage.request_body_ref = Some("usage://request/req-1/request_body".to_string());
usage.request_body_state = Some(UsageBodyCaptureState::Disabled);
let result =
usage.body_capture_result(UsageBodyField::RequestBody, usage.request_body.as_ref());
assert!(!result.available);
assert_eq!(result.storage, UsageBodyCaptureStorage::Disabled);
assert_eq!(result.state_label(), "disabled");
assert_eq!(result.request_capture_source(), "disabled");
}
#[test]
fn body_capture_result_infers_legacy_inline_storage_when_typed_state_is_missing() {
let mut usage = sample_usage();
usage.request_body = Some(json!({"model": "gpt-5"}));
let result =
usage.body_capture_result(UsageBodyField::RequestBody, usage.request_body.as_ref());
assert!(result.available);
assert_eq!(result.storage, UsageBodyCaptureStorage::Inline);
assert_eq!(result.state_label(), "legacy_unknown");
assert_eq!(result.request_capture_source(), "stored_original");
}
#[test]
fn request_body_capture_json_entry_includes_capture_source_and_body_ref() {
let mut usage = sample_usage();
usage.request_body_ref = Some("usage://request/req-1/request_body".to_string());
usage.request_body_state = Some(UsageBodyCaptureState::Reference);
let entry = usage.request_body_capture_json_entry();
assert_eq!(entry.get("available"), Some(&Value::Bool(true)));
assert_eq!(
entry.get("storage"),
Some(&Value::String("reference".to_string()))
);
assert_eq!(
entry.get("state"),
Some(&Value::String("reference".to_string()))
);
assert_eq!(
entry.get("body_ref"),
Some(&Value::String(
"usage://request/req-1/request_body".to_string()
))
);
assert_eq!(
entry.get("capture_source"),
Some(&Value::String("stored_reference".to_string()))
);
}
#[test]
fn curl_body_source_prefers_provider_request_body_over_request_body() {
let mut usage = sample_usage();
usage.request_body = Some(json!({"client": true}));
usage.provider_request_body = Some(json!({"provider": true}));
assert_eq!(
usage.preferred_request_body_source_field(),
Some(UsageBodyField::ProviderRequestBody)
);
assert_eq!(usage.curl_body_source(), "provider_request");
usage.provider_request_body = None;
assert_eq!(
usage.preferred_request_body_source_field(),
Some(UsageBodyField::RequestBody)
);
assert_eq!(usage.curl_body_source(), "request");
usage.request_body = None;
assert_eq!(usage.preferred_request_body_source_field(), None);
assert_eq!(usage.curl_body_source(), "unavailable");
}
}

4
crates/aether-data/bootstrap/20260413020000_baseline_v2.sql
View File

@@ -4623,6 +4623,10 @@ CREATE TABLE IF NOT EXISTS public.usage_http_audits (
provider_request_body_ref character varying(160),
response_body_ref character varying(160),
client_response_body_ref character varying(160),
request_body_state character varying(32),
provider_request_body_state character varying(32),
response_body_state character varying(32),
client_response_body_state character varying(32),
body_capture_mode character varying(32) DEFAULT 'none' NOT NULL,
created_at timestamp with time zone DEFAULT now() NOT NULL,
updated_at timestamp with time zone DEFAULT now() NOT NULL,

4
crates/aether-data/migrations/20260413020000_squash_usage_schema_split.sql
View File

@@ -35,6 +35,10 @@ CREATE TABLE IF NOT EXISTS public.usage_http_audits (
provider_request_body_ref character varying(160),
response_body_ref character varying(160),
client_response_body_ref character varying(160),
request_body_state character varying(32),
provider_request_body_state character varying(32),
response_body_state character varying(32),
client_response_body_state character varying(32),
body_capture_mode character varying(32) DEFAULT 'none' NOT NULL,
created_at timestamp with time zone DEFAULT now() NOT NULL,
updated_at timestamp with time zone DEFAULT now() NOT NULL,

5
crates/aether-data/migrations/20260418000000_add_usage_http_audit_body_states.sql Normal file
View File

@@ -0,0 +1,5 @@
ALTER TABLE public.usage_http_audits
ADD COLUMN IF NOT EXISTS request_body_state character varying(32),
ADD COLUMN IF NOT EXISTS provider_request_body_state character varying(32),
ADD COLUMN IF NOT EXISTS response_body_state character varying(32),
ADD COLUMN IF NOT EXISTS client_response_body_state character varying(32);

10
crates/aether-data/src/migrate.rs
View File

@@ -8,7 +8,7 @@ use tracing::{error, info, warn};
static MIGRATOR: Migrator = sqlx::migrate!("./migrations");
static BASELINE_V2_SQL: &str = include_str!("../bootstrap/20260413020000_baseline_v2.sql");
const BASELINE_V2_CUTOFF_VERSION: i64 = 20260415000000;
const BASELINE_V2_CUTOFF_VERSION: i64 = 20260418000000;
const MIGRATIONS_TABLE_EXISTS_SQL: &str =
"SELECT to_regclass('public._sqlx_migrations') IS NOT NULL";
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
@@ -623,6 +623,7 @@ mod tests {
20260413020000,
20260413030000,
20260415000000,
20260418000000,
]
);
}
@@ -632,6 +633,10 @@ mod tests {
assert!(BASELINE_V2_SQL.contains("CREATE TABLE IF NOT EXISTS public.usage_body_blobs"));
assert!(BASELINE_V2_SQL.contains("ix_usage_body_blobs_request_id"));
assert!(BASELINE_V2_SQL.contains("CREATE TABLE IF NOT EXISTS public.usage_http_audits"));
assert!(BASELINE_V2_SQL.contains("request_body_state character varying(32)"));
assert!(BASELINE_V2_SQL.contains("provider_request_body_state character varying(32)"));
assert!(BASELINE_V2_SQL.contains("response_body_state character varying(32)"));
assert!(BASELINE_V2_SQL.contains("client_response_body_state character varying(32)"));
assert!(
BASELINE_V2_SQL.contains("CREATE TABLE IF NOT EXISTS public.usage_routing_snapshots")
);
@@ -717,7 +722,8 @@ mod tests {
20260410000000,
20260413020000,
20260413030000,
20260415000000
20260415000000,
20260418000000,
]
);
}

68
crates/aether-data/src/repository/usage/memory.rs
View File

@@ -2194,6 +2194,9 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
.request_body
.or_else(|| existing.and_then(|existing| existing.request_body.clone())),
request_body_ref,
request_body_state: usage
.request_body_state
.or_else(|| existing.and_then(|existing| existing.request_body_state)),
provider_request_headers: usage.provider_request_headers.or_else(|| {
existing.and_then(|existing| existing.provider_request_headers.clone())
}),
@@ -2201,6 +2204,9 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
.provider_request_body
.or_else(|| existing.and_then(|existing| existing.provider_request_body.clone())),
provider_request_body_ref,
provider_request_body_state: usage
.provider_request_body_state
.or_else(|| existing.and_then(|existing| existing.provider_request_body_state)),
response_headers: usage
.response_headers
.or_else(|| existing.and_then(|existing| existing.response_headers.clone())),
@@ -2208,6 +2214,9 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
.response_body
.or_else(|| existing.and_then(|existing| existing.response_body.clone())),
response_body_ref,
response_body_state: usage
.response_body_state
.or_else(|| existing.and_then(|existing| existing.response_body_state)),
client_response_headers: usage
.client_response_headers
.or_else(|| existing.and_then(|existing| existing.client_response_headers.clone())),
@@ -2215,6 +2224,9 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
.client_response_body
.or_else(|| existing.and_then(|existing| existing.client_response_body.clone())),
client_response_body_ref,
client_response_body_state: usage
.client_response_body_state
.or_else(|| existing.and_then(|existing| existing.client_response_body_state)),
candidate_id: usage.candidate_id.or_else(|| {
existing.and_then(|existing| existing.routing_candidate_id().map(ToOwned::to_owned))
}),
@@ -2382,6 +2394,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -2451,6 +2467,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -2534,6 +2554,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -2603,6 +2627,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: Some("cand-1".to_string()),
candidate_index: Some(1),
key_name: Some("primary".to_string()),
@@ -2694,6 +2722,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: Some("cand-1".to_string()),
candidate_index: Some(1),
key_name: Some("primary".to_string()),
@@ -2765,6 +2797,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -2965,6 +3001,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -3052,6 +3092,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -3127,6 +3171,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -3205,6 +3253,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -3288,6 +3340,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: Some("cand-typed".to_string()),
candidate_index: Some(2),
key_name: Some("primary".to_string()),
@@ -3382,6 +3438,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -3457,6 +3517,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -3523,6 +3587,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,

8
crates/aether-data/src/repository/usage/mod.rs
View File

@@ -45,6 +45,10 @@ pub(crate) fn usage_can_recover_terminal_failure(
&& incoming_usage_can_recover_terminal_failure(incoming_status, incoming_billing_status)
}
/// Clear legacy display-cache fields that still exist on `public.usage` for compatibility.
///
/// These values are no longer treated as authoritative read-model inputs; new writes should not
/// repopulate the deprecated mirror columns.
pub(crate) fn strip_deprecated_usage_display_fields(
mut usage: UpsertUsageRecord,
) -> UpsertUsageRecord {
@@ -114,6 +118,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,

171
crates/aether-data/src/repository/usage/sql.rs
View File

@@ -6,8 +6,8 @@ use aether_data_contracts::repository::usage::{
StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
StoredUsagePerformancePercentilesRow, StoredUsageSettledCostSummary,
StoredUsageTimeSeriesBucket, UsageAuditAggregationGroupBy, UsageAuditAggregationQuery,
UsageAuditKeywordSearchQuery, UsageAuditSummaryQuery, UsageBodyField, UsageBreakdownGroupBy,
UsageBreakdownSummaryQuery, UsageCacheAffinityHitSummaryQuery,
UsageAuditKeywordSearchQuery, UsageAuditSummaryQuery, UsageBodyCaptureState, UsageBodyField,
UsageBreakdownGroupBy, UsageBreakdownSummaryQuery, UsageCacheAffinityHitSummaryQuery,
UsageCacheAffinityIntervalGroupBy, UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery,
UsageCostSavingsSummaryQuery, UsageDashboardDailyBreakdownQuery,
UsageDashboardProviderCountsQuery, UsageDashboardSummaryQuery, UsageErrorDistributionQuery,
@@ -89,6 +89,10 @@ INSERT INTO usage_http_audits (
provider_request_body_ref,
response_body_ref,
client_response_body_ref,
request_body_state,
provider_request_body_state,
response_body_state,
client_response_body_state,
body_capture_mode
) VALUES (
$1,
@@ -100,7 +104,11 @@ INSERT INTO usage_http_audits (
$7,
$8,
$9,
$10
$10,
$11,
$12,
$13,
$14
)
ON CONFLICT (request_id)
DO UPDATE SET
@@ -124,6 +132,22 @@ DO UPDATE SET
EXCLUDED.client_response_body_ref,
usage_http_audits.client_response_body_ref
),
request_body_state = COALESCE(
EXCLUDED.request_body_state,
usage_http_audits.request_body_state
),
provider_request_body_state = COALESCE(
EXCLUDED.provider_request_body_state,
usage_http_audits.provider_request_body_state
),
response_body_state = COALESCE(
EXCLUDED.response_body_state,
usage_http_audits.response_body_state
),
client_response_body_state = COALESCE(
EXCLUDED.client_response_body_state,
usage_http_audits.client_response_body_state
),
body_capture_mode = COALESCE(
NULLIF(EXCLUDED.body_capture_mode, 'none'),
usage_http_audits.body_capture_mode,
@@ -330,6 +354,10 @@ SELECT
usage_http_audits.provider_request_body_ref AS http_provider_request_body_ref,
usage_http_audits.response_body_ref AS http_response_body_ref,
usage_http_audits.client_response_body_ref AS http_client_response_body_ref,
usage_http_audits.request_body_state AS http_request_body_state,
usage_http_audits.provider_request_body_state AS http_provider_request_body_state,
usage_http_audits.response_body_state AS http_response_body_state,
usage_http_audits.client_response_body_state AS http_client_response_body_state,
usage_routing_snapshots.candidate_id AS routing_candidate_id,
usage_routing_snapshots.candidate_index AS routing_candidate_index,
usage_routing_snapshots.key_name AS routing_key_name,
@@ -597,6 +625,10 @@ SELECT
NULL::varchar AS http_provider_request_body_ref,
NULL::varchar AS http_response_body_ref,
NULL::varchar AS http_client_response_body_ref,
NULL::varchar AS http_request_body_state,
NULL::varchar AS http_provider_request_body_state,
NULL::varchar AS http_response_body_state,
NULL::varchar AS http_client_response_body_state,
COALESCE(
usage_routing_snapshots.candidate_id,
NULLIF(BTRIM("usage".request_metadata->>'candidate_id'), '')
@@ -813,6 +845,10 @@ SELECT
NULL::varchar AS http_provider_request_body_ref,
NULL::varchar AS http_response_body_ref,
NULL::varchar AS http_client_response_body_ref,
NULL::varchar AS http_request_body_state,
NULL::varchar AS http_provider_request_body_state,
NULL::varchar AS http_response_body_state,
NULL::varchar AS http_client_response_body_state,
COALESCE(
usage_routing_snapshots.candidate_id,
NULLIF(BTRIM("usage".request_metadata->>'candidate_id'), '')
@@ -1148,6 +1184,10 @@ RETURNING
NULL::varchar AS http_provider_request_body_ref,
NULL::varchar AS http_response_body_ref,
NULL::varchar AS http_client_response_body_ref,
NULL::varchar AS http_request_body_state,
NULL::varchar AS http_provider_request_body_state,
NULL::varchar AS http_response_body_state,
NULL::varchar AS http_client_response_body_state,
NULL::varchar AS routing_candidate_id,
NULL::integer AS routing_candidate_index,
NULL::varchar AS routing_key_name,
@@ -3991,6 +4031,12 @@ WHERE "usage".created_at >= TO_TIMESTAMP($1::double precision)"#,
None,
),
};
let http_audit_states = UsageHttpAuditStates {
request_body_state: usage.request_body_state,
provider_request_body_state: usage.provider_request_body_state,
response_body_state: usage.response_body_state,
client_response_body_state: usage.client_response_body_state,
};
let request_metadata_value = prepare_request_metadata_for_body_storage(
usage.request_metadata.clone(),
[
@@ -4163,6 +4209,7 @@ WHERE "usage".created_at >= TO_TIMESTAMP($1::double precision)"#,
&usage.request_id,
&http_audit_headers,
&http_audit_refs,
&http_audit_states,
http_audit_capture_mode,
)
.await?;
@@ -4224,6 +4271,16 @@ WHERE "usage".created_at >= TO_TIMESTAMP($1::double precision)"#,
client_response_body_storage.has_detached_blob(),
http_audit_refs.client_response_body_ref.as_deref(),
);
stored.request_body_state =
usage.request_body_state.or(stored.request_body_state);
stored.provider_request_body_state = usage
.provider_request_body_state
.or(stored.provider_request_body_state);
stored.response_body_state =
usage.response_body_state.or(stored.response_body_state);
stored.client_response_body_state = usage
.client_response_body_state
.or(stored.client_response_body_state);
stored.candidate_id = routing_snapshot.candidate_id.clone();
stored.candidate_index = routing_snapshot.candidate_index;
stored.key_name = routing_snapshot.key_name.clone();
@@ -4462,6 +4519,10 @@ impl UsageWriteRepository for SqlxUsageReadRepository {
}
}
// Build the usage read model from the split storage layout.
//
// Query projections already prefer the newer audit/snapshot owners and only fall back to
// deprecated `public.usage` mirror columns for historical rows that predate the split schema.
fn map_usage_row(
row: &sqlx::postgres::PgRow,
resolve_compressed_bodies: bool,
@@ -4574,6 +4635,28 @@ fn map_usage_row(
.try_get("http_client_response_body_ref")
.map_postgres_err()?,
};
let http_audit_states = UsageHttpAuditStates {
request_body_state: row
.try_get::<Option<String>, _>("http_request_body_state")
.map_postgres_err()?
.as_deref()
.and_then(parse_usage_body_capture_state),
provider_request_body_state: row
.try_get::<Option<String>, _>("http_provider_request_body_state")
.map_postgres_err()?
.as_deref()
.and_then(parse_usage_body_capture_state),
response_body_state: row
.try_get::<Option<String>, _>("http_response_body_state")
.map_postgres_err()?
.as_deref()
.and_then(parse_usage_body_capture_state),
client_response_body_state: row
.try_get::<Option<String>, _>("http_client_response_body_state")
.map_postgres_err()?
.as_deref()
.and_then(parse_usage_body_capture_state),
};
let routing_snapshot = usage_routing_snapshot_from_row(row)?;
let settlement_pricing_snapshot = usage_settlement_pricing_snapshot_from_row(row)?;
usage.request_body = request_body.value;
@@ -4613,6 +4696,10 @@ fn map_usage_row(
client_response_body.has_compressed_storage,
http_audit_refs.client_response_body_ref.as_deref(),
);
usage.request_body_state = http_audit_states.request_body_state;
usage.provider_request_body_state = http_audit_states.provider_request_body_state;
usage.response_body_state = http_audit_states.response_body_state;
usage.client_response_body_state = http_audit_states.client_response_body_state;
usage.candidate_id = routing_snapshot.candidate_id.clone();
usage.candidate_index = routing_snapshot.candidate_index;
usage.key_name = routing_snapshot.key_name.clone();
@@ -4675,6 +4762,23 @@ impl UsageHttpAuditRefs {
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct UsageHttpAuditStates {
request_body_state: Option<UsageBodyCaptureState>,
provider_request_body_state: Option<UsageBodyCaptureState>,
response_body_state: Option<UsageBodyCaptureState>,
client_response_body_state: Option<UsageBodyCaptureState>,
}
impl UsageHttpAuditStates {
fn any_present(&self) -> bool {
self.request_body_state.is_some()
|| self.provider_request_body_state.is_some()
|| self.response_body_state.is_some()
|| self.client_response_body_state.is_some()
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
struct UsageHttpAuditHeaders<'a> {
request_headers_json: Option<&'a str>,
@@ -4802,6 +4906,24 @@ fn json_bind_text(value: Option<&Value>) -> Result<Option<String>, DataLayerErro
.transpose()
}
fn usage_body_capture_state_bind_text(
value: Option<UsageBodyCaptureState>,
) -> Option<&'static str> {
value.map(UsageBodyCaptureState::as_str)
}
fn parse_usage_body_capture_state(value: &str) -> Option<UsageBodyCaptureState> {
match value.trim() {
"none" => Some(UsageBodyCaptureState::None),
"inline" => Some(UsageBodyCaptureState::Inline),
"reference" => Some(UsageBodyCaptureState::Reference),
"truncated" => Some(UsageBodyCaptureState::Truncated),
"disabled" => Some(UsageBodyCaptureState::Disabled),
"unavailable" => Some(UsageBodyCaptureState::Unavailable),
_ => None,
}
}
#[cfg(test)]
fn usage_http_audit_body_refs(metadata: Option<&Value>) -> UsageHttpAuditRefs {
let object = metadata.and_then(Value::as_object);
@@ -5042,6 +5164,10 @@ fn usage_settlement_pricing_snapshot_from_usage(
}
}
// Decode deprecated inline/compressed body columns from `public.usage`.
//
// New writes keep these columns empty by forcing body storage through `usage_body_blobs` and
// `usage_http_audits`; this helper exists only so older rows remain readable without backfill.
fn usage_json_column(
row: &sqlx::postgres::PgRow,
inline_column: &str,
@@ -5286,12 +5412,17 @@ async fn sync_usage_http_audit_storage<'e, E>(
request_id: &str,
headers: &UsageHttpAuditHeaders<'_>,
refs: &UsageHttpAuditRefs,
states: &UsageHttpAuditStates,
body_capture_mode: &str,
) -> Result<(), DataLayerError>
where
E: sqlx::Executor<'e, Database = Postgres>,
{
if !headers.any_present() && !refs.any_present() && body_capture_mode == "none" {
if !headers.any_present()
&& !refs.any_present()
&& !states.any_present()
&& body_capture_mode == "none"
{
return Ok(());
}
@@ -5305,6 +5436,18 @@ where
.bind(refs.provider_request_body_ref.as_deref())
.bind(refs.response_body_ref.as_deref())
.bind(refs.client_response_body_ref.as_deref())
.bind(usage_body_capture_state_bind_text(
states.request_body_state,
))
.bind(usage_body_capture_state_bind_text(
states.provider_request_body_state,
))
.bind(usage_body_capture_state_bind_text(
states.response_body_state,
))
.bind(usage_body_capture_state_bind_text(
states.client_response_body_state,
))
.bind(body_capture_mode)
.execute(executor)
.await
@@ -5661,6 +5804,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -6276,6 +6423,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -6371,6 +6522,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,
@@ -6445,6 +6600,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: Some("cand-typed".to_string()),
candidate_index: Some(2),
key_name: Some("primary".to_string()),
@@ -6576,6 +6735,10 @@ mod tests {
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
request_body_state: None,
provider_request_body_state: None,
response_body_state: None,
client_response_body_state: None,
candidate_id: None,
candidate_index: None,
key_name: None,

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()