mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-12 14:10:19 +08:00
refactor(workspace): enforce layered crate boundaries
This commit is contained in:
elky
@@ -0,0 +1,829 @@
|
||||
use std::io::{self, Write};
|
||||
|
||||
use aether_data_contracts::repository::usage::{
|
||||
UpsertUsageRecord, UsageBodyCaptureState, UsageBodyField,
|
||||
};
|
||||
use serde::Serialize;
|
||||
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(Default)]
|
||||
struct CountingWriter {
|
||||
bytes: u64,
|
||||
}
|
||||
|
||||
impl Write for CountingWriter {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
self.bytes = self.bytes.saturating_add(buf.len() as u64);
|
||||
Ok(buf.len())
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[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 = json_serialized_len(&value);
|
||||
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 = json_serialized_len(&truncated_value);
|
||||
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 {
|
||||
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 json_serialized_len(&candidate).is_some_and(|bytes| bytes <= max_bytes as u64) {
|
||||
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()
|
||||
}
|
||||
|
||||
fn json_serialized_len<T: Serialize>(value: &T) -> Option<u64> {
|
||||
let mut writer = CountingWriter::default();
|
||||
serde_json::to_writer(&mut writer, value).ok()?;
|
||||
Some(writer.bytes)
|
||||
}
|
||||
|
||||
pub(crate) fn sync_usage_body_ref_metadata(
|
||||
metadata: &mut Option<Value>,
|
||||
field: UsageBodyField,
|
||||
body_ref: Option<&str>,
|
||||
) {
|
||||
let key = field.as_ref_key();
|
||||
let Some(body_ref) = body_ref.map(str::trim).filter(|value| !value.is_empty()) else {
|
||||
let clear_metadata = match metadata.as_mut() {
|
||||
Some(Value::Object(object)) => {
|
||||
object.remove(key);
|
||||
object.is_empty()
|
||||
}
|
||||
_ => false,
|
||||
};
|
||||
if clear_metadata {
|
||||
*metadata = None;
|
||||
}
|
||||
return;
|
||||
};
|
||||
if let Some(Value::Object(object)) = metadata.as_mut() {
|
||||
if object.get(key).and_then(Value::as_str) == Some(body_ref) {
|
||||
return;
|
||||
}
|
||||
object.insert(key.to_owned(), Value::String(body_ref.to_owned()));
|
||||
return;
|
||||
}
|
||||
let object = metadata
|
||||
.get_or_insert_with(|| Value::Object(Map::new()))
|
||||
.as_object_mut();
|
||||
let Some(object) = object else {
|
||||
return;
|
||||
};
|
||||
object.insert(key.to_owned(), Value::String(body_ref.to_owned()));
|
||||
}
|
||||
|
||||
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 provider_decoded_len = provider_body_base64.and_then(decoded_base64_len_hint);
|
||||
let client_decoded_len = client_body_base64.and_then(decoded_base64_len_hint);
|
||||
let body_capture_capacity =
|
||||
usize::from(provider_body_state.is_some()) + usize::from(client_body_state.is_some());
|
||||
let mut metadata = Map::with_capacity(
|
||||
usize::from(provider_decoded_len.is_some())
|
||||
+ usize::from(client_decoded_len.is_some())
|
||||
+ usize::from(body_capture_capacity > 0),
|
||||
);
|
||||
if let Some(decoded_len) = provider_decoded_len {
|
||||
metadata.insert(
|
||||
"provider_response_body_base64_bytes".to_string(),
|
||||
Value::Number(decoded_len.into()),
|
||||
);
|
||||
}
|
||||
if let Some(decoded_len) = client_decoded_len {
|
||||
metadata.insert(
|
||||
"client_response_body_base64_bytes".to_string(),
|
||||
Value::Number(decoded_len.into()),
|
||||
);
|
||||
}
|
||||
|
||||
if body_capture_capacity > 0 {
|
||||
let mut body_capture = Map::with_capacity(body_capture_capacity);
|
||||
append_body_capture_metadata_entry(
|
||||
&mut body_capture,
|
||||
"response",
|
||||
provider_body_state,
|
||||
provider_decoded_len,
|
||||
provider_decoded_len,
|
||||
);
|
||||
append_body_capture_metadata_entry(
|
||||
&mut body_capture,
|
||||
"client_response",
|
||||
client_body_state,
|
||||
client_decoded_len,
|
||||
client_decoded_len,
|
||||
);
|
||||
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,
|
||||
);
|
||||
let Some(body_capture_object) = body_capture_object_mut(metadata, 2) else {
|
||||
return;
|
||||
};
|
||||
body_capture_object.insert(
|
||||
"request".to_string(),
|
||||
build_body_capture_metadata_entry(states.request, None, None, None),
|
||||
);
|
||||
body_capture_object.insert(
|
||||
"provider_request".to_string(),
|
||||
build_body_capture_metadata_entry(
|
||||
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> {
|
||||
provider_request_body_base64?;
|
||||
let mut metadata = Map::with_capacity(2);
|
||||
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()),
|
||||
);
|
||||
}
|
||||
let Some(body_capture_object) = body_capture_object_mut(metadata, 1) else {
|
||||
return;
|
||||
};
|
||||
body_capture_object.insert(
|
||||
"provider_request".to_string(),
|
||||
build_body_capture_metadata_entry(
|
||||
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;
|
||||
};
|
||||
target.insert(
|
||||
key.to_string(),
|
||||
build_body_capture_metadata_entry(
|
||||
state,
|
||||
stored_bytes,
|
||||
source_bytes,
|
||||
matches!(state, UsageBodyCaptureState::Truncated)
|
||||
.then_some("body_capture_limit_exceeded"),
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
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 Some(body_capture_object) = body_capture_value_object_mut(metadata, 1) else {
|
||||
return;
|
||||
};
|
||||
body_capture_object.insert(
|
||||
key.to_string(),
|
||||
build_body_capture_metadata_entry(state, stored_bytes, source_bytes, reason),
|
||||
);
|
||||
}
|
||||
|
||||
fn body_capture_object_mut(
|
||||
metadata: &mut Map<String, Value>,
|
||||
capacity: usize,
|
||||
) -> Option<&mut Map<String, Value>> {
|
||||
let body_capture = metadata
|
||||
.entry("body_capture".to_string())
|
||||
.or_insert_with(|| Value::Object(Map::with_capacity(capacity)));
|
||||
body_capture.as_object_mut()
|
||||
}
|
||||
|
||||
fn body_capture_value_object_mut(
|
||||
metadata: &mut Option<Value>,
|
||||
capacity: usize,
|
||||
) -> Option<&mut Map<String, Value>> {
|
||||
let metadata_object = metadata
|
||||
.get_or_insert_with(|| Value::Object(Map::with_capacity(1)))
|
||||
.as_object_mut();
|
||||
let metadata_object = metadata_object?;
|
||||
body_capture_object_mut(metadata_object, capacity)
|
||||
}
|
||||
|
||||
fn build_body_capture_metadata_entry(
|
||||
state: UsageBodyCaptureState,
|
||||
stored_bytes: Option<u64>,
|
||||
source_bytes: Option<u64>,
|
||||
reason: Option<&str>,
|
||||
) -> Value {
|
||||
let mut entry = Map::with_capacity(
|
||||
1 + usize::from(stored_bytes.is_some())
|
||||
+ usize::from(source_bytes.is_some())
|
||||
+ usize::from(reason.is_some()),
|
||||
);
|
||||
entry.insert(
|
||||
"state".to_string(),
|
||||
Value::String(state.as_str().to_owned()),
|
||||
);
|
||||
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_owned()));
|
||||
}
|
||||
Value::Object(entry)
|
||||
}
|
||||
|
||||
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.and_then(trim_owned_non_empty_string)
|
||||
}
|
||||
|
||||
fn trim_owned_non_empty_string(value: String) -> Option<String> {
|
||||
let trimmed = value.trim();
|
||||
if trimmed.is_empty() {
|
||||
return None;
|
||||
}
|
||||
if trimmed.len() == value.len() {
|
||||
return Some(value);
|
||||
}
|
||||
Some(trimmed.to_string())
|
||||
}
|
||||
|
||||
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",
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{
|
||||
build_plan_body_capture_metadata, sync_usage_body_ref_metadata,
|
||||
trim_owned_non_empty_string, truncate_usage_body_string,
|
||||
upsert_body_capture_metadata_value_entry,
|
||||
};
|
||||
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
|
||||
use aether_data_contracts::repository::usage::UsageBodyField;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
#[test]
|
||||
fn build_plan_body_capture_metadata_returns_none_without_base64_body() {
|
||||
assert!(build_plan_body_capture_metadata(None).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn trim_owned_non_empty_string_preserves_clean_values_and_drops_blank_ones() {
|
||||
assert_eq!(
|
||||
trim_owned_non_empty_string("blob://body-ref-1".to_string()),
|
||||
Some("blob://body-ref-1".to_string()),
|
||||
);
|
||||
assert_eq!(
|
||||
trim_owned_non_empty_string(" blob://body-ref-1 ".to_string()),
|
||||
Some("blob://body-ref-1".to_string()),
|
||||
);
|
||||
assert_eq!(trim_owned_non_empty_string(" ".to_string()), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upsert_body_capture_metadata_value_entry_ignores_none_state() {
|
||||
let mut metadata = Some(Value::Object(Map::<String, Value>::new()));
|
||||
upsert_body_capture_metadata_value_entry(&mut metadata, "response", None, None, None, None);
|
||||
assert_eq!(metadata, Some(Value::Object(Map::new())));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn upsert_body_capture_metadata_value_entry_preserves_existing_metadata_fields() {
|
||||
let mut metadata = Some(Value::Object(Map::from_iter([(
|
||||
"request_body_ref".to_string(),
|
||||
Value::String("blob://body-ref-1".to_string()),
|
||||
)])));
|
||||
|
||||
upsert_body_capture_metadata_value_entry(
|
||||
&mut metadata,
|
||||
"response",
|
||||
Some(UsageBodyCaptureState::Reference),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
metadata,
|
||||
Some(Value::Object(Map::from_iter([
|
||||
(
|
||||
"request_body_ref".to_string(),
|
||||
Value::String("blob://body-ref-1".to_string()),
|
||||
),
|
||||
(
|
||||
"body_capture".to_string(),
|
||||
Value::Object(Map::from_iter([(
|
||||
"response".to_string(),
|
||||
Value::Object(Map::from_iter([(
|
||||
"state".to_string(),
|
||||
Value::String("reference".to_string()),
|
||||
)])),
|
||||
)])),
|
||||
),
|
||||
]))),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_usage_body_ref_metadata_clears_empty_metadata_object() {
|
||||
let mut metadata = Some(Value::Object(Map::from_iter([(
|
||||
"request_body_ref".to_string(),
|
||||
Value::String("blob://body-ref-1".to_string()),
|
||||
)])));
|
||||
|
||||
sync_usage_body_ref_metadata(&mut metadata, UsageBodyField::RequestBody, None);
|
||||
|
||||
assert!(metadata.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sync_usage_body_ref_metadata_preserves_existing_ref_value() {
|
||||
let mut metadata = Some(Value::Object(Map::from_iter([(
|
||||
"request_body_ref".to_string(),
|
||||
Value::String("blob://body-ref-1".to_string()),
|
||||
)])));
|
||||
|
||||
sync_usage_body_ref_metadata(
|
||||
&mut metadata,
|
||||
UsageBodyField::RequestBody,
|
||||
Some("blob://body-ref-1"),
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
metadata,
|
||||
Some(Value::Object(Map::from_iter([(
|
||||
"request_body_ref".to_string(),
|
||||
Value::String("blob://body-ref-1".to_string()),
|
||||
)]))),
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_usage_body_string_respects_json_byte_limit() {
|
||||
let limit = 32usize;
|
||||
let truncated = truncate_usage_body_string("x".repeat(256).as_str(), limit);
|
||||
|
||||
assert!(truncated.ends_with("...[truncated]"));
|
||||
assert!(serde_json::to_vec(&truncated)
|
||||
.ok()
|
||||
.is_some_and(|bytes| bytes.len() <= limit));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,208 @@
|
||||
use aether_data_contracts::DataLayerError;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct UsageRuntimeConfig {
|
||||
pub enabled: bool,
|
||||
pub queue_terminal_events: bool,
|
||||
pub queue_lifecycle_events: bool,
|
||||
pub worker_count: usize,
|
||||
pub worker_autoscale_enabled: bool,
|
||||
pub worker_max_count: usize,
|
||||
pub worker_record_concurrency_limit: Option<usize>,
|
||||
pub worker_scale_interval_ms: u64,
|
||||
pub worker_idle_scale_down_ticks: u64,
|
||||
pub stream_key: String,
|
||||
pub consumer_group: String,
|
||||
pub dlq_stream_key: String,
|
||||
pub stream_maxlen: usize,
|
||||
pub consumer_batch_size: usize,
|
||||
pub consumer_block_ms: u64,
|
||||
pub reclaim_idle_ms: u64,
|
||||
pub reclaim_count: usize,
|
||||
pub reclaim_interval_ms: u64,
|
||||
pub terminal_enqueue_max_in_flight: u64,
|
||||
pub lifecycle_enqueue_max_in_flight: u64,
|
||||
pub lifecycle_enqueue_delay_ms: u64,
|
||||
pub retry_deferred_lifecycle_events: bool,
|
||||
pub enqueue_retry_buffer_capacity: usize,
|
||||
pub enqueue_retry_workers: usize,
|
||||
pub enqueue_retry_initial_backoff_ms: u64,
|
||||
pub enqueue_retry_max_backoff_ms: u64,
|
||||
}
|
||||
|
||||
impl Default for UsageRuntimeConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false,
|
||||
queue_terminal_events: false,
|
||||
queue_lifecycle_events: false,
|
||||
worker_count: 4,
|
||||
worker_autoscale_enabled: true,
|
||||
worker_max_count: 32,
|
||||
worker_record_concurrency_limit: Some(32),
|
||||
worker_scale_interval_ms: 1_000,
|
||||
worker_idle_scale_down_ticks: 30,
|
||||
stream_key: "usage:events".to_string(),
|
||||
consumer_group: "usage_consumers".to_string(),
|
||||
dlq_stream_key: "usage:events:dlq".to_string(),
|
||||
stream_maxlen: 200_000,
|
||||
consumer_batch_size: 128,
|
||||
consumer_block_ms: 500,
|
||||
reclaim_idle_ms: 60_000,
|
||||
reclaim_count: 128,
|
||||
reclaim_interval_ms: 5_000,
|
||||
terminal_enqueue_max_in_flight: 1_024,
|
||||
lifecycle_enqueue_max_in_flight: 512,
|
||||
lifecycle_enqueue_delay_ms: 1_000,
|
||||
retry_deferred_lifecycle_events: true,
|
||||
enqueue_retry_buffer_capacity: 131_072,
|
||||
enqueue_retry_workers: 8,
|
||||
enqueue_retry_initial_backoff_ms: 3_000,
|
||||
enqueue_retry_max_backoff_ms: 10_000,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl UsageRuntimeConfig {
|
||||
pub fn disabled() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn validate(&self) -> Result<(), DataLayerError> {
|
||||
if !self.enabled {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
if self.stream_key.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime stream_key cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.consumer_group.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime consumer_group cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.dlq_stream_key.trim().is_empty() {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime dlq_stream_key cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.worker_count == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime worker_count must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.worker_max_count == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime worker_max_count must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self
|
||||
.worker_record_concurrency_limit
|
||||
.is_some_and(|limit| limit == 0)
|
||||
{
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime worker_record_concurrency_limit must be positive when set"
|
||||
.to_string(),
|
||||
));
|
||||
}
|
||||
if self.worker_scale_interval_ms == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime worker_scale_interval_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.worker_idle_scale_down_ticks == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime worker_idle_scale_down_ticks must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.stream_maxlen == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime stream_maxlen must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.consumer_batch_size == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime consumer_batch_size must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.consumer_block_ms == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime consumer_block_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.reclaim_idle_ms == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime reclaim_idle_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.reclaim_count == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime reclaim_count must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.reclaim_interval_ms == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime reclaim_interval_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.terminal_enqueue_max_in_flight == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime terminal_enqueue_max_in_flight must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.lifecycle_enqueue_max_in_flight == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime lifecycle_enqueue_max_in_flight must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.enqueue_retry_buffer_capacity == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime enqueue_retry_buffer_capacity must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.enqueue_retry_workers == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime enqueue_retry_workers must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.enqueue_retry_initial_backoff_ms == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime enqueue_retry_initial_backoff_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.enqueue_retry_max_backoff_ms == 0 {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime enqueue_retry_max_backoff_ms must be positive".to_string(),
|
||||
));
|
||||
}
|
||||
if self.enqueue_retry_initial_backoff_ms > self.enqueue_retry_max_backoff_ms {
|
||||
return Err(DataLayerError::InvalidConfiguration(
|
||||
"usage runtime enqueue retry initial backoff cannot exceed max backoff".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::UsageRuntimeConfig;
|
||||
|
||||
#[test]
|
||||
fn disabled_config_is_valid() {
|
||||
assert!(UsageRuntimeConfig::disabled().validate().is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn enabled_config_rejects_empty_stream_key() {
|
||||
let config = UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
stream_key: String::new(),
|
||||
..UsageRuntimeConfig::default()
|
||||
};
|
||||
assert!(config.validate().is_err());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,257 @@
|
||||
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;
|
||||
|
||||
pub const USAGE_EVENT_VERSION: u8 = 1;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum UsageEventType {
|
||||
Pending,
|
||||
Streaming,
|
||||
Completed,
|
||||
Failed,
|
||||
Cancelled,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
|
||||
pub struct UsageEventData {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub user_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub api_key_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub username: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub api_key_name: Option<String>,
|
||||
pub provider_name: String,
|
||||
pub model: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub target_model: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub model_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub global_model_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_endpoint_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_api_key_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub request_type: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub api_format: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub api_family: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub endpoint_kind: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub endpoint_api_format: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_api_family: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_endpoint_kind: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub has_format_conversion: Option<bool>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub is_stream: Option<bool>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub input_tokens: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub output_tokens: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub total_tokens: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub cache_creation_input_tokens: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub cache_creation_ephemeral_5m_input_tokens: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub cache_creation_ephemeral_1h_input_tokens: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub cache_read_input_tokens: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub cache_creation_cost_usd: Option<f64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub cache_read_cost_usd: Option<f64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub output_price_per_1m: Option<f64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub total_cost_usd: Option<f64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub actual_total_cost_usd: Option<f64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub status_code: Option<u16>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub error_message: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub error_category: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub response_time_ms: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub first_byte_time_ms: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub request_headers: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub request_body: Option<Value>,
|
||||
#[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>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub key_name: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub planner_kind: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub route_family: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub route_kind: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub execution_path: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub local_execution_runtime_miss_reason: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub request_metadata: Option<Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
pub struct UsageEvent {
|
||||
pub event_type: UsageEventType,
|
||||
pub request_id: String,
|
||||
pub timestamp_ms: u64,
|
||||
pub data: UsageEventData,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
|
||||
struct UsageEventEnvelope {
|
||||
v: u8,
|
||||
#[serde(rename = "type")]
|
||||
event_type: UsageEventType,
|
||||
request_id: String,
|
||||
timestamp_ms: u64,
|
||||
data: UsageEventData,
|
||||
}
|
||||
|
||||
impl UsageEvent {
|
||||
pub fn new(
|
||||
event_type: UsageEventType,
|
||||
request_id: impl Into<String>,
|
||||
data: UsageEventData,
|
||||
) -> Self {
|
||||
Self {
|
||||
event_type,
|
||||
request_id: request_id.into(),
|
||||
timestamp_ms: now_ms(),
|
||||
data,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_stream_fields(&self) -> Result<BTreeMap<String, String>, DataLayerError> {
|
||||
let payload = UsageEventEnvelope {
|
||||
v: USAGE_EVENT_VERSION,
|
||||
event_type: self.event_type,
|
||||
request_id: self.request_id.clone(),
|
||||
timestamp_ms: self.timestamp_ms,
|
||||
data: self.data.clone(),
|
||||
};
|
||||
let payload = serde_json::to_string(&payload).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"failed to serialize usage event payload: {err}"
|
||||
))
|
||||
})?;
|
||||
Ok(BTreeMap::from([("payload".to_string(), payload)]))
|
||||
}
|
||||
|
||||
pub fn from_stream_fields(fields: &BTreeMap<String, String>) -> Result<Self, DataLayerError> {
|
||||
let payload = fields.get("payload").ok_or_else(|| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"usage event stream entry missing payload field".to_string(),
|
||||
)
|
||||
})?;
|
||||
let envelope: UsageEventEnvelope = serde_json::from_str(payload).map_err(|err| {
|
||||
DataLayerError::UnexpectedValue(format!(
|
||||
"failed to deserialize usage event payload: {err}"
|
||||
))
|
||||
})?;
|
||||
if envelope.v != USAGE_EVENT_VERSION {
|
||||
return Err(DataLayerError::UnexpectedValue(format!(
|
||||
"unsupported usage event version: {}",
|
||||
envelope.v
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
event_type: envelope.event_type,
|
||||
request_id: envelope.request_id,
|
||||
timestamp_ms: envelope.timestamp_ms,
|
||||
data: envelope.data,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
pub fn now_ms() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_millis() as u64
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{UsageEvent, UsageEventData, UsageEventType};
|
||||
|
||||
#[test]
|
||||
fn usage_event_round_trips_through_stream_fields() {
|
||||
let event = UsageEvent::new(
|
||||
UsageEventType::Completed,
|
||||
"req-1",
|
||||
UsageEventData {
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(20),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
);
|
||||
|
||||
let fields = event.to_stream_fields().expect("event should serialize");
|
||||
let parsed = UsageEvent::from_stream_fields(&fields).expect("event should parse");
|
||||
|
||||
assert_eq!(parsed.request_id, "req-1");
|
||||
assert_eq!(parsed.event_type, UsageEventType::Completed);
|
||||
assert_eq!(parsed.data.total_tokens, None);
|
||||
assert_eq!(parsed.data.output_tokens, Some(20));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,80 @@
|
||||
use std::future::Future;
|
||||
use std::sync::OnceLock;
|
||||
|
||||
const DEFAULT_USAGE_BACKGROUND_RUNTIME_THREADS: usize = 8;
|
||||
const MAX_USAGE_BACKGROUND_RUNTIME_THREADS: usize = 64;
|
||||
const GATEWAY_USAGE_BACKGROUND_RUNTIME_THREADS_ENV: &str = "AETHER_GATEWAY_USAGE_RUNTIME_THREADS";
|
||||
const USAGE_BACKGROUND_RUNTIME_THREADS_ENV: &str = "AETHER_USAGE_RUNTIME_THREADS";
|
||||
const USAGE_BACKGROUND_RUNTIME_STACK_BYTES: usize = 8 * 1024 * 1024;
|
||||
const USAGE_BACKGROUND_RUNTIME_THREAD_NAME: &str = "aether-usage-runtime";
|
||||
|
||||
pub(crate) fn spawn_on_usage_background_runtime<F>(task: F) -> tokio::task::JoinHandle<F::Output>
|
||||
where
|
||||
F: Future + Send + 'static,
|
||||
F::Output: Send + 'static,
|
||||
{
|
||||
usage_background_runtime().handle().spawn(task)
|
||||
}
|
||||
|
||||
fn usage_background_runtime() -> &'static tokio::runtime::Runtime {
|
||||
static RUNTIME: OnceLock<&'static tokio::runtime::Runtime> = OnceLock::new();
|
||||
|
||||
RUNTIME.get_or_init(|| {
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.worker_threads(usage_background_runtime_threads())
|
||||
.thread_name(USAGE_BACKGROUND_RUNTIME_THREAD_NAME)
|
||||
.thread_stack_size(USAGE_BACKGROUND_RUNTIME_STACK_BYTES)
|
||||
.build()
|
||||
.expect("usage background runtime should build");
|
||||
Box::leak(Box::new(runtime))
|
||||
})
|
||||
}
|
||||
|
||||
fn usage_background_runtime_threads() -> usize {
|
||||
parse_usage_background_runtime_threads(
|
||||
std::env::var(GATEWAY_USAGE_BACKGROUND_RUNTIME_THREADS_ENV)
|
||||
.ok()
|
||||
.or_else(|| std::env::var(USAGE_BACKGROUND_RUNTIME_THREADS_ENV).ok())
|
||||
.as_deref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn parse_usage_background_runtime_threads(value: Option<&str>) -> usize {
|
||||
value
|
||||
.and_then(|value| value.trim().parse::<usize>().ok())
|
||||
.filter(|threads| *threads > 0)
|
||||
.unwrap_or(DEFAULT_USAGE_BACKGROUND_RUNTIME_THREADS)
|
||||
.clamp(1, MAX_USAGE_BACKGROUND_RUNTIME_THREADS)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{parse_usage_background_runtime_threads, spawn_on_usage_background_runtime};
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_background_runtime_runs_on_dedicated_named_threads() {
|
||||
let thread_name = spawn_on_usage_background_runtime(async move {
|
||||
std::thread::current()
|
||||
.name()
|
||||
.unwrap_or_default()
|
||||
.to_string()
|
||||
})
|
||||
.await
|
||||
.expect("background task should complete");
|
||||
|
||||
assert_eq!(thread_name, "aether-usage-runtime");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_background_runtime_threads_are_configurable() {
|
||||
assert_eq!(parse_usage_background_runtime_threads(None), 8);
|
||||
assert_eq!(parse_usage_background_runtime_threads(Some("12")), 12);
|
||||
assert_eq!(parse_usage_background_runtime_threads(Some("0")), 8);
|
||||
assert_eq!(
|
||||
parse_usage_background_runtime_threads(Some("not-a-number")),
|
||||
8
|
||||
);
|
||||
assert_eq!(parse_usage_background_runtime_threads(Some("999")), 64);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::{Arc, Mutex as StdMutex, Weak};
|
||||
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
const DEFAULT_PRUNE_THRESHOLD: usize = 8_192;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub(crate) struct KeyedAsyncLockPool {
|
||||
entries: StdMutex<HashMap<String, Weak<Mutex<()>>>>,
|
||||
prune_threshold: usize,
|
||||
}
|
||||
|
||||
impl Default for KeyedAsyncLockPool {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
entries: StdMutex::new(HashMap::new()),
|
||||
prune_threshold: DEFAULT_PRUNE_THRESHOLD,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl KeyedAsyncLockPool {
|
||||
pub(crate) fn lock_for(&self, key: &str) -> Arc<Mutex<()>> {
|
||||
let mut entries = self
|
||||
.entries
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner());
|
||||
if let Some(existing) = entries.get(key).and_then(Weak::upgrade) {
|
||||
return existing;
|
||||
}
|
||||
if entries.len() >= self.prune_threshold {
|
||||
entries.retain(|_, lock| lock.strong_count() > 0);
|
||||
}
|
||||
let lock = Arc::new(Mutex::new(()));
|
||||
entries.insert(key.to_string(), Arc::downgrade(&lock));
|
||||
lock
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn tracked_keys(&self) -> usize {
|
||||
self.entries
|
||||
.lock()
|
||||
.unwrap_or_else(|poisoned| poisoned.into_inner())
|
||||
.len()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::KeyedAsyncLockPool;
|
||||
use std::sync::Arc;
|
||||
|
||||
#[test]
|
||||
fn reuses_active_key_without_serializing_different_keys() {
|
||||
let pool = KeyedAsyncLockPool::default();
|
||||
let first = pool.lock_for("request-a");
|
||||
let same = pool.lock_for("request-a");
|
||||
let different = pool.lock_for("request-b");
|
||||
|
||||
assert!(Arc::ptr_eq(&first, &same));
|
||||
assert!(!Arc::ptr_eq(&first, &different));
|
||||
assert_eq!(pool.tracked_keys(), 2);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
mod body_capture;
|
||||
pub mod config;
|
||||
pub mod event;
|
||||
mod executor;
|
||||
mod keyed_lock;
|
||||
pub mod queue;
|
||||
pub mod record;
|
||||
pub mod report;
|
||||
pub mod report_context;
|
||||
mod request_metadata;
|
||||
pub mod runtime;
|
||||
pub mod settlement;
|
||||
pub mod standardized_usage;
|
||||
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;
|
||||
pub use record::build_upsert_usage_record_from_event;
|
||||
pub use report::{
|
||||
extract_gemini_file_mapping_entries, gemini_file_mapping_cache_key,
|
||||
infer_internal_finalize_signature, is_local_ai_stream_report_kind,
|
||||
is_local_ai_sync_report_kind, normalize_gemini_file_name, report_request_id,
|
||||
resolve_internal_finalize_route, should_handle_local_stream_report,
|
||||
should_handle_local_sync_report, stream_capture_terminal_state,
|
||||
stream_report_missing_terminal_event, stream_report_represents_failure,
|
||||
stream_report_requires_observed_terminal_event, sync_report_represents_failure,
|
||||
GatewayStreamReportRequest, GatewaySyncReportRequest, GeminiFileMappingEntry,
|
||||
InternalFinalizeRoute, StreamCapturedTerminalState, GEMINI_FILE_MAPPING_TTL_SECONDS,
|
||||
STREAM_MISSING_TERMINAL_EVENT_CATEGORY, STREAM_MISSING_TERMINAL_EVENT_MESSAGE,
|
||||
STREAM_TERMINAL_ERROR_CATEGORY, STREAM_TERMINAL_ERROR_MESSAGE,
|
||||
};
|
||||
pub use report_context::{
|
||||
build_locally_actionable_report_context_from_request_candidate,
|
||||
build_locally_actionable_report_context_from_video_task, report_context_is_locally_actionable,
|
||||
};
|
||||
pub use runtime::{
|
||||
UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageQueueHealthSnapshot,
|
||||
UsageRequestRecordLevel, UsageRuntime, UsageRuntimeAccess, UsageRuntimeMetricsSnapshot,
|
||||
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;
|
||||
pub use usage_mapper::{map_usage, map_usage_from_response, UsageMapper};
|
||||
pub use worker::{
|
||||
build_usage_queue_worker, write_event_record, ManualProxyNodeCounter, UsageDataEventRecorder,
|
||||
UsageEventRecorder, UsageQueueWorker, UsageRecordWriter,
|
||||
};
|
||||
pub use write::{
|
||||
build_lifecycle_usage_seed, build_pending_usage_record, build_pending_usage_record_from_seed,
|
||||
build_stream_terminal_usage_event, build_stream_terminal_usage_outcome,
|
||||
build_stream_terminal_usage_payload_seed, build_stream_terminal_usage_seed,
|
||||
build_streaming_usage_record, build_streaming_usage_record_from_seed,
|
||||
build_sync_terminal_usage_event, build_sync_terminal_usage_outcome,
|
||||
build_sync_terminal_usage_payload_seed, build_sync_terminal_usage_seed,
|
||||
build_terminal_usage_context_seed, build_terminal_usage_event_from_outcome,
|
||||
build_terminal_usage_event_from_seed, build_usage_event_data_seed, LifecycleUsageSeed,
|
||||
StreamTerminalUsagePayloadSeed, SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed,
|
||||
TerminalUsageOutcome, TerminalUsageSeed, UsageTerminalState,
|
||||
};
|
||||
@@ -0,0 +1,167 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde_json::json;
|
||||
|
||||
use aether_data_contracts::DataLayerError;
|
||||
use aether_runtime_state::{
|
||||
RuntimeQueueEntry, RuntimeQueueReclaimConfig, RuntimeQueueStats, RuntimeQueueStore,
|
||||
};
|
||||
|
||||
use super::config::UsageRuntimeConfig;
|
||||
use super::event::UsageEvent;
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct UsageQueue {
|
||||
runner: Arc<dyn RuntimeQueueStore>,
|
||||
config: UsageRuntimeConfig,
|
||||
stream: String,
|
||||
group: String,
|
||||
dlq_stream: String,
|
||||
}
|
||||
|
||||
impl UsageQueue {
|
||||
pub fn new(
|
||||
runner: Arc<dyn RuntimeQueueStore>,
|
||||
config: UsageRuntimeConfig,
|
||||
) -> Result<Self, DataLayerError> {
|
||||
config.validate()?;
|
||||
Ok(Self {
|
||||
runner,
|
||||
stream: config.stream_key.clone(),
|
||||
group: config.consumer_group.clone(),
|
||||
dlq_stream: config.dlq_stream_key.clone(),
|
||||
config,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn ensure_consumer_group(&self) -> Result<(), DataLayerError> {
|
||||
self.runner
|
||||
.ensure_consumer_group(&self.stream, &self.group, "0-0")
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn enqueue(&self, event: &UsageEvent) -> Result<String, DataLayerError> {
|
||||
let fields = event.to_stream_fields()?;
|
||||
self.runner
|
||||
.append_fields_with_maxlen(&self.stream, &fields, Some(self.config.stream_maxlen))
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn read_group(
|
||||
&self,
|
||||
consumer: &str,
|
||||
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
|
||||
self.runner
|
||||
.read_group(
|
||||
&self.stream,
|
||||
&self.group,
|
||||
consumer,
|
||||
self.config.consumer_batch_size.max(1),
|
||||
Some(self.config.consumer_block_ms.max(1)),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn claim_stale(
|
||||
&self,
|
||||
consumer: &str,
|
||||
start_id: &str,
|
||||
) -> Result<Vec<RuntimeQueueEntry>, DataLayerError> {
|
||||
self.runner
|
||||
.claim_stale(
|
||||
&self.stream,
|
||||
&self.group,
|
||||
consumer,
|
||||
start_id,
|
||||
RuntimeQueueReclaimConfig {
|
||||
min_idle_ms: self.config.reclaim_idle_ms,
|
||||
count: self.config.reclaim_count,
|
||||
},
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn ack_and_delete(&self, ids: &[String]) -> Result<(), DataLayerError> {
|
||||
self.runner.ack(&self.stream, &self.group, ids).await?;
|
||||
self.runner.delete(&self.stream, ids).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn push_dead_letter(
|
||||
&self,
|
||||
entry: &RuntimeQueueEntry,
|
||||
error: &str,
|
||||
) -> Result<String, DataLayerError> {
|
||||
let fields = std::collections::BTreeMap::from([(
|
||||
"payload".to_string(),
|
||||
serde_json::to_string(&json!({
|
||||
"entry_id": entry.id,
|
||||
"fields": entry.fields,
|
||||
"error": error,
|
||||
}))
|
||||
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?,
|
||||
)]);
|
||||
self.runner
|
||||
.append_fields_with_maxlen(&self.dlq_stream, &fields, None)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn stats(&self) -> Result<RuntimeQueueStats, DataLayerError> {
|
||||
self.runner.stats(&self.stream, Some(&self.group)).await
|
||||
}
|
||||
|
||||
pub async fn dlq_stats(&self) -> Result<RuntimeQueueStats, DataLayerError> {
|
||||
self.runner.stats(&self.dlq_stream, None).await
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
#[cfg(test)]
|
||||
struct UsageQueueRuntimeSettings {
|
||||
command_timeout_ms: Option<u64>,
|
||||
read_block_ms: Option<u64>,
|
||||
read_count: usize,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
fn usage_queue_runtime_settings(config: &UsageRuntimeConfig) -> UsageQueueRuntimeSettings {
|
||||
let read_block_ms = config.consumer_block_ms.max(1);
|
||||
let command_timeout_ms = read_block_ms.saturating_add(2_000).max(5_000);
|
||||
UsageQueueRuntimeSettings {
|
||||
command_timeout_ms: Some(command_timeout_ms),
|
||||
read_block_ms: Some(read_block_ms),
|
||||
read_count: config.consumer_batch_size.max(1),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{usage_queue_runtime_settings, UsageQueue, UsageQueueRuntimeSettings};
|
||||
use crate::UsageRuntimeConfig;
|
||||
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeState};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[test]
|
||||
fn usage_queue_applies_runtime_block_and_batch_settings() {
|
||||
let config = UsageRuntimeConfig {
|
||||
enabled: true,
|
||||
consumer_block_ms: 750,
|
||||
consumer_batch_size: 123,
|
||||
..UsageRuntimeConfig::default()
|
||||
};
|
||||
let queue = UsageQueue::new(
|
||||
Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default())),
|
||||
config,
|
||||
)
|
||||
.expect("usage queue should build from runtime config");
|
||||
|
||||
assert_eq!(
|
||||
usage_queue_runtime_settings(&queue.config),
|
||||
UsageQueueRuntimeSettings {
|
||||
command_timeout_ms: Some(5_000),
|
||||
read_block_ms: Some(750),
|
||||
read_count: 123,
|
||||
}
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,294 @@
|
||||
use aether_data_contracts::repository::usage::UpsertUsageRecord;
|
||||
use aether_data_contracts::DataLayerError;
|
||||
|
||||
use crate::request_metadata::{
|
||||
attach_provider_request_body_metadata, sanitize_usage_request_metadata,
|
||||
};
|
||||
use crate::{UsageEvent, UsageEventType};
|
||||
|
||||
fn metadata_string(metadata: Option<&serde_json::Value>, key: &str) -> Option<String> {
|
||||
metadata
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|object| object.get(key))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn metadata_u64(metadata: Option<&serde_json::Value>, key: &str) -> Option<u64> {
|
||||
metadata
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|object| object.get(key))
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_u64()
|
||||
.or_else(|| value.as_i64().and_then(|number| u64::try_from(number).ok()))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_upsert_usage_record_from_event(
|
||||
event: &UsageEvent,
|
||||
) -> Result<UpsertUsageRecord, DataLayerError> {
|
||||
let (status, billing_status) = lifecycle_status_and_billing(event.event_type);
|
||||
let finalized_at_unix_secs = match event.event_type {
|
||||
UsageEventType::Pending | UsageEventType::Streaming => None,
|
||||
UsageEventType::Completed | UsageEventType::Failed | UsageEventType::Cancelled => {
|
||||
Some(event.timestamp_ms / 1_000)
|
||||
}
|
||||
};
|
||||
let mut data = event.data.clone();
|
||||
data.request_metadata = attach_provider_request_body_metadata(
|
||||
data.request_metadata,
|
||||
data.endpoint_api_format
|
||||
.as_deref()
|
||||
.or(data.api_format.as_deref()),
|
||||
data.target_model.as_deref().or(Some(data.model.as_str())),
|
||||
Some(data.model.as_str()),
|
||||
data.provider_request_body.as_ref(),
|
||||
);
|
||||
let now_unix_secs = event.timestamp_ms / 1_000;
|
||||
|
||||
Ok(UpsertUsageRecord {
|
||||
request_id: event.request_id.clone(),
|
||||
user_id: data.user_id,
|
||||
api_key_id: data.api_key_id,
|
||||
username: data.username,
|
||||
api_key_name: data.api_key_name,
|
||||
provider_name: data.provider_name,
|
||||
model: data.model,
|
||||
target_model: data.target_model,
|
||||
provider_id: empty_to_none(data.provider_id),
|
||||
provider_endpoint_id: empty_to_none(data.provider_endpoint_id),
|
||||
provider_api_key_id: empty_to_none(data.provider_api_key_id),
|
||||
request_type: data.request_type,
|
||||
api_format: data.api_format,
|
||||
api_family: data.api_family,
|
||||
endpoint_kind: data.endpoint_kind,
|
||||
endpoint_api_format: data.endpoint_api_format,
|
||||
provider_api_family: data.provider_api_family,
|
||||
provider_endpoint_kind: data.provider_endpoint_kind,
|
||||
has_format_conversion: data.has_format_conversion,
|
||||
is_stream: data.is_stream,
|
||||
input_tokens: data.input_tokens,
|
||||
output_tokens: data.output_tokens,
|
||||
total_tokens: data.total_tokens,
|
||||
cache_creation_input_tokens: data.cache_creation_input_tokens,
|
||||
cache_creation_ephemeral_5m_input_tokens: data.cache_creation_ephemeral_5m_input_tokens,
|
||||
cache_creation_ephemeral_1h_input_tokens: data.cache_creation_ephemeral_1h_input_tokens,
|
||||
cache_read_input_tokens: data.cache_read_input_tokens,
|
||||
cache_creation_cost_usd: data.cache_creation_cost_usd,
|
||||
cache_read_cost_usd: data.cache_read_cost_usd,
|
||||
output_price_per_1m: data.output_price_per_1m,
|
||||
total_cost_usd: data.total_cost_usd,
|
||||
actual_total_cost_usd: data.actual_total_cost_usd,
|
||||
status_code: data.status_code,
|
||||
error_message: data.error_message,
|
||||
error_category: data.error_category,
|
||||
response_time_ms: data.response_time_ms,
|
||||
first_byte_time_ms: data.first_byte_time_ms,
|
||||
status: status.to_string(),
|
||||
billing_status: billing_status.to_string(),
|
||||
request_headers: data.request_headers,
|
||||
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")),
|
||||
candidate_index: data
|
||||
.candidate_index
|
||||
.or_else(|| metadata_u64(data.request_metadata.as_ref(), "candidate_index")),
|
||||
key_name: data
|
||||
.key_name
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "key_name")),
|
||||
planner_kind: data
|
||||
.planner_kind
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "planner_kind")),
|
||||
route_family: data
|
||||
.route_family
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "route_family")),
|
||||
route_kind: data
|
||||
.route_kind
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "route_kind")),
|
||||
execution_path: data
|
||||
.execution_path
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "execution_path")),
|
||||
local_execution_runtime_miss_reason: data.local_execution_runtime_miss_reason.or_else(
|
||||
|| {
|
||||
metadata_string(
|
||||
data.request_metadata.as_ref(),
|
||||
"local_execution_runtime_miss_reason",
|
||||
)
|
||||
},
|
||||
),
|
||||
request_metadata: sanitize_usage_request_metadata(data.request_metadata),
|
||||
finalized_at_unix_secs,
|
||||
created_at_unix_ms: Some(now_unix_secs),
|
||||
updated_at_unix_secs: now_unix_secs,
|
||||
})
|
||||
}
|
||||
|
||||
fn lifecycle_status_and_billing(event_type: UsageEventType) -> (&'static str, &'static str) {
|
||||
match event_type {
|
||||
UsageEventType::Pending => ("pending", "pending"),
|
||||
UsageEventType::Streaming => ("streaming", "pending"),
|
||||
UsageEventType::Completed => ("completed", "pending"),
|
||||
UsageEventType::Failed => ("failed", "void"),
|
||||
UsageEventType::Cancelled => ("cancelled", "void"),
|
||||
}
|
||||
}
|
||||
|
||||
fn empty_to_none(value: Option<String>) -> Option<String> {
|
||||
value
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::{UsageEvent, UsageEventData, UsageEventType};
|
||||
|
||||
use super::build_upsert_usage_record_from_event;
|
||||
|
||||
#[test]
|
||||
fn builds_upsert_record_from_terminal_event() {
|
||||
let record = build_upsert_usage_record_from_event(&UsageEvent {
|
||||
event_type: UsageEventType::Completed,
|
||||
request_id: "req-1".to_string(),
|
||||
timestamp_ms: 1_700_000_000_000,
|
||||
data: UsageEventData {
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("key-1".to_string()),
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(20),
|
||||
total_tokens: Some(30),
|
||||
status_code: Some(200),
|
||||
provider_request_body: Some(serde_json::json!({
|
||||
"reasoning": { "effort": "max" },
|
||||
"service_tier": "priority"
|
||||
})),
|
||||
request_metadata: Some(serde_json::json!({
|
||||
"provider_actual_service_tier": "default"
|
||||
})),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
})
|
||||
.expect("record should build");
|
||||
|
||||
assert_eq!(record.request_id, "req-1");
|
||||
assert_eq!(record.status, "completed");
|
||||
assert_eq!(record.billing_status, "pending");
|
||||
assert_eq!(record.total_tokens, Some(30));
|
||||
assert_eq!(
|
||||
record
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("provider_reasoning_effort"))
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("max")
|
||||
);
|
||||
assert_eq!(
|
||||
record
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("provider_service_tier"))
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("priority")
|
||||
);
|
||||
assert_eq!(
|
||||
record
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|value| value.get("provider_actual_service_tier"))
|
||||
.and_then(serde_json::Value::as_str),
|
||||
Some("default")
|
||||
);
|
||||
assert_eq!(record.finalized_at_unix_secs, Some(1_700_000_000));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cancelled_terminal_record_is_void_for_billing() {
|
||||
let record = build_upsert_usage_record_from_event(&UsageEvent {
|
||||
event_type: UsageEventType::Cancelled,
|
||||
request_id: "req-cancelled".to_string(),
|
||||
timestamp_ms: 1_700_000_000_000,
|
||||
data: UsageEventData {
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(20),
|
||||
total_tokens: Some(30),
|
||||
total_cost_usd: Some(0.03),
|
||||
actual_total_cost_usd: Some(0.02),
|
||||
status_code: Some(499),
|
||||
response_time_ms: Some(200),
|
||||
first_byte_time_ms: Some(50),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
})
|
||||
.expect("record should build");
|
||||
|
||||
assert_eq!(record.status, "cancelled");
|
||||
assert_eq!(record.billing_status, "void");
|
||||
assert_eq!(record.total_tokens, Some(30));
|
||||
assert_eq!(record.total_cost_usd, Some(0.03));
|
||||
assert_eq!(record.actual_total_cost_usd, Some(0.02));
|
||||
assert_eq!(record.status_code, Some(499));
|
||||
assert_eq!(record.response_time_ms, Some(200));
|
||||
assert_eq!(record.first_byte_time_ms, Some(50));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitizes_request_metadata_before_building_upsert_record() {
|
||||
let record = build_upsert_usage_record_from_event(&UsageEvent {
|
||||
event_type: UsageEventType::Completed,
|
||||
request_id: "req-2".to_string(),
|
||||
timestamp_ms: 1_700_000_000_000,
|
||||
data: UsageEventData {
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
request_metadata: Some(serde_json::json!({
|
||||
"request_id": "req-2",
|
||||
"provider_id": "provider-1",
|
||||
"candidate_id": "cand-2",
|
||||
"key_name": "upstream-primary",
|
||||
"billing_snapshot": { "status": "complete" }
|
||||
})),
|
||||
..UsageEventData::default()
|
||||
},
|
||||
})
|
||||
.expect("record should build");
|
||||
|
||||
assert_eq!(record.candidate_id.as_deref(), Some("cand-2"));
|
||||
assert_eq!(record.key_name.as_deref(), Some("upstream-primary"));
|
||||
assert_eq!(
|
||||
record.request_metadata,
|
||||
Some(serde_json::json!({
|
||||
"billing_snapshot": { "status": "complete" }
|
||||
}))
|
||||
);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,237 @@
|
||||
use aether_data_contracts::repository::candidates::StoredRequestCandidate;
|
||||
use aether_data_contracts::repository::video_tasks::StoredVideoTask;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
pub fn report_context_is_locally_actionable(report_context: Option<&Value>) -> bool {
|
||||
let Some(context) = report_context else {
|
||||
return false;
|
||||
};
|
||||
|
||||
has_non_empty_str(context, "request_id")
|
||||
&& (has_non_empty_str(context, "candidate_id")
|
||||
|| has_u64(context, "candidate_index")
|
||||
|| has_non_empty_str(context, "provider_id")
|
||||
|| has_non_empty_str(context, "endpoint_id")
|
||||
|| has_non_empty_str(context, "key_id"))
|
||||
}
|
||||
|
||||
pub fn build_locally_actionable_report_context_from_request_candidate(
|
||||
context: &Value,
|
||||
candidate: &StoredRequestCandidate,
|
||||
) -> Option<Value> {
|
||||
let mut object = context.as_object()?.clone();
|
||||
insert_missing_string_value(&mut object, "candidate_id", Some(candidate.id.as_str()));
|
||||
if !object.contains_key("candidate_index") {
|
||||
object.insert(
|
||||
"candidate_index".to_string(),
|
||||
Value::Number(candidate.candidate_index.into()),
|
||||
);
|
||||
}
|
||||
insert_missing_optional_string_value(
|
||||
&mut object,
|
||||
"provider_id",
|
||||
candidate.provider_id.as_deref(),
|
||||
);
|
||||
insert_missing_optional_string_value(
|
||||
&mut object,
|
||||
"endpoint_id",
|
||||
candidate.endpoint_id.as_deref(),
|
||||
);
|
||||
insert_missing_optional_string_value(&mut object, "key_id", candidate.key_id.as_deref());
|
||||
insert_missing_optional_string_value(&mut object, "user_id", candidate.user_id.as_deref());
|
||||
insert_missing_optional_string_value(
|
||||
&mut object,
|
||||
"api_key_id",
|
||||
candidate.api_key_id.as_deref(),
|
||||
);
|
||||
|
||||
let resolved = Value::Object(object);
|
||||
report_context_is_locally_actionable(Some(&resolved)).then_some(resolved)
|
||||
}
|
||||
|
||||
pub fn build_locally_actionable_report_context_from_video_task(
|
||||
context: &Value,
|
||||
task: &StoredVideoTask,
|
||||
) -> Option<Value> {
|
||||
let mut object = context.as_object()?.clone();
|
||||
insert_missing_string_value(&mut object, "request_id", Some(task.request_id.as_str()));
|
||||
insert_missing_optional_string_value(&mut object, "provider_id", task.provider_id.as_deref());
|
||||
insert_missing_optional_string_value(&mut object, "endpoint_id", task.endpoint_id.as_deref());
|
||||
insert_missing_optional_string_value(&mut object, "key_id", task.key_id.as_deref());
|
||||
insert_missing_optional_string_value(&mut object, "user_id", task.user_id.as_deref());
|
||||
insert_missing_optional_string_value(&mut object, "api_key_id", task.api_key_id.as_deref());
|
||||
insert_missing_optional_string_value(
|
||||
&mut object,
|
||||
"client_api_format",
|
||||
task.client_api_format.as_deref(),
|
||||
);
|
||||
insert_missing_optional_string_value(
|
||||
&mut object,
|
||||
"provider_api_format",
|
||||
task.provider_api_format.as_deref(),
|
||||
);
|
||||
Some(Value::Object(object))
|
||||
}
|
||||
|
||||
fn insert_missing_string_value(object: &mut Map<String, Value>, key: &str, value: Option<&str>) {
|
||||
if object.contains_key(key) {
|
||||
return;
|
||||
}
|
||||
let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) else {
|
||||
return;
|
||||
};
|
||||
object.insert(key.to_string(), Value::String(value.to_string()));
|
||||
}
|
||||
|
||||
fn insert_missing_optional_string_value(
|
||||
object: &mut Map<String, Value>,
|
||||
key: &str,
|
||||
value: Option<&str>,
|
||||
) {
|
||||
insert_missing_string_value(object, key, value);
|
||||
}
|
||||
|
||||
fn has_non_empty_str(value: &Value, key: &str) -> bool {
|
||||
value
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn has_u64(value: &Value, key: &str) -> bool {
|
||||
value.get(key).and_then(Value::as_u64).is_some()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use aether_data_contracts::repository::candidates::{
|
||||
RequestCandidateStatus, StoredRequestCandidate,
|
||||
};
|
||||
use aether_data_contracts::repository::video_tasks::{StoredVideoTask, VideoTaskStatus};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
use super::{
|
||||
build_locally_actionable_report_context_from_request_candidate,
|
||||
build_locally_actionable_report_context_from_video_task,
|
||||
report_context_is_locally_actionable,
|
||||
};
|
||||
|
||||
fn sample_candidate() -> StoredRequestCandidate {
|
||||
StoredRequestCandidate {
|
||||
id: "cand-1".to_string(),
|
||||
request_id: "req-1".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("api-key-1".to_string()),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
candidate_index: 0,
|
||||
retry_index: 0,
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
endpoint_id: Some("endpoint-1".to_string()),
|
||||
key_id: Some("key-1".to_string()),
|
||||
status: RequestCandidateStatus::Pending,
|
||||
skip_reason: None,
|
||||
is_cached: false,
|
||||
status_code: None,
|
||||
error_type: None,
|
||||
error_message: None,
|
||||
latency_ms: None,
|
||||
concurrent_requests: None,
|
||||
extra_data: None,
|
||||
required_capabilities: None,
|
||||
created_at_unix_ms: 1,
|
||||
started_at_unix_ms: None,
|
||||
finished_at_unix_ms: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_video_task() -> StoredVideoTask {
|
||||
StoredVideoTask {
|
||||
id: "task-1".to_string(),
|
||||
short_id: Some("short-1".to_string()),
|
||||
request_id: "req-1".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("api-key-1".to_string()),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
external_task_id: Some("ext-1".to_string()),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
endpoint_id: Some("endpoint-1".to_string()),
|
||||
key_id: Some("key-1".to_string()),
|
||||
client_api_format: Some("openai:video".to_string()),
|
||||
provider_api_format: Some("openai:video".to_string()),
|
||||
format_converted: false,
|
||||
model: Some("sora".to_string()),
|
||||
prompt: None,
|
||||
original_request_body: None,
|
||||
duration_seconds: None,
|
||||
resolution: None,
|
||||
aspect_ratio: None,
|
||||
size: None,
|
||||
status: VideoTaskStatus::Submitted,
|
||||
progress_percent: 0,
|
||||
progress_message: None,
|
||||
retry_count: 0,
|
||||
poll_interval_seconds: 10,
|
||||
next_poll_at_unix_secs: None,
|
||||
poll_count: 0,
|
||||
max_poll_count: 360,
|
||||
created_at_unix_ms: 1,
|
||||
submitted_at_unix_secs: Some(1),
|
||||
completed_at_unix_secs: None,
|
||||
updated_at_unix_secs: 1,
|
||||
error_code: None,
|
||||
error_message: None,
|
||||
video_url: None,
|
||||
request_metadata: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn detects_locally_actionable_report_context() {
|
||||
assert!(report_context_is_locally_actionable(Some(&json!({
|
||||
"request_id": "req-1",
|
||||
"provider_id": "provider-1"
|
||||
}))));
|
||||
assert!(!report_context_is_locally_actionable(Some(&json!({
|
||||
"request_id": "req-1"
|
||||
}))));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn patches_locally_actionable_report_context_from_candidate() {
|
||||
let resolved = build_locally_actionable_report_context_from_request_candidate(
|
||||
&json!({"request_id": "req-1"}),
|
||||
&sample_candidate(),
|
||||
)
|
||||
.expect("candidate context should resolve");
|
||||
|
||||
assert_eq!(
|
||||
resolved.get("candidate_id").and_then(Value::as_str),
|
||||
Some("cand-1")
|
||||
);
|
||||
assert_eq!(
|
||||
resolved.get("provider_id").and_then(Value::as_str),
|
||||
Some("provider-1")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn patches_locally_actionable_report_context_from_video_task() {
|
||||
let resolved = build_locally_actionable_report_context_from_video_task(
|
||||
&json!({"local_task_id": "task-1"}),
|
||||
&sample_video_task(),
|
||||
)
|
||||
.expect("video task context should resolve");
|
||||
|
||||
assert_eq!(
|
||||
resolved.get("request_id").and_then(Value::as_str),
|
||||
Some("req-1")
|
||||
);
|
||||
assert_eq!(
|
||||
resolved.get("provider_api_format").and_then(Value::as_str),
|
||||
Some("openai:video")
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,967 @@
|
||||
use aether_ai_formats::api::{
|
||||
sanitize_request_path, sanitize_request_path_and_query, sanitize_request_query_string,
|
||||
};
|
||||
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
|
||||
use aether_contracts::ExecutionPlan;
|
||||
use aether_data_contracts::repository::usage::{
|
||||
extract_provider_actual_service_tier_from_response,
|
||||
extract_provider_reasoning_effort_from_body, extract_provider_service_tier_from_body,
|
||||
normalize_provider_service_tier, resolve_provider_cache_ttl_minutes,
|
||||
PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY,
|
||||
PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY,
|
||||
};
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
const MAX_USAGE_REQUEST_METADATA_DEPTH: usize = 32;
|
||||
const MAX_USAGE_REQUEST_METADATA_NODES: usize = 4_000;
|
||||
const MAX_USAGE_REQUEST_METADATA_BYTES: usize = 16 * 1024;
|
||||
const MAX_USAGE_REQUEST_METADATA_STRING_BYTES: usize = 1_024;
|
||||
|
||||
pub(crate) fn build_usage_request_metadata_seed(
|
||||
_plan: &ExecutionPlan,
|
||||
context: Option<&Map<String, Value>>,
|
||||
) -> Option<Value> {
|
||||
let mut metadata = Map::new();
|
||||
if let Some(context) = context {
|
||||
copy_allowed_metadata_fields(context, &mut metadata);
|
||||
}
|
||||
(!metadata.is_empty()).then_some(Value::Object(metadata))
|
||||
}
|
||||
|
||||
pub(crate) fn merge_usage_request_metadata(
|
||||
base: Option<Value>,
|
||||
override_value: Option<Value>,
|
||||
) -> Option<Value> {
|
||||
let mut metadata = Map::new();
|
||||
if let Some(Value::Object(base)) = base.as_ref() {
|
||||
copy_allowed_metadata_fields(base, &mut metadata);
|
||||
}
|
||||
if let Some(Value::Object(override_object)) = override_value.as_ref() {
|
||||
copy_allowed_metadata_fields(override_object, &mut metadata);
|
||||
}
|
||||
(!metadata.is_empty()).then_some(Value::Object(metadata))
|
||||
}
|
||||
|
||||
pub(crate) fn merge_usage_request_metadata_owned(
|
||||
base: Option<Value>,
|
||||
override_value: Option<Value>,
|
||||
) -> Option<Value> {
|
||||
let mut metadata = match base {
|
||||
Some(Value::Object(base)) => base,
|
||||
_ => Map::new(),
|
||||
};
|
||||
if let Some(Value::Object(override_object)) = override_value {
|
||||
move_allowed_metadata_fields(override_object, &mut metadata);
|
||||
}
|
||||
(!metadata.is_empty()).then_some(Value::Object(metadata))
|
||||
}
|
||||
|
||||
pub(crate) fn sanitize_usage_request_metadata(value: Option<Value>) -> Option<Value> {
|
||||
let Value::Object(object) = value? else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let mut filtered = Map::new();
|
||||
move_allowed_metadata_fields(object, &mut filtered);
|
||||
|
||||
(!filtered.is_empty()).then_some(Value::Object(filtered))
|
||||
}
|
||||
|
||||
pub(crate) fn sanitize_usage_request_metadata_ref(value: Option<&Value>) -> Option<Value> {
|
||||
let object = value.and_then(Value::as_object)?;
|
||||
|
||||
let mut filtered = Map::new();
|
||||
copy_allowed_metadata_fields(object, &mut filtered);
|
||||
|
||||
(!filtered.is_empty()).then_some(Value::Object(filtered))
|
||||
}
|
||||
|
||||
pub(crate) fn attach_provider_request_body_metadata(
|
||||
metadata: Option<Value>,
|
||||
provider_api_format: Option<&str>,
|
||||
provider_model: Option<&str>,
|
||||
source_model: Option<&str>,
|
||||
provider_request_body: Option<&Value>,
|
||||
) -> Option<Value> {
|
||||
let provider_body_is_object = provider_request_body.and_then(Value::as_object).is_some();
|
||||
let reasoning_effort = extract_provider_reasoning_effort_from_body(provider_request_body);
|
||||
let service_tier = extract_provider_service_tier_from_body(provider_request_body);
|
||||
let cache_ttl_minutes = resolve_provider_cache_ttl_minutes(
|
||||
provider_api_format,
|
||||
provider_model,
|
||||
source_model,
|
||||
provider_request_body,
|
||||
);
|
||||
if !provider_body_is_object
|
||||
&& reasoning_effort.is_none()
|
||||
&& service_tier.is_none()
|
||||
&& cache_ttl_minutes.is_none()
|
||||
{
|
||||
return metadata;
|
||||
}
|
||||
let mut object = match metadata {
|
||||
Some(Value::Object(object)) => object,
|
||||
_ => Map::new(),
|
||||
};
|
||||
if provider_body_is_object {
|
||||
object.remove(PROVIDER_REASONING_EFFORT_METADATA_KEY);
|
||||
object.remove(PROVIDER_SERVICE_TIER_METADATA_KEY);
|
||||
object.remove(PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY);
|
||||
}
|
||||
if let Some(reasoning_effort) = reasoning_effort {
|
||||
object.insert(
|
||||
PROVIDER_REASONING_EFFORT_METADATA_KEY.to_string(),
|
||||
Value::String(reasoning_effort),
|
||||
);
|
||||
}
|
||||
if let Some(service_tier) = service_tier {
|
||||
object.insert(
|
||||
PROVIDER_SERVICE_TIER_METADATA_KEY.to_string(),
|
||||
Value::String(service_tier),
|
||||
);
|
||||
}
|
||||
if let Some(cache_ttl_minutes) = cache_ttl_minutes {
|
||||
object.insert(
|
||||
PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY.to_string(),
|
||||
Value::Number(cache_ttl_minutes.into()),
|
||||
);
|
||||
}
|
||||
(!object.is_empty()).then_some(Value::Object(object))
|
||||
}
|
||||
|
||||
pub(crate) fn attach_provider_response_body_metadata(
|
||||
metadata: Option<Value>,
|
||||
provider_response_body: Option<&Value>,
|
||||
) -> Option<Value> {
|
||||
if metadata
|
||||
.as_ref()
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|object| object.get(PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY))
|
||||
.and_then(Value::as_str)
|
||||
.and_then(normalize_provider_service_tier)
|
||||
.is_some()
|
||||
{
|
||||
return metadata;
|
||||
}
|
||||
let actual_service_tier =
|
||||
extract_provider_actual_service_tier_from_response(provider_response_body);
|
||||
attach_provider_actual_service_tier_metadata(metadata, actual_service_tier.as_deref())
|
||||
}
|
||||
|
||||
pub(crate) fn attach_provider_actual_service_tier_metadata(
|
||||
metadata: Option<Value>,
|
||||
actual_service_tier: Option<&str>,
|
||||
) -> Option<Value> {
|
||||
let Some(actual_service_tier) = actual_service_tier.and_then(normalize_provider_service_tier)
|
||||
else {
|
||||
return metadata;
|
||||
};
|
||||
let mut object = match metadata {
|
||||
Some(Value::Object(object)) => object,
|
||||
_ => Map::new(),
|
||||
};
|
||||
object.insert(
|
||||
PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY.to_string(),
|
||||
Value::String(actual_service_tier),
|
||||
);
|
||||
(!object.is_empty()).then_some(Value::Object(object))
|
||||
}
|
||||
|
||||
fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<String, Value>) {
|
||||
copy_non_empty_string(source, target, "trace_id");
|
||||
copy_non_empty_string(source, target, "client_ip");
|
||||
copy_non_empty_string(source, target, "user_agent");
|
||||
copy_non_empty_string(source, target, "client_family");
|
||||
copy_bool(source, target, "client_requested_stream");
|
||||
copy_bool(source, target, UPSTREAM_IS_STREAM_KEY);
|
||||
copy_non_null_value(source, target, "client_session_affinity");
|
||||
copy_bool(source, target, "api_key_is_standalone");
|
||||
copy_non_empty_string(source, target, "request_path");
|
||||
copy_non_empty_string(source, target, "request_query_string");
|
||||
copy_non_empty_string(source, target, "request_path_and_query");
|
||||
copy_non_empty_string(source, target, PROVIDER_REASONING_EFFORT_METADATA_KEY);
|
||||
copy_non_empty_string(source, target, PROVIDER_SERVICE_TIER_METADATA_KEY);
|
||||
copy_non_empty_string(source, target, PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY);
|
||||
copy_number(source, target, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY);
|
||||
copy_number(source, target, "provider_request_body_base64_bytes");
|
||||
copy_number(source, target, "provider_response_body_base64_bytes");
|
||||
copy_number(source, target, "client_response_body_base64_bytes");
|
||||
copy_non_null_value(source, target, "body_size");
|
||||
copy_number(source, target, "client_response_status_code");
|
||||
copy_non_null_value(source, target, "billing_snapshot");
|
||||
copy_non_empty_string(source, target, "billing_snapshot_schema_version");
|
||||
copy_non_empty_string(source, target, "billing_snapshot_status");
|
||||
copy_non_null_value(source, target, "settlement_snapshot");
|
||||
copy_non_empty_string(source, target, "settlement_snapshot_schema_version");
|
||||
copy_non_null_value(source, target, "billing_dimensions");
|
||||
copy_non_empty_string(source, target, "model_id");
|
||||
copy_non_empty_string(source, target, "global_model_id");
|
||||
copy_non_empty_string(source, target, "global_model_name");
|
||||
copy_non_null_value(source, target, "dimensions");
|
||||
copy_non_null_value(source, target, "billing_rule_snapshot");
|
||||
copy_non_null_value(source, target, "scheduling_audit");
|
||||
copy_non_null_value(source, target, "tls_fingerprint");
|
||||
copy_number(source, target, "rate_multiplier");
|
||||
copy_bool(source, target, "is_free_tier");
|
||||
copy_number(source, target, "input_price_per_1m");
|
||||
copy_number(source, target, "output_price_per_1m");
|
||||
copy_number(source, target, "cache_creation_price_per_1m");
|
||||
copy_number(source, target, "cache_read_price_per_1m");
|
||||
copy_number(source, target, "price_per_request");
|
||||
copy_non_null_value(source, target, "proxy");
|
||||
copy_non_null_value(source, target, "stage_timings_ms");
|
||||
copy_non_null_value(source, target, "db_timings_ms");
|
||||
sanitize_request_path_metadata_fields(target);
|
||||
}
|
||||
|
||||
fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map<String, Value>) {
|
||||
remove_non_empty_string(&mut source, target, "trace_id");
|
||||
remove_non_empty_string(&mut source, target, "client_ip");
|
||||
remove_non_empty_string(&mut source, target, "user_agent");
|
||||
remove_non_empty_string(&mut source, target, "client_family");
|
||||
remove_bool(&mut source, target, "client_requested_stream");
|
||||
remove_bool(&mut source, target, UPSTREAM_IS_STREAM_KEY);
|
||||
remove_non_null_value(&mut source, target, "client_session_affinity");
|
||||
remove_bool(&mut source, target, "api_key_is_standalone");
|
||||
remove_non_empty_string(&mut source, target, "request_path");
|
||||
remove_non_empty_string(&mut source, target, "request_query_string");
|
||||
remove_non_empty_string(&mut source, target, "request_path_and_query");
|
||||
remove_non_empty_string(&mut source, target, PROVIDER_REASONING_EFFORT_METADATA_KEY);
|
||||
remove_non_empty_string(&mut source, target, PROVIDER_SERVICE_TIER_METADATA_KEY);
|
||||
remove_non_empty_string(
|
||||
&mut source,
|
||||
target,
|
||||
PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY,
|
||||
);
|
||||
remove_number(&mut source, target, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY);
|
||||
remove_number(&mut source, target, "provider_request_body_base64_bytes");
|
||||
remove_number(&mut source, target, "provider_response_body_base64_bytes");
|
||||
remove_number(&mut source, target, "client_response_body_base64_bytes");
|
||||
remove_non_null_value(&mut source, target, "body_size");
|
||||
remove_number(&mut source, target, "client_response_status_code");
|
||||
remove_non_null_value(&mut source, target, "billing_snapshot");
|
||||
remove_non_empty_string(&mut source, target, "billing_snapshot_schema_version");
|
||||
remove_non_empty_string(&mut source, target, "billing_snapshot_status");
|
||||
remove_non_null_value(&mut source, target, "settlement_snapshot");
|
||||
remove_non_empty_string(&mut source, target, "settlement_snapshot_schema_version");
|
||||
remove_non_null_value(&mut source, target, "billing_dimensions");
|
||||
remove_non_empty_string(&mut source, target, "model_id");
|
||||
remove_non_empty_string(&mut source, target, "global_model_id");
|
||||
remove_non_empty_string(&mut source, target, "global_model_name");
|
||||
remove_non_null_value(&mut source, target, "dimensions");
|
||||
remove_non_null_value(&mut source, target, "billing_rule_snapshot");
|
||||
remove_non_null_value(&mut source, target, "scheduling_audit");
|
||||
remove_non_null_value(&mut source, target, "tls_fingerprint");
|
||||
remove_number(&mut source, target, "rate_multiplier");
|
||||
remove_bool(&mut source, target, "is_free_tier");
|
||||
remove_number(&mut source, target, "input_price_per_1m");
|
||||
remove_number(&mut source, target, "output_price_per_1m");
|
||||
remove_number(&mut source, target, "cache_creation_price_per_1m");
|
||||
remove_number(&mut source, target, "cache_read_price_per_1m");
|
||||
remove_number(&mut source, target, "price_per_request");
|
||||
remove_non_null_value(&mut source, target, "proxy");
|
||||
remove_non_null_value(&mut source, target, "stage_timings_ms");
|
||||
remove_non_null_value(&mut source, target, "db_timings_ms");
|
||||
sanitize_request_path_metadata_fields(target);
|
||||
}
|
||||
|
||||
fn sanitize_request_path_metadata_fields(target: &mut Map<String, Value>) {
|
||||
let path = target
|
||||
.get("request_path")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(sanitize_request_path);
|
||||
let query = target
|
||||
.get("request_query_string")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(sanitize_request_query_string);
|
||||
let path_and_query = target
|
||||
.get("request_path_and_query")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(|value| sanitize_request_path_and_query(value, None))
|
||||
.or_else(|| {
|
||||
path.as_deref()
|
||||
.and_then(|path| sanitize_request_path_and_query(path, query.as_deref()))
|
||||
});
|
||||
|
||||
apply_optional_string_field(target, "request_path", path.as_deref());
|
||||
apply_optional_string_field(target, "request_query_string", query.as_deref());
|
||||
apply_optional_string_field(target, "request_path_and_query", path_and_query.as_deref());
|
||||
}
|
||||
|
||||
fn apply_optional_string_field(target: &mut Map<String, Value>, key: &str, value: Option<&str>) {
|
||||
if let Some(value) = value {
|
||||
target.insert(key.to_string(), Value::String(value.to_string()));
|
||||
} else {
|
||||
target.remove(key);
|
||||
}
|
||||
}
|
||||
|
||||
fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
|
||||
let Some(value) = source
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
else {
|
||||
return;
|
||||
};
|
||||
target.insert(
|
||||
key.to_string(),
|
||||
Value::String(truncate_usage_request_metadata_string(value)),
|
||||
);
|
||||
}
|
||||
|
||||
fn remove_non_empty_string(
|
||||
source: &mut Map<String, Value>,
|
||||
target: &mut Map<String, Value>,
|
||||
key: &str,
|
||||
) {
|
||||
let Some(Value::String(value)) = source.remove(key) else {
|
||||
return;
|
||||
};
|
||||
let Some(value) = trim_and_truncate_usage_request_metadata_string_owned(value) else {
|
||||
return;
|
||||
};
|
||||
target.insert(key.to_string(), Value::String(value));
|
||||
}
|
||||
|
||||
fn copy_number(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
|
||||
let Some(value) = source.get(key).filter(|value| value.is_number()) else {
|
||||
return;
|
||||
};
|
||||
target.insert(key.to_string(), value.clone());
|
||||
}
|
||||
|
||||
fn remove_number(source: &mut Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
|
||||
let Some(value) = source.remove(key).filter(|value| value.is_number()) else {
|
||||
return;
|
||||
};
|
||||
target.insert(key.to_string(), value);
|
||||
}
|
||||
|
||||
fn copy_bool(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
|
||||
let Some(value) = source.get(key).filter(|value| value.is_boolean()) else {
|
||||
return;
|
||||
};
|
||||
target.insert(key.to_string(), value.clone());
|
||||
}
|
||||
|
||||
fn remove_bool(source: &mut Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
|
||||
let Some(value) = source.remove(key).filter(|value| value.is_boolean()) else {
|
||||
return;
|
||||
};
|
||||
target.insert(key.to_string(), value);
|
||||
}
|
||||
|
||||
fn copy_non_null_value(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
|
||||
let Some(value) = source.get(key).filter(|value| !value.is_null()) else {
|
||||
return;
|
||||
};
|
||||
target.insert(
|
||||
key.to_string(),
|
||||
sanitize_usage_request_metadata_value(value),
|
||||
);
|
||||
}
|
||||
|
||||
fn remove_non_null_value(
|
||||
source: &mut Map<String, Value>,
|
||||
target: &mut Map<String, Value>,
|
||||
key: &str,
|
||||
) {
|
||||
let Some(value) = source.remove(key).filter(|value| !value.is_null()) else {
|
||||
return;
|
||||
};
|
||||
target.insert(
|
||||
key.to_string(),
|
||||
sanitize_usage_request_metadata_value_owned(value),
|
||||
);
|
||||
}
|
||||
|
||||
fn sanitize_usage_request_metadata_value(value: &Value) -> Value {
|
||||
match value {
|
||||
Value::String(text) => Value::String(truncate_usage_request_metadata_string(text)),
|
||||
_ if usage_request_metadata_within_limits(value) => value.clone(),
|
||||
_ => truncated_usage_request_metadata_value(value),
|
||||
}
|
||||
}
|
||||
|
||||
fn sanitize_usage_request_metadata_value_owned(value: Value) -> Value {
|
||||
match value {
|
||||
Value::String(text) => Value::String(truncate_usage_request_metadata_string_owned(text)),
|
||||
_ if usage_request_metadata_within_limits(&value) => value,
|
||||
_ => truncated_usage_request_metadata_value(&value),
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate_usage_request_metadata_string(value: &str) -> String {
|
||||
const TRUNCATED_SUFFIX: &str = "...[truncated]";
|
||||
|
||||
if value.len() <= MAX_USAGE_REQUEST_METADATA_STRING_BYTES {
|
||||
return value.to_string();
|
||||
}
|
||||
|
||||
let target_bytes =
|
||||
MAX_USAGE_REQUEST_METADATA_STRING_BYTES.saturating_sub(TRUNCATED_SUFFIX.len());
|
||||
let mut end = 0usize;
|
||||
for (idx, ch) in value.char_indices() {
|
||||
let next = idx + ch.len_utf8();
|
||||
if next > target_bytes {
|
||||
break;
|
||||
}
|
||||
end = next;
|
||||
}
|
||||
|
||||
if end == 0 {
|
||||
return TRUNCATED_SUFFIX.to_string();
|
||||
}
|
||||
|
||||
format!("{}{TRUNCATED_SUFFIX}", &value[..end])
|
||||
}
|
||||
|
||||
fn trim_and_truncate_usage_request_metadata_string_owned(value: String) -> Option<String> {
|
||||
let trimmed = value.trim();
|
||||
if trimmed.is_empty() {
|
||||
return None;
|
||||
}
|
||||
if trimmed.len() == value.len() {
|
||||
return Some(truncate_usage_request_metadata_string_owned(value));
|
||||
}
|
||||
Some(truncate_usage_request_metadata_string(trimmed))
|
||||
}
|
||||
|
||||
fn truncate_usage_request_metadata_string_owned(value: String) -> String {
|
||||
if value.len() <= MAX_USAGE_REQUEST_METADATA_STRING_BYTES {
|
||||
return value;
|
||||
}
|
||||
truncate_usage_request_metadata_string(value.as_str())
|
||||
}
|
||||
|
||||
fn truncated_usage_request_metadata_value(value: &Value) -> Value {
|
||||
json!({
|
||||
"truncated": true,
|
||||
"reason": "usage_request_metadata_limits_exceeded",
|
||||
"max_depth": MAX_USAGE_REQUEST_METADATA_DEPTH,
|
||||
"max_nodes": MAX_USAGE_REQUEST_METADATA_NODES,
|
||||
"max_bytes": MAX_USAGE_REQUEST_METADATA_BYTES,
|
||||
"value_kind": usage_request_metadata_value_kind(value),
|
||||
})
|
||||
}
|
||||
|
||||
fn usage_request_metadata_within_limits(value: &Value) -> bool {
|
||||
let mut nodes = 0usize;
|
||||
let mut estimated_bytes = 0usize;
|
||||
let mut stack = vec![(value, 1usize)];
|
||||
|
||||
while let Some((current, depth)) = stack.pop() {
|
||||
nodes = nodes.saturating_add(1);
|
||||
estimated_bytes =
|
||||
estimated_bytes.saturating_add(usage_request_metadata_value_size_hint(current));
|
||||
if depth > MAX_USAGE_REQUEST_METADATA_DEPTH
|
||||
|| nodes > MAX_USAGE_REQUEST_METADATA_NODES
|
||||
|| estimated_bytes > MAX_USAGE_REQUEST_METADATA_BYTES
|
||||
{
|
||||
return false;
|
||||
}
|
||||
match current {
|
||||
Value::Array(items) => {
|
||||
estimated_bytes = estimated_bytes.saturating_add(items.len().saturating_mul(2));
|
||||
for item in items.iter().rev() {
|
||||
stack.push((item, depth + 1));
|
||||
}
|
||||
}
|
||||
Value::Object(object) => {
|
||||
estimated_bytes = estimated_bytes
|
||||
.saturating_add(object.len().saturating_mul(3))
|
||||
.saturating_add(
|
||||
object
|
||||
.keys()
|
||||
.map(|key| key.len().saturating_add(2))
|
||||
.sum::<usize>(),
|
||||
);
|
||||
for item in object.values() {
|
||||
stack.push((item, depth + 1));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
fn usage_request_metadata_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",
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_request_metadata_value_size_hint(value: &Value) -> usize {
|
||||
match value {
|
||||
Value::Null => 4,
|
||||
Value::Bool(false) => 5,
|
||||
Value::Bool(true) => 4,
|
||||
Value::Number(number) => number.to_string().len(),
|
||||
Value::String(text) => text.len().saturating_add(2),
|
||||
Value::Array(_) | Value::Object(_) => 2,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use aether_contracts::{ExecutionPlan, RequestBody};
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{
|
||||
attach_provider_actual_service_tier_metadata, attach_provider_request_body_metadata,
|
||||
attach_provider_response_body_metadata, build_usage_request_metadata_seed,
|
||||
merge_usage_request_metadata, merge_usage_request_metadata_owned,
|
||||
sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref,
|
||||
MAX_USAGE_REQUEST_METADATA_BYTES, MAX_USAGE_REQUEST_METADATA_DEPTH,
|
||||
MAX_USAGE_REQUEST_METADATA_NODES,
|
||||
};
|
||||
|
||||
fn sample_plan() -> ExecutionPlan {
|
||||
ExecutionPlan {
|
||||
request_id: "req-1".to_string(),
|
||||
candidate_id: Some("cand-1".to_string()),
|
||||
provider_name: Some("OpenAI".to_string()),
|
||||
provider_id: "provider-1".to_string(),
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
method: "POST".to_string(),
|
||||
url: "https://example.com/v1/chat/completions".to_string(),
|
||||
headers: BTreeMap::new(),
|
||||
content_type: None,
|
||||
content_encoding: None,
|
||||
body: RequestBody::from_json(json!({"model": "gpt-5"})),
|
||||
stream: false,
|
||||
client_api_format: "openai:chat".to_string(),
|
||||
provider_api_format: "openai:chat".to_string(),
|
||||
model_name: Some("gpt-5".to_string()),
|
||||
proxy: None,
|
||||
transport_profile: None,
|
||||
timeouts: None,
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_stage_timings_metadata() -> Value {
|
||||
json!({
|
||||
"stream_candidate_slot": 1,
|
||||
"stream_provider_in_flight": 2,
|
||||
"stream_upstream_headers": 180,
|
||||
"stream_first_data": 8210
|
||||
})
|
||||
}
|
||||
|
||||
fn sample_db_timings_metadata() -> Value {
|
||||
json!({
|
||||
"query_count": 2,
|
||||
"query_total": 950,
|
||||
"query_max": 650,
|
||||
"operations": {
|
||||
"request_candidate_upsert": {"count": 1, "sum": 650, "max": 650},
|
||||
"usage_upsert": {"count": 1, "sum": 300, "max": 300}
|
||||
},
|
||||
"pool": {
|
||||
"max_checked_out": 20,
|
||||
"max_pool_size": 20,
|
||||
"min_idle": 0,
|
||||
"max_connections": 20,
|
||||
"max_usage_rate": 100.0
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitizes_request_metadata_to_allowlist() {
|
||||
let stage_timings_ms = sample_stage_timings_metadata();
|
||||
let db_timings_ms = sample_db_timings_metadata();
|
||||
let metadata = sanitize_usage_request_metadata(Some(json!({
|
||||
"request_id": "req-1",
|
||||
"provider_id": "provider-1",
|
||||
"provider_name": "OpenAI",
|
||||
"model": "gpt-5",
|
||||
"candidate_index": 2,
|
||||
"trace_id": "trace-1",
|
||||
"client_ip": "203.0.113.8",
|
||||
"user_agent": "Claude-Code/1.0",
|
||||
"client_requested_stream": false,
|
||||
"upstream_is_stream": true,
|
||||
"api_key_is_standalone": true,
|
||||
"provider_request_body_base64_bytes": 512,
|
||||
"provider_response_body_base64_bytes": 1024,
|
||||
"client_response_body_base64_bytes": 2048,
|
||||
"body_size": {
|
||||
"client_request_body": "1 KB",
|
||||
"provider_request_body": "4 KB",
|
||||
"provider_over_client": "4x"
|
||||
},
|
||||
"billing_snapshot": {"status": "complete"},
|
||||
"billing_snapshot_schema_version": "2.0",
|
||||
"billing_snapshot_status": "complete",
|
||||
"model_id": "model-1",
|
||||
"global_model_id": "global-model-1",
|
||||
"global_model_name": "gpt-5",
|
||||
"dimensions": {"total_input_context": 10},
|
||||
"rate_multiplier": 1.25,
|
||||
"is_free_tier": false,
|
||||
"input_price_per_1m": 3.0,
|
||||
"output_price_per_1m": 15.0,
|
||||
"cache_creation_price_per_1m": 3.75,
|
||||
"cache_read_price_per_1m": 0.3,
|
||||
"price_per_request": 0.02,
|
||||
"stage_timings_ms": stage_timings_ms.clone(),
|
||||
"db_timings_ms": db_timings_ms.clone(),
|
||||
"original_headers": {"authorization": "Bearer secret"},
|
||||
"original_request_body": {"messages": []},
|
||||
"provider_request_headers": {"authorization": "Bearer secret"},
|
||||
"upstream_url": "https://example.com/v1/chat/completions"
|
||||
})))
|
||||
.expect("metadata should remain");
|
||||
|
||||
assert_eq!(
|
||||
metadata,
|
||||
json!({
|
||||
"trace_id": "trace-1",
|
||||
"client_ip": "203.0.113.8",
|
||||
"user_agent": "Claude-Code/1.0",
|
||||
"client_requested_stream": false,
|
||||
"upstream_is_stream": true,
|
||||
"api_key_is_standalone": true,
|
||||
"provider_request_body_base64_bytes": 512,
|
||||
"provider_response_body_base64_bytes": 1024,
|
||||
"client_response_body_base64_bytes": 2048,
|
||||
"body_size": {
|
||||
"client_request_body": "1 KB",
|
||||
"provider_request_body": "4 KB",
|
||||
"provider_over_client": "4x"
|
||||
},
|
||||
"billing_snapshot": {"status": "complete"},
|
||||
"billing_snapshot_schema_version": "2.0",
|
||||
"billing_snapshot_status": "complete",
|
||||
"model_id": "model-1",
|
||||
"global_model_id": "global-model-1",
|
||||
"global_model_name": "gpt-5",
|
||||
"dimensions": {"total_input_context": 10},
|
||||
"rate_multiplier": 1.25,
|
||||
"is_free_tier": false,
|
||||
"input_price_per_1m": 3.0,
|
||||
"output_price_per_1m": 15.0,
|
||||
"cache_creation_price_per_1m": 3.75,
|
||||
"cache_read_price_per_1m": 0.3,
|
||||
"price_per_request": 0.02,
|
||||
"stage_timings_ms": stage_timings_ms,
|
||||
"db_timings_ms": db_timings_ms
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitizes_request_path_query_metadata() {
|
||||
let metadata = sanitize_usage_request_metadata(Some(json!({
|
||||
"request_path": "/v1beta/models/gemini-2.5-pro:streamGenerateContent?key=secret",
|
||||
"request_query_string": "key=secret&alt=sse&pageSize=10&token=hidden",
|
||||
"request_path_and_query": "/v1beta/models/gemini-2.5-pro:streamGenerateContent?key=secret&alt=sse&pageSize=10&token=hidden",
|
||||
})))
|
||||
.expect("metadata should remain");
|
||||
|
||||
assert_eq!(
|
||||
metadata,
|
||||
json!({
|
||||
"request_path": "/v1beta/models/gemini-2.5-pro:streamGenerateContent",
|
||||
"request_query_string": "alt=sse&pageSize=10",
|
||||
"request_path_and_query": "/v1beta/models/gemini-2.5-pro:streamGenerateContent?alt=sse&pageSize=10",
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitizes_large_allowed_metadata_values_to_bounded_representations() {
|
||||
let metadata = sanitize_usage_request_metadata(Some(json!({
|
||||
"trace_id": "t".repeat(2_048),
|
||||
"billing_snapshot": {
|
||||
"payload": "x".repeat(32 * 1024)
|
||||
}
|
||||
})))
|
||||
.expect("metadata should remain");
|
||||
|
||||
assert!(metadata
|
||||
.get("trace_id")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| value.ends_with("...[truncated]")));
|
||||
assert_eq!(
|
||||
metadata.get("billing_snapshot"),
|
||||
Some(&json!({
|
||||
"truncated": true,
|
||||
"reason": "usage_request_metadata_limits_exceeded",
|
||||
"max_depth": MAX_USAGE_REQUEST_METADATA_DEPTH,
|
||||
"max_nodes": MAX_USAGE_REQUEST_METADATA_NODES,
|
||||
"max_bytes": MAX_USAGE_REQUEST_METADATA_BYTES,
|
||||
"value_kind": "object",
|
||||
}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sanitizes_request_metadata_preserves_tls_fingerprint() {
|
||||
let metadata = sanitize_usage_request_metadata(Some(json!({
|
||||
"tls_fingerprint": {
|
||||
"incoming": {
|
||||
"source": "forwarded_header",
|
||||
"ja3": "incoming-ja3",
|
||||
"ja4": "incoming-ja4"
|
||||
},
|
||||
"outgoing": {
|
||||
"source": "aether_transport_config",
|
||||
"backend": "reqwest_rustls",
|
||||
"observed": false
|
||||
}
|
||||
},
|
||||
"untrusted_tls_fingerprint": {
|
||||
"ja3": "spoofed"
|
||||
}
|
||||
})))
|
||||
.expect("metadata should remain");
|
||||
|
||||
assert_eq!(
|
||||
metadata,
|
||||
json!({
|
||||
"tls_fingerprint": {
|
||||
"incoming": {
|
||||
"source": "forwarded_header",
|
||||
"ja3": "incoming-ja3",
|
||||
"ja4": "incoming-ja4"
|
||||
},
|
||||
"outgoing": {
|
||||
"source": "aether_transport_config",
|
||||
"backend": "reqwest_rustls",
|
||||
"observed": false
|
||||
}
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_seed_from_context_and_allowlisted_metadata_only() {
|
||||
let metadata = build_usage_request_metadata_seed(
|
||||
&sample_plan(),
|
||||
Some(
|
||||
json!({
|
||||
"request_id": "req-1",
|
||||
"candidate_index": 0,
|
||||
"client_requested_stream": false,
|
||||
"upstream_is_stream": true,
|
||||
"api_key_is_standalone": true,
|
||||
"provider_id": "provider-1",
|
||||
"model_id": "model-1",
|
||||
"global_model_id": "global-model-1",
|
||||
"global_model_name": "gpt-5",
|
||||
"client_ip": "203.0.113.8",
|
||||
"user_agent": "Claude-Code/1.0",
|
||||
"billing_snapshot": {"status": "complete"},
|
||||
"stage_timings_ms": {
|
||||
"stream_candidate_slot": 0,
|
||||
"stream_upstream_headers": 180,
|
||||
"stream_first_data": 8210
|
||||
},
|
||||
"db_timings_ms": {
|
||||
"query_count": 1,
|
||||
"query_total": 42,
|
||||
"query_max": 42,
|
||||
"operations": {
|
||||
"auth_api_key_snapshot": {"count": 1, "sum": 42, "max": 42}
|
||||
}
|
||||
}
|
||||
})
|
||||
.as_object()
|
||||
.expect("object"),
|
||||
),
|
||||
)
|
||||
.expect("metadata should remain");
|
||||
|
||||
assert_eq!(
|
||||
metadata,
|
||||
json!({
|
||||
"client_requested_stream": false,
|
||||
"upstream_is_stream": true,
|
||||
"api_key_is_standalone": true,
|
||||
"model_id": "model-1",
|
||||
"global_model_id": "global-model-1",
|
||||
"global_model_name": "gpt-5",
|
||||
"client_ip": "203.0.113.8",
|
||||
"user_agent": "Claude-Code/1.0",
|
||||
"billing_snapshot": {"status": "complete"},
|
||||
"stage_timings_ms": {
|
||||
"stream_candidate_slot": 0,
|
||||
"stream_upstream_headers": 180,
|
||||
"stream_first_data": 8210
|
||||
},
|
||||
"db_timings_ms": {
|
||||
"query_count": 1,
|
||||
"query_total": 42,
|
||||
"query_max": 42,
|
||||
"operations": {
|
||||
"auth_api_key_snapshot": {"count": 1, "sum": 42, "max": 42}
|
||||
}
|
||||
}
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merges_and_filters_request_metadata() {
|
||||
let metadata = merge_usage_request_metadata(
|
||||
Some(json!({
|
||||
"request_id": "req-1"
|
||||
})),
|
||||
Some(json!({
|
||||
"candidate_index": 0,
|
||||
"provider_name": "OpenAI"
|
||||
})),
|
||||
);
|
||||
|
||||
assert_eq!(metadata, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_request_body_metadata_uses_final_provider_body_as_source_of_truth() {
|
||||
let metadata = Some(json!({
|
||||
"trace_id": "trace-1",
|
||||
"provider_reasoning_effort": "high",
|
||||
"provider_service_tier": "priority"
|
||||
}));
|
||||
|
||||
let updated = attach_provider_request_body_metadata(
|
||||
metadata.clone(),
|
||||
Some("openai:responses"),
|
||||
Some("gpt-5.6-sol"),
|
||||
Some("gpt-5.6-sol"),
|
||||
Some(&json!({
|
||||
"model": "gpt-5.6-sol",
|
||||
"reasoning": { "effort": "low" },
|
||||
"service_tier": "standard"
|
||||
})),
|
||||
)
|
||||
.expect("metadata should remain");
|
||||
|
||||
assert_eq!(
|
||||
updated,
|
||||
json!({
|
||||
"trace_id": "trace-1",
|
||||
"provider_reasoning_effort": "low",
|
||||
"provider_service_tier": "standard",
|
||||
"provider_cache_ttl_minutes": 30
|
||||
})
|
||||
);
|
||||
|
||||
let cleared = attach_provider_request_body_metadata(
|
||||
metadata,
|
||||
Some("openai:responses"),
|
||||
Some("gpt-5"),
|
||||
Some("gpt-5"),
|
||||
Some(&json!({
|
||||
"model": "gpt-5"
|
||||
})),
|
||||
)
|
||||
.expect("metadata should retain unrelated fields");
|
||||
|
||||
assert_eq!(
|
||||
cleared,
|
||||
json!({
|
||||
"trace_id": "trace-1"
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_response_metadata_preserves_terminal_actual_service_tier() {
|
||||
let metadata = attach_provider_response_body_metadata(
|
||||
Some(json!({"provider_service_tier": "priority"})),
|
||||
Some(&json!({
|
||||
"chunks": [
|
||||
{"service_tier": "priority"},
|
||||
{"service_tier": "Default", "usage": {"total_tokens": 12}}
|
||||
]
|
||||
})),
|
||||
)
|
||||
.expect("requested and actual provider tiers should be preserved");
|
||||
|
||||
assert_eq!(
|
||||
metadata,
|
||||
json!({
|
||||
"provider_service_tier": "priority",
|
||||
"provider_actual_service_tier": "default"
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminal_summary_tier_uses_the_same_normalized_metadata_field() {
|
||||
let metadata = attach_provider_actual_service_tier_metadata(
|
||||
Some(json!({"trace_id": "trace-1"})),
|
||||
Some(" Flex "),
|
||||
)
|
||||
.expect("terminal summary tier should be retained");
|
||||
|
||||
assert_eq!(
|
||||
metadata,
|
||||
json!({
|
||||
"trace_id": "trace-1",
|
||||
"provider_actual_service_tier": "flex"
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn terminal_summary_tier_precedes_truncated_response_capture() {
|
||||
let metadata = attach_provider_response_body_metadata(
|
||||
Some(json!({"provider_actual_service_tier": "default"})),
|
||||
Some(&json!({"chunks": [{"service_tier": "priority"}]})),
|
||||
)
|
||||
.expect("terminal summary tier should remain");
|
||||
|
||||
assert_eq!(
|
||||
metadata.get("provider_actual_service_tier"),
|
||||
Some(&Value::String("default".to_string()))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn owned_merge_matches_filtered_merge_for_trusted_objects() {
|
||||
let base = Some(json!({
|
||||
"trace_id": "trace-1",
|
||||
"provider_request_body_base64_bytes": 128
|
||||
}));
|
||||
let override_value = Some(json!({
|
||||
"billing_snapshot_status": "complete",
|
||||
"trace_id": "trace-2",
|
||||
"provider_actual_service_tier": "default"
|
||||
}));
|
||||
|
||||
assert_eq!(
|
||||
merge_usage_request_metadata_owned(base.clone(), override_value.clone()),
|
||||
merge_usage_request_metadata(base, override_value)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn borrowed_sanitize_matches_owned_sanitize() {
|
||||
let value = json!({
|
||||
"trace_id": "trace-1",
|
||||
"billing_snapshot": {"status": "complete"},
|
||||
"provider_name": "OpenAI"
|
||||
});
|
||||
|
||||
assert_eq!(
|
||||
sanitize_usage_request_metadata_ref(Some(&value)),
|
||||
sanitize_usage_request_metadata(Some(value))
|
||||
);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,325 @@
|
||||
use std::sync::{Arc, OnceLock};
|
||||
|
||||
use aether_data_contracts::repository::settlement::{StoredUsageSettlement, UsageSettlementInput};
|
||||
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
||||
use aether_data_contracts::{DataLayerError, DataLayerError::InvalidInput};
|
||||
use async_trait::async_trait;
|
||||
|
||||
use crate::keyed_lock::KeyedAsyncLockPool;
|
||||
|
||||
#[async_trait]
|
||||
pub trait UsageSettlementWriter: Send + Sync {
|
||||
fn has_usage_settlement_writer(&self) -> bool;
|
||||
|
||||
async fn settle_usage(
|
||||
&self,
|
||||
input: UsageSettlementInput,
|
||||
) -> Result<Option<StoredUsageSettlement>, DataLayerError>;
|
||||
}
|
||||
|
||||
pub async fn settle_usage_if_needed(
|
||||
writer: &dyn UsageSettlementWriter,
|
||||
usage: &StoredRequestUsageAudit,
|
||||
) -> Result<(), DataLayerError> {
|
||||
if !writer.has_usage_settlement_writer() || usage.billing_status != "pending" {
|
||||
return Ok(());
|
||||
}
|
||||
if !matches!(usage.status.as_str(), "completed" | "failed") {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let finalized_at_unix_secs = usage
|
||||
.finalized_at_unix_secs
|
||||
.or(Some(usage.updated_at_unix_secs));
|
||||
let input = UsageSettlementInput {
|
||||
request_id: usage.request_id.clone(),
|
||||
user_id: usage.user_id.clone(),
|
||||
api_key_id: usage.api_key_id.clone(),
|
||||
api_key_is_standalone: usage_api_key_is_standalone(usage),
|
||||
provider_id: usage.provider_id.clone(),
|
||||
status: usage.status.clone(),
|
||||
billing_status: usage.billing_status.clone(),
|
||||
total_cost_usd: finite_cost(usage.total_cost_usd)?,
|
||||
actual_total_cost_usd: finite_cost(usage.actual_total_cost_usd)?,
|
||||
finalized_at_unix_secs,
|
||||
};
|
||||
let settlement_key = usage_settlement_lock_key(&input);
|
||||
let settlement_lock = usage_settlement_lock(&settlement_key);
|
||||
let _guard = settlement_lock.lock().await;
|
||||
let _ = writer.settle_usage(input).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn usage_settlement_lock(key: &str) -> Arc<tokio::sync::Mutex<()>> {
|
||||
static LOCKS: OnceLock<KeyedAsyncLockPool> = OnceLock::new();
|
||||
LOCKS.get_or_init(KeyedAsyncLockPool::default).lock_for(key)
|
||||
}
|
||||
|
||||
fn usage_settlement_lock_key(input: &UsageSettlementInput) -> String {
|
||||
if input.api_key_is_standalone {
|
||||
if let Some(api_key_id) = input.api_key_id.as_deref().and_then(non_empty_trimmed) {
|
||||
return format!("api-key:{api_key_id}");
|
||||
}
|
||||
}
|
||||
if let Some(user_id) = input.user_id.as_deref().and_then(non_empty_trimmed) {
|
||||
return format!("user:{user_id}");
|
||||
}
|
||||
if let Some(api_key_id) = input.api_key_id.as_deref().and_then(non_empty_trimmed) {
|
||||
return format!("api-key:{api_key_id}");
|
||||
}
|
||||
format!("request:{}", input.request_id.trim())
|
||||
}
|
||||
|
||||
fn non_empty_trimmed(value: &str) -> Option<&str> {
|
||||
let value = value.trim();
|
||||
(!value.is_empty()).then_some(value)
|
||||
}
|
||||
|
||||
fn usage_api_key_is_standalone(usage: &StoredRequestUsageAudit) -> bool {
|
||||
usage
|
||||
.request_metadata
|
||||
.as_ref()
|
||||
.and_then(|metadata| metadata.get("api_key_is_standalone"))
|
||||
.and_then(serde_json::Value::as_bool)
|
||||
.unwrap_or(false)
|
||||
}
|
||||
|
||||
fn finite_cost(value: f64) -> Result<f64, DataLayerError> {
|
||||
if value.is_finite() {
|
||||
Ok(value)
|
||||
} else {
|
||||
Err(InvalidInput(
|
||||
"wallet settlement cost must be finite".to_string(),
|
||||
))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::sync::Mutex;
|
||||
use std::time::Duration;
|
||||
|
||||
use super::{settle_usage_if_needed, UsageSettlementWriter};
|
||||
use aether_data_contracts::repository::settlement::UsageSettlementInput;
|
||||
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
||||
use async_trait::async_trait;
|
||||
use serde_json::json;
|
||||
|
||||
#[derive(Default)]
|
||||
struct TestSettlementWriter {
|
||||
has_writer: bool,
|
||||
inputs: Mutex<Vec<UsageSettlementInput>>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct SlowSettlementWriter {
|
||||
active: AtomicUsize,
|
||||
max_active: AtomicUsize,
|
||||
inputs: Mutex<Vec<UsageSettlementInput>>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageSettlementWriter for TestSettlementWriter {
|
||||
fn has_usage_settlement_writer(&self) -> bool {
|
||||
self.has_writer
|
||||
}
|
||||
|
||||
async fn settle_usage(
|
||||
&self,
|
||||
input: UsageSettlementInput,
|
||||
) -> Result<
|
||||
Option<aether_data_contracts::repository::settlement::StoredUsageSettlement>,
|
||||
aether_data_contracts::DataLayerError,
|
||||
> {
|
||||
self.inputs
|
||||
.lock()
|
||||
.expect("settlement inputs lock")
|
||||
.push(input);
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageSettlementWriter for SlowSettlementWriter {
|
||||
fn has_usage_settlement_writer(&self) -> bool {
|
||||
true
|
||||
}
|
||||
|
||||
async fn settle_usage(
|
||||
&self,
|
||||
input: UsageSettlementInput,
|
||||
) -> Result<
|
||||
Option<aether_data_contracts::repository::settlement::StoredUsageSettlement>,
|
||||
aether_data_contracts::DataLayerError,
|
||||
> {
|
||||
let active = self.active.fetch_add(1, Ordering::AcqRel) + 1;
|
||||
self.max_active.fetch_max(active, Ordering::AcqRel);
|
||||
tokio::time::sleep(Duration::from_millis(30)).await;
|
||||
self.inputs
|
||||
.lock()
|
||||
.expect("settlement inputs lock")
|
||||
.push(input);
|
||||
self.active.fetch_sub(1, Ordering::AcqRel);
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
fn sample_usage() -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
None,
|
||||
None,
|
||||
"openai".to_string(),
|
||||
"gpt-5".to_string(),
|
||||
None,
|
||||
Some("provider-1".to_string()),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
false,
|
||||
false,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
1.25,
|
||||
0.75,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
"completed".to_string(),
|
||||
"pending".to_string(),
|
||||
100,
|
||||
200,
|
||||
None,
|
||||
)
|
||||
.expect("usage should build")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn settles_pending_terminal_usage() {
|
||||
let writer = TestSettlementWriter {
|
||||
has_writer: true,
|
||||
..Default::default()
|
||||
};
|
||||
let usage = sample_usage();
|
||||
|
||||
settle_usage_if_needed(&writer, &usage)
|
||||
.await
|
||||
.expect("settlement should succeed");
|
||||
|
||||
let inputs = writer.inputs.lock().expect("settlement inputs lock");
|
||||
assert_eq!(inputs.len(), 1);
|
||||
assert_eq!(inputs[0].request_id, "req-1");
|
||||
assert_eq!(inputs[0].status, "completed");
|
||||
assert_eq!(inputs[0].billing_status, "pending");
|
||||
assert_eq!(inputs[0].finalized_at_unix_secs, Some(200));
|
||||
assert_eq!(inputs[0].total_cost_usd, 1.25);
|
||||
assert_eq!(inputs[0].actual_total_cost_usd, 0.75);
|
||||
assert!(!inputs[0].api_key_is_standalone);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn skips_pending_cancelled_usage() {
|
||||
let writer = TestSettlementWriter {
|
||||
has_writer: true,
|
||||
..Default::default()
|
||||
};
|
||||
let mut usage = sample_usage();
|
||||
usage.status = "cancelled".to_string();
|
||||
usage.status_code = Some(499);
|
||||
|
||||
settle_usage_if_needed(&writer, &usage)
|
||||
.await
|
||||
.expect("skipped settlement should succeed");
|
||||
|
||||
let inputs = writer.inputs.lock().expect("settlement inputs lock");
|
||||
assert!(inputs.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn propagates_standalone_key_flag_from_usage_metadata() {
|
||||
let writer = TestSettlementWriter {
|
||||
has_writer: true,
|
||||
..Default::default()
|
||||
};
|
||||
let mut usage = sample_usage();
|
||||
usage.request_metadata = Some(json!({ "api_key_is_standalone": true }));
|
||||
|
||||
settle_usage_if_needed(&writer, &usage)
|
||||
.await
|
||||
.expect("settlement should succeed");
|
||||
|
||||
let inputs = writer.inputs.lock().expect("settlement inputs lock");
|
||||
assert_eq!(inputs.len(), 1);
|
||||
assert!(inputs[0].api_key_is_standalone);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn skips_when_usage_is_not_pending_or_terminal() {
|
||||
let writer = TestSettlementWriter {
|
||||
has_writer: true,
|
||||
..Default::default()
|
||||
};
|
||||
let mut usage = sample_usage();
|
||||
usage.billing_status = "settled".to_string();
|
||||
usage.status = "streaming".to_string();
|
||||
|
||||
settle_usage_if_needed(&writer, &usage)
|
||||
.await
|
||||
.expect("skipped settlement should succeed");
|
||||
|
||||
let inputs = writer.inputs.lock().expect("settlement inputs lock");
|
||||
assert!(inputs.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rejects_non_finite_costs_before_writing() {
|
||||
let writer = TestSettlementWriter {
|
||||
has_writer: true,
|
||||
..Default::default()
|
||||
};
|
||||
let mut usage = sample_usage();
|
||||
usage.total_cost_usd = f64::NAN;
|
||||
|
||||
let err = settle_usage_if_needed(&writer, &usage)
|
||||
.await
|
||||
.expect_err("non-finite costs should be rejected");
|
||||
|
||||
assert!(matches!(
|
||||
err,
|
||||
aether_data_contracts::DataLayerError::InvalidInput(_)
|
||||
));
|
||||
let inputs = writer.inputs.lock().expect("settlement inputs lock");
|
||||
assert!(inputs.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn serializes_settlements_for_same_billing_subject() {
|
||||
let writer = SlowSettlementWriter::default();
|
||||
let mut first = sample_usage();
|
||||
first.request_id = "req-same-subject-1".to_string();
|
||||
let mut second = sample_usage();
|
||||
second.request_id = "req-same-subject-2".to_string();
|
||||
|
||||
tokio::try_join!(
|
||||
settle_usage_if_needed(&writer, &first),
|
||||
settle_usage_if_needed(&writer, &second)
|
||||
)
|
||||
.expect("settlements should succeed");
|
||||
|
||||
assert_eq!(writer.max_active.load(Ordering::Acquire), 1);
|
||||
assert_eq!(writer.inputs.lock().expect("inputs lock").len(), 2);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub use aether_contracts::StandardizedUsage;
|
||||
@@ -0,0 +1,774 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::StandardizedUsage;
|
||||
|
||||
pub struct UsageMapper;
|
||||
|
||||
impl UsageMapper {
|
||||
pub fn map(
|
||||
raw_usage: &serde_json::Value,
|
||||
api_format: &str,
|
||||
extra_mapping: Option<&BTreeMap<String, String>>,
|
||||
) -> StandardizedUsage {
|
||||
if !raw_usage.is_object() {
|
||||
return StandardizedUsage::new();
|
||||
}
|
||||
|
||||
let mut usage = StandardizedUsage::new();
|
||||
let mut mapping = base_mapping(api_format);
|
||||
if let Some(extra_mapping) = extra_mapping {
|
||||
mapping.extend(extra_mapping.clone());
|
||||
}
|
||||
|
||||
for (source_path, target_field) in mapping {
|
||||
if let Some(value) = get_nested_value(raw_usage, &source_path) {
|
||||
usage.set(&target_field, value.clone());
|
||||
}
|
||||
}
|
||||
|
||||
apply_openai_cache_write_tokens(raw_usage, api_format, &mut usage);
|
||||
derive_missing_input_tokens(raw_usage, api_format, &mut usage);
|
||||
copy_explicit_total_tokens(raw_usage, api_format, &mut usage);
|
||||
usage.normalize_cache_creation_breakdown()
|
||||
}
|
||||
|
||||
pub fn map_from_response(response: &serde_json::Value, api_format: &str) -> StandardizedUsage {
|
||||
let family = api_family(api_format);
|
||||
let mut usage = if let Some(usage_value) = resolve_usage_value(response, family.as_str()) {
|
||||
Self::map(usage_value, api_format, None)
|
||||
} else {
|
||||
StandardizedUsage::new()
|
||||
};
|
||||
if is_openai_image_api(api_format) {
|
||||
apply_openai_image_response_dimensions(response, &mut usage);
|
||||
}
|
||||
usage
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_openai_cache_write_tokens(
|
||||
raw_usage: &serde_json::Value,
|
||||
api_format: &str,
|
||||
usage: &mut StandardizedUsage,
|
||||
) {
|
||||
if api_family(api_format).as_str() != "openai" {
|
||||
return;
|
||||
}
|
||||
for details_key in ["prompt_tokens_details", "input_tokens_details"] {
|
||||
if let Some(value) = raw_usage
|
||||
.get(details_key)
|
||||
.and_then(|details| details.get("cache_write_tokens"))
|
||||
{
|
||||
usage.set("cache_creation_tokens", value.clone());
|
||||
return;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn map_usage(raw_usage: &serde_json::Value, api_format: &str) -> StandardizedUsage {
|
||||
UsageMapper::map(raw_usage, api_format, None)
|
||||
}
|
||||
|
||||
pub fn map_usage_from_response(
|
||||
response: &serde_json::Value,
|
||||
api_format: &str,
|
||||
) -> StandardizedUsage {
|
||||
UsageMapper::map_from_response(response, api_format)
|
||||
}
|
||||
|
||||
fn api_family(api_format: &str) -> String {
|
||||
api_format
|
||||
.split(':')
|
||||
.next()
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn api_kind(api_format: &str) -> String {
|
||||
api_format
|
||||
.split(':')
|
||||
.nth(1)
|
||||
.unwrap_or_default()
|
||||
.trim()
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn is_openai_image_api(api_format: &str) -> bool {
|
||||
api_family(api_format).as_str() == "openai" && api_kind(api_format).as_str() == "image"
|
||||
}
|
||||
|
||||
fn apply_openai_image_response_dimensions(
|
||||
response: &serde_json::Value,
|
||||
usage: &mut StandardizedUsage,
|
||||
) {
|
||||
let image_count = openai_image_response_image_count(response);
|
||||
if image_count <= 0 {
|
||||
return;
|
||||
}
|
||||
|
||||
usage.request_count = image_count;
|
||||
usage
|
||||
.dimensions
|
||||
.insert("image_count".to_string(), serde_json::json!(image_count));
|
||||
}
|
||||
|
||||
fn openai_image_response_image_count(response: &serde_json::Value) -> i64 {
|
||||
response
|
||||
.get("data")
|
||||
.and_then(serde_json::Value::as_array)
|
||||
.map(|items| items.len() as i64)
|
||||
.filter(|value| *value > 0)
|
||||
.or_else(|| image_result_count(response.get("result")))
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn image_result_count(value: Option<&serde_json::Value>) -> Option<i64> {
|
||||
match value? {
|
||||
serde_json::Value::Array(items) => Some(items.len() as i64).filter(|count| *count > 0),
|
||||
serde_json::Value::Object(object) if !object.is_empty() => Some(1),
|
||||
serde_json::Value::String(text) if !text.trim().is_empty() => Some(1),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn base_mapping(api_format: &str) -> BTreeMap<String, String> {
|
||||
let mut mapping = BTreeMap::new();
|
||||
match api_family(api_format).as_str() {
|
||||
"openai" => {
|
||||
mapping.insert("prompt_tokens".to_string(), "input_tokens".to_string());
|
||||
mapping.insert("completion_tokens".to_string(), "output_tokens".to_string());
|
||||
mapping.insert("input_tokens".to_string(), "input_tokens".to_string());
|
||||
mapping.insert("output_tokens".to_string(), "output_tokens".to_string());
|
||||
mapping.insert(
|
||||
"cache_creation_input_tokens".to_string(),
|
||||
"cache_creation_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"cache_creation.ephemeral_5m_input_tokens".to_string(),
|
||||
"cache_creation_ephemeral_5m_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"cache_creation.ephemeral_1h_input_tokens".to_string(),
|
||||
"cache_creation_ephemeral_1h_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"cache_read_input_tokens".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"prompt_tokens_details.cached_tokens".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"input_tokens_details.cached_tokens".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"prompt_tokens_details.cached_creation_tokens".to_string(),
|
||||
"cache_creation_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"prompt_tokens_details.cache_write_tokens".to_string(),
|
||||
"cache_creation_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"input_tokens_details.cached_creation_tokens".to_string(),
|
||||
"cache_creation_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"input_tokens_details.cache_write_tokens".to_string(),
|
||||
"cache_creation_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"completion_tokens_details.reasoning_tokens".to_string(),
|
||||
"reasoning_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"output_tokens_details.reasoning_tokens".to_string(),
|
||||
"reasoning_tokens".to_string(),
|
||||
);
|
||||
}
|
||||
"gemini" => {
|
||||
mapping.insert("promptTokenCount".to_string(), "input_tokens".to_string());
|
||||
mapping.insert(
|
||||
"candidatesTokenCount".to_string(),
|
||||
"output_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"cachedContentTokenCount".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"usageMetadata.promptTokenCount".to_string(),
|
||||
"input_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"usageMetadata.candidatesTokenCount".to_string(),
|
||||
"output_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"usageMetadata.cachedContentTokenCount".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
}
|
||||
"claude" | "anthropic" => {
|
||||
mapping.insert("input_tokens".to_string(), "input_tokens".to_string());
|
||||
mapping.insert("output_tokens".to_string(), "output_tokens".to_string());
|
||||
mapping.insert(
|
||||
"cache_creation_input_tokens".to_string(),
|
||||
"cache_creation_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"cache_creation.ephemeral_5m_input_tokens".to_string(),
|
||||
"cache_creation_ephemeral_5m_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"cache_creation.ephemeral_1h_input_tokens".to_string(),
|
||||
"cache_creation_ephemeral_1h_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"cache_read_input_tokens".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
}
|
||||
_ => {
|
||||
mapping.insert("input_tokens".to_string(), "input_tokens".to_string());
|
||||
mapping.insert("output_tokens".to_string(), "output_tokens".to_string());
|
||||
mapping.insert(
|
||||
"cache_creation_input_tokens".to_string(),
|
||||
"cache_creation_tokens".to_string(),
|
||||
);
|
||||
mapping.insert(
|
||||
"cache_read_input_tokens".to_string(),
|
||||
"cache_read_tokens".to_string(),
|
||||
);
|
||||
}
|
||||
}
|
||||
mapping
|
||||
}
|
||||
|
||||
fn derive_missing_input_tokens(
|
||||
raw_usage: &serde_json::Value,
|
||||
api_format: &str,
|
||||
usage: &mut StandardizedUsage,
|
||||
) {
|
||||
if usage.input_tokens > 0 || api_family(api_format).as_str() != "openai" {
|
||||
return;
|
||||
}
|
||||
|
||||
let Some(total_tokens) = numeric_i64(raw_usage.get("total_tokens")) else {
|
||||
return;
|
||||
};
|
||||
let output_tokens = usage
|
||||
.output_tokens
|
||||
.max(numeric_i64(raw_usage.get("completion_tokens")).unwrap_or_default())
|
||||
.max(numeric_i64(raw_usage.get("output_tokens")).unwrap_or_default());
|
||||
let inferred_input_tokens = total_tokens.saturating_sub(output_tokens);
|
||||
if inferred_input_tokens > 0 {
|
||||
usage.input_tokens = inferred_input_tokens;
|
||||
}
|
||||
}
|
||||
|
||||
fn copy_explicit_total_tokens(
|
||||
raw_usage: &serde_json::Value,
|
||||
api_format: &str,
|
||||
usage: &mut StandardizedUsage,
|
||||
) {
|
||||
let total_tokens = match api_family(api_format).as_str() {
|
||||
"gemini" => numeric_i64(raw_usage.get("totalTokenCount")),
|
||||
_ => numeric_i64(raw_usage.get("total_tokens")),
|
||||
};
|
||||
if let Some(total_tokens) = total_tokens.filter(|value| *value > 0) {
|
||||
usage
|
||||
.dimensions
|
||||
.insert("total_tokens".to_string(), serde_json::json!(total_tokens));
|
||||
}
|
||||
}
|
||||
|
||||
fn numeric_i64(value: Option<&serde_json::Value>) -> Option<i64> {
|
||||
value.and_then(|value| {
|
||||
value
|
||||
.as_i64()
|
||||
.or_else(|| value.as_u64().and_then(|number| i64::try_from(number).ok()))
|
||||
})
|
||||
}
|
||||
|
||||
fn get_nested_value<'a>(value: &'a serde_json::Value, path: &str) -> Option<&'a serde_json::Value> {
|
||||
let mut current = value;
|
||||
for segment in path.split('.') {
|
||||
current = current.get(segment)?;
|
||||
}
|
||||
Some(current)
|
||||
}
|
||||
|
||||
fn resolve_usage_value<'a>(
|
||||
response: &'a serde_json::Value,
|
||||
family: &str,
|
||||
) -> Option<&'a serde_json::Value> {
|
||||
match family {
|
||||
"gemini" => {
|
||||
if let Some(usage) = response.get("usageMetadata") {
|
||||
return Some(usage);
|
||||
}
|
||||
if let Some(usage) = response
|
||||
.get("candidates")
|
||||
.and_then(|value| value.get(0))
|
||||
.and_then(|value| value.get("usageMetadata"))
|
||||
{
|
||||
return Some(usage);
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
if let Some(usage) = response.get("usage") {
|
||||
return Some(usage);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for nested_key in ["response", "message", "item"] {
|
||||
if let Some(nested) = response.get(nested_key) {
|
||||
if let Some(usage) = resolve_usage_value(nested, family) {
|
||||
return Some(usage);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(chunks) = response.get("chunks").and_then(serde_json::Value::as_array) {
|
||||
for chunk in chunks.iter().rev() {
|
||||
if let Some(usage) = resolve_usage_value(chunk, family) {
|
||||
return Some(usage);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{map_usage, map_usage_from_response};
|
||||
|
||||
#[test]
|
||||
fn maps_openai_usage() {
|
||||
let usage = map_usage(
|
||||
&serde_json::json!({
|
||||
"prompt_tokens": 12,
|
||||
"completion_tokens": 8,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 2,
|
||||
"cached_creation_tokens": 1
|
||||
},
|
||||
"completion_tokens_details": { "reasoning_tokens": 3 }
|
||||
}),
|
||||
"openai:chat",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 12);
|
||||
assert_eq!(usage.output_tokens, 8);
|
||||
assert_eq!(usage.cache_creation_tokens, 1);
|
||||
assert_eq!(usage.cache_read_tokens, 2);
|
||||
assert_eq!(usage.reasoning_tokens, 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_openai_responses_usage_from_response() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"usage": {
|
||||
"input_tokens": 14,
|
||||
"output_tokens": 6,
|
||||
"total_tokens": 20,
|
||||
"input_tokens_details": {
|
||||
"cached_tokens": 3,
|
||||
"cached_creation_tokens": 2
|
||||
},
|
||||
"output_tokens_details": {
|
||||
"reasoning_tokens": 1
|
||||
}
|
||||
}
|
||||
}),
|
||||
"openai:responses",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 14);
|
||||
assert_eq!(usage.output_tokens, 6);
|
||||
assert_eq!(usage.cache_creation_tokens, 2);
|
||||
assert_eq!(usage.cache_read_tokens, 3);
|
||||
assert_eq!(usage.reasoning_tokens, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_openai_responses_cache_write_tokens() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"usage": {
|
||||
"input_tokens": 32_963,
|
||||
"input_tokens_details": {
|
||||
"cache_write_tokens": 512,
|
||||
"cached_creation_tokens": 1,
|
||||
"cached_tokens": 30_336
|
||||
},
|
||||
"output_tokens": 129,
|
||||
"output_tokens_details": {
|
||||
"reasoning_tokens": 8
|
||||
},
|
||||
"total_tokens": 33_092
|
||||
}
|
||||
}),
|
||||
"openai:responses",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 32_963);
|
||||
assert_eq!(usage.output_tokens, 129);
|
||||
assert_eq!(usage.cache_creation_tokens, 512);
|
||||
assert_eq!(usage.cache_read_tokens, 30_336);
|
||||
assert_eq!(usage.reasoning_tokens, 8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_openai_responses_usage_with_missing_input_from_total() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"usage": {
|
||||
"output_tokens": 899,
|
||||
"total_tokens": 53_499,
|
||||
"input_tokens_details": {
|
||||
"cached_tokens": 52_600
|
||||
}
|
||||
}
|
||||
}),
|
||||
"openai:responses",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 52_600);
|
||||
assert_eq!(usage.output_tokens, 899);
|
||||
assert_eq!(usage.cache_read_tokens, 52_600);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn keeps_missing_input_derivation_scoped_to_openai() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"usage": {
|
||||
"output_tokens": 7,
|
||||
"total_tokens": 17,
|
||||
"cache_read_input_tokens": 10
|
||||
}
|
||||
}),
|
||||
"claude:messages",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 0);
|
||||
assert_eq!(usage.output_tokens, 7);
|
||||
assert_eq!(usage.cache_read_tokens, 10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_openai_responses_with_top_level_cache_fields() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"usage": {
|
||||
"input_tokens": 6,
|
||||
"output_tokens": 20,
|
||||
"cache_creation_input_tokens": 42_262,
|
||||
"cache_read_input_tokens": 0
|
||||
}
|
||||
}),
|
||||
"openai:chat",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 6);
|
||||
assert_eq!(usage.output_tokens, 20);
|
||||
assert_eq!(usage.cache_creation_tokens, 42_262);
|
||||
assert_eq!(usage.cache_read_tokens, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_openai_responses_usage_from_stream_chunks() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"chunks": [
|
||||
{
|
||||
"type": "response.created",
|
||||
"response": {
|
||||
"id": "resp_123",
|
||||
"object": "response"
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": "resp_123",
|
||||
"object": "response",
|
||||
"usage": {
|
||||
"input_tokens": 9,
|
||||
"output_tokens": 4,
|
||||
"total_tokens": 13,
|
||||
"input_tokens_details": {
|
||||
"cached_tokens": 5,
|
||||
"cached_creation_tokens": 2
|
||||
},
|
||||
"output_tokens_details": {
|
||||
"reasoning_tokens": 1
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}),
|
||||
"openai:responses",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 9);
|
||||
assert_eq!(usage.output_tokens, 4);
|
||||
assert_eq!(usage.cache_creation_tokens, 2);
|
||||
assert_eq!(usage.cache_read_tokens, 5);
|
||||
assert_eq!(usage.reasoning_tokens, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_claude_usage() {
|
||||
let usage = map_usage(
|
||||
&serde_json::json!({
|
||||
"input_tokens": 10,
|
||||
"output_tokens": 5,
|
||||
"cache_creation_input_tokens": 4,
|
||||
"cache_read_input_tokens": 1
|
||||
}),
|
||||
"claude:messages",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 10);
|
||||
assert_eq!(usage.output_tokens, 5);
|
||||
assert_eq!(usage.cache_creation_tokens, 4);
|
||||
assert_eq!(usage.cache_read_tokens, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_claude_usage_from_stream_chunks() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"chunks": [
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"usage": {
|
||||
"input_tokens": 5,
|
||||
"cache_creation_input_tokens": 59_573,
|
||||
"cache_read_input_tokens": 0,
|
||||
"output_tokens": 0
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"type": "message_delta",
|
||||
"usage": {
|
||||
"input_tokens": 5,
|
||||
"cache_creation_input_tokens": 59_573,
|
||||
"cache_read_input_tokens": 0,
|
||||
"output_tokens": 162
|
||||
}
|
||||
}
|
||||
]
|
||||
}),
|
||||
"claude:messages",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 5);
|
||||
assert_eq!(usage.output_tokens, 162);
|
||||
assert_eq!(usage.cache_creation_tokens, 59_573);
|
||||
assert_eq!(usage.cache_read_tokens, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_claude_usage_from_message_start_chunk() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"chunks": [
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"usage": {
|
||||
"input_tokens": 5,
|
||||
"cache_creation_input_tokens": 59_573,
|
||||
"cache_read_input_tokens": 0,
|
||||
"output_tokens": 0
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}),
|
||||
"claude:messages",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 5);
|
||||
assert_eq!(usage.output_tokens, 0);
|
||||
assert_eq!(usage.cache_creation_tokens, 59_573);
|
||||
assert_eq!(usage.cache_read_tokens, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_claude_usage_with_large_cache_read_tokens_without_subtracting_input() {
|
||||
let usage = map_usage(
|
||||
&serde_json::json!({
|
||||
"input_tokens": 4941,
|
||||
"cache_creation_input_tokens": 687,
|
||||
"cache_read_input_tokens": 52873,
|
||||
"output_tokens": 973
|
||||
}),
|
||||
"claude:messages",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 4941);
|
||||
assert_eq!(usage.cache_creation_tokens, 687);
|
||||
assert_eq!(usage.cache_read_tokens, 52873);
|
||||
assert_eq!(usage.output_tokens, 973);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_claude_usage_with_cache_creation_total_and_zero_ttl_breakdown() {
|
||||
let usage = map_usage(
|
||||
&serde_json::json!({
|
||||
"cache_creation": {
|
||||
"ephemeral_1h_input_tokens": 0,
|
||||
"ephemeral_5m_input_tokens": 0
|
||||
},
|
||||
"cache_creation_input_tokens": 2051,
|
||||
"cache_read_input_tokens": 2051,
|
||||
"inference_geo": "inference_geo",
|
||||
"input_tokens": 2095,
|
||||
"output_tokens": 503,
|
||||
"server_tool_use": {
|
||||
"web_fetch_requests": 2,
|
||||
"web_search_requests": 0
|
||||
},
|
||||
"service_tier": "standard"
|
||||
}),
|
||||
"claude:messages",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 2095);
|
||||
assert_eq!(usage.cache_creation_tokens, 2051);
|
||||
assert_eq!(usage.cache_creation_ephemeral_5m_tokens, 0);
|
||||
assert_eq!(usage.cache_creation_ephemeral_1h_tokens, 0);
|
||||
assert_eq!(usage.cache_read_tokens, 2051);
|
||||
assert_eq!(usage.output_tokens, 503);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_claude_usage_with_ephemeral_cache_breakdown() {
|
||||
let usage = map_usage(
|
||||
&serde_json::json!({
|
||||
"input_tokens": 1,
|
||||
"output_tokens": 8,
|
||||
"cache_creation": {
|
||||
"ephemeral_1h_input_tokens": 0,
|
||||
"ephemeral_5m_input_tokens": 5191
|
||||
},
|
||||
"cache_creation_input_tokens": 5191,
|
||||
"cache_read_input_tokens": 97634
|
||||
}),
|
||||
"claude:messages",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 1);
|
||||
assert_eq!(usage.output_tokens, 8);
|
||||
assert_eq!(usage.cache_creation_tokens, 5191);
|
||||
assert_eq!(usage.cache_creation_ephemeral_5m_tokens, 5191);
|
||||
assert_eq!(usage.cache_creation_ephemeral_1h_tokens, 0);
|
||||
assert_eq!(usage.cache_read_tokens, 97634);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_gemini_usage_from_response() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 14,
|
||||
"candidatesTokenCount": 6,
|
||||
"cachedContentTokenCount": 2
|
||||
}
|
||||
}),
|
||||
"gemini:generate_content",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 14);
|
||||
assert_eq!(usage.output_tokens, 6);
|
||||
assert_eq!(usage.cache_read_tokens, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_gemini_usage_from_stream_chunks() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"chunks": [
|
||||
{
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [{ "text": "hello" }]
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 14,
|
||||
"candidatesTokenCount": 6,
|
||||
"cachedContentTokenCount": 2
|
||||
}
|
||||
}
|
||||
]
|
||||
}),
|
||||
"gemini:generate_content",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 14);
|
||||
assert_eq!(usage.output_tokens, 6);
|
||||
assert_eq!(usage.cache_read_tokens, 2);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_openai_image_response_dimensions_without_usage() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"created": 1_700_000_000,
|
||||
"data": [
|
||||
{ "b64_json": "abc" },
|
||||
{ "url": "https://example.test/image.png" }
|
||||
]
|
||||
}),
|
||||
"openai:image",
|
||||
);
|
||||
|
||||
assert_eq!(usage.request_count, 2);
|
||||
assert_eq!(
|
||||
usage.dimensions.get("image_count"),
|
||||
Some(&serde_json::json!(2))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn maps_openai_image_response_dimensions_with_native_usage() {
|
||||
let usage = map_usage_from_response(
|
||||
&serde_json::json!({
|
||||
"usage": {
|
||||
"input_tokens": 11,
|
||||
"output_tokens": 22,
|
||||
"total_tokens": 33
|
||||
},
|
||||
"data": [{ "b64_json": "abc" }]
|
||||
}),
|
||||
"openai:image",
|
||||
);
|
||||
|
||||
assert_eq!(usage.input_tokens, 11);
|
||||
assert_eq!(usage.output_tokens, 22);
|
||||
assert_eq!(usage.request_count, 1);
|
||||
assert_eq!(
|
||||
usage.dimensions.get("image_count"),
|
||||
Some(&serde_json::json!(1))
|
||||
);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user