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

refactor: 优化调度候选排序与用量写入链路并改进 Fernet 缓存与前端批量列表

Browse Source
This commit is contained in:
fawney19
2026-04-21 16:19:07 +08:00
parent c5c56ff92f
commit 25a2b417be
77 changed files with 4801 additions and 2881 deletions
Download Patch File Download Diff File Expand all files Collapse all files

416
crates/aether-usage-runtime/src/body_capture.rs
View File

@@ -1,6 +1,9 @@
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;
@@ -76,6 +79,22 @@ 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,
@@ -286,9 +305,7 @@ fn limit_usage_body_capture_value(
value: Value,
max_bytes: Option<usize>,
) -> LimitedUsageBodyCapture {
let source_bytes = serde_json::to_vec(&value)
.ok()
.map(|bytes| bytes.len() as u64);
let source_bytes = json_serialized_len(&value);
let Some(limit) = max_bytes.filter(|value| *value > 0) else {
return LimitedUsageBodyCapture {
stored_bytes: source_bytes,
@@ -327,9 +344,7 @@ fn limit_usage_body_capture_value(
"value_kind": usage_value_kind(&other),
}),
};
let stored_bytes = serde_json::to_vec(&truncated_value)
.ok()
.map(|bytes| bytes.len() as u64);
let stored_bytes = json_serialized_len(&truncated_value);
LimitedUsageBodyCapture {
value: truncated_value,
source_bytes: Some(source_len),
@@ -340,13 +355,6 @@ fn limit_usage_body_capture_value(
}
fn truncate_usage_body_string(value: &str, max_bytes: usize) -> String {
if serde_json::to_vec(&Value::String(value.to_string()))
.ok()
.is_some_and(|bytes| bytes.len() <= max_bytes)
{
return value.to_string();
}
let mut end = value.len();
while end > 0 {
while end > 0 && !value.is_char_boundary(end) {
@@ -354,10 +362,7 @@ fn truncate_usage_body_string(value: &str, max_bytes: usize) -> String {
}
let mut candidate = value[..end].to_string();
candidate.push_str(TRUNCATED_BODY_STRING_SUFFIX);
if serde_json::to_vec(&Value::String(candidate.clone()))
.ok()
.is_some_and(|bytes| bytes.len() <= max_bytes)
{
if json_serialized_len(&candidate).is_some_and(|bytes| bytes <= max_bytes as u64) {
return candidate;
}
end = value[..end]
@@ -379,27 +384,45 @@ fn truncate_usage_body_string(value: &str, max_bytes: usize) -> 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 {
if let Some(object) = metadata.as_mut().and_then(Value::as_object_mut) {
object.remove(field.as_ref_key());
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(
field.as_ref_key().to_string(),
Value::String(body_ref.to_string()),
);
object.insert(key.to_owned(), Value::String(body_ref.to_owned()));
}
pub(crate) fn build_payload_body_capture_metadata(
@@ -408,36 +431,44 @@ pub(crate) fn build_payload_body_capture_metadata(
provider_body_state: Option<UsageBodyCaptureState>,
client_body_state: Option<UsageBodyCaptureState>,
) -> Option<Value> {
let mut metadata = Map::new();
if let Some(decoded_len) = provider_body_base64.and_then(decoded_base64_len_hint) {
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_body_base64.and_then(decoded_base64_len_hint) {
if let Some(decoded_len) = client_decoded_len {
metadata.insert(
"client_response_body_base64_bytes".to_string(),
Value::Number(decoded_len.into()),
);
}
let mut body_capture = Map::new();
append_body_capture_metadata_entry(
&mut body_capture,
"response",
provider_body_state,
provider_body_base64.and_then(decoded_base64_len_hint),
provider_body_base64.and_then(decoded_base64_len_hint),
);
append_body_capture_metadata_entry(
&mut body_capture,
"client_response",
client_body_state,
client_body_base64.and_then(decoded_base64_len_hint),
client_body_base64.and_then(decoded_base64_len_hint),
);
if !body_capture.is_empty() {
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));
}
@@ -476,21 +507,29 @@ pub(crate) fn append_runtime_body_capture_metadata(
input.provider_request_body_ref,
input.provider_request_unavailable,
);
upsert_body_capture_metadata_entry(metadata, "request", Some(states.request), None, None, None);
upsert_body_capture_metadata_entry(
metadata,
"provider_request",
Some(states.provider_request),
input.provider_request_source_bytes,
input.provider_request_source_bytes,
input.provider_request_unavailable_reason,
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> {
let mut metadata = Map::new();
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))
}
@@ -507,13 +546,17 @@ pub(crate) fn append_plan_body_capture_metadata(
Value::Number(decoded_len.into()),
);
}
upsert_body_capture_metadata_entry(
metadata,
"provider_request",
Some(UsageBodyCaptureState::Unavailable),
decoded_len,
decoded_len,
Some("body_bytes_base64_only"),
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"),
),
);
}
}
@@ -528,58 +571,16 @@ fn append_body_capture_metadata_entry(
let Some(state) = state else {
return;
};
let mut entry = Map::new();
entry.insert(
"state".to_string(),
Value::String(state.as_str().to_string()),
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"),
),
);
if let Some(stored_bytes) = stored_bytes {
entry.insert("stored_bytes".to_string(), json!(stored_bytes));
}
if let Some(source_bytes) = source_bytes {
entry.insert("source_bytes".to_string(), json!(source_bytes));
}
if matches!(state, UsageBodyCaptureState::Truncated) {
entry.insert(
"reason".to_string(),
Value::String("body_capture_limit_exceeded".to_string()),
);
}
target.insert(key.to_string(), Value::Object(entry));
}
pub(crate) fn upsert_body_capture_metadata_entry(
metadata: &mut Map<String, Value>,
key: &str,
state: Option<UsageBodyCaptureState>,
stored_bytes: Option<u64>,
source_bytes: Option<u64>,
reason: Option<&str>,
) {
let body_capture = metadata
.entry("body_capture".to_string())
.or_insert_with(|| Value::Object(Map::new()));
let Some(body_capture_object) = body_capture.as_object_mut() else {
return;
};
let Some(state) = state else {
return;
};
let mut entry = Map::new();
entry.insert(
"state".to_string(),
Value::String(state.as_str().to_string()),
);
if let Some(bytes) = stored_bytes {
entry.insert("stored_bytes".to_string(), json!(bytes));
}
if let Some(bytes) = source_bytes {
entry.insert("source_bytes".to_string(), json!(bytes));
}
if let Some(reason) = reason {
entry.insert("reason".to_string(), Value::String(reason.to_string()));
}
body_capture_object.insert(key.to_string(), Value::Object(entry));
}
fn upsert_body_capture_metadata_value_entry(
@@ -593,22 +594,63 @@ fn upsert_body_capture_metadata_value_entry(
let Some(state) = state else {
return;
};
let metadata_object = metadata
.get_or_insert_with(|| Value::Object(Map::new()))
.as_object_mut();
let Some(metadata_object) = metadata_object else {
let Some(body_capture_object) = body_capture_value_object_mut(metadata, 1) else {
return;
};
upsert_body_capture_metadata_entry(
metadata_object,
key,
Some(state),
stored_bytes,
source_bytes,
reason,
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() {
@@ -642,9 +684,18 @@ pub(crate) fn decoded_base64_len_hint(body_base64: &str) -> Option<u64> {
}
fn sanitize_usage_body_ref(value: Option<String>) -> Option<String> {
value
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
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 {
@@ -657,3 +708,122 @@ fn usage_value_kind(value: &Value) -> &'static str {
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));
}
}

147
crates/aether-usage-runtime/src/request_metadata.rs
View File

@@ -31,13 +31,36 @@ pub(crate) fn merge_usage_request_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();
copy_allowed_metadata_fields(&object, &mut filtered);
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))
}
@@ -62,6 +85,26 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
copy_number(source, target, "price_per_request");
}
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_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, "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, "dimensions");
remove_non_null_value(&mut source, target, "billing_rule_snapshot");
remove_non_null_value(&mut source, target, "scheduling_audit");
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");
}
fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
let Some(value) = source
.get(key)
@@ -77,6 +120,20 @@ fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, V
);
}
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;
@@ -84,6 +141,13 @@ fn copy_number(source: &Map<String, Value>, target: &mut Map<String, Value>, key
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;
@@ -91,6 +155,13 @@ fn copy_bool(source: &Map<String, Value>, target: &mut Map<String, Value>, key:
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;
@@ -101,6 +172,20 @@ fn copy_non_null_value(source: &Map<String, Value>, target: &mut Map<String, Val
);
}
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)),
@@ -109,6 +194,14 @@ fn sanitize_usage_request_metadata_value(value: &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]";
@@ -134,6 +227,24 @@ fn truncate_usage_request_metadata_string(value: &str) -> 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,
@@ -217,7 +328,8 @@ mod tests {
use super::{
build_usage_request_metadata_seed, merge_usage_request_metadata,
sanitize_usage_request_metadata, MAX_USAGE_REQUEST_METADATA_BYTES,
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,
};
@@ -363,4 +475,35 @@ mod tests {
assert_eq!(metadata, None);
}
#[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"
}));
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))
);
}
}

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

@@ -12,8 +12,7 @@ use tracing::warn;
use crate::executor::spawn_on_usage_background_runtime;
use crate::{
apply_usage_body_capture_policy_to_event, apply_usage_body_capture_policy_to_record,
build_pending_usage_record_from_seed, build_stream_terminal_usage_seed,
build_streaming_usage_record_from_seed, build_sync_terminal_usage_seed,
build_stream_terminal_usage_seed, build_sync_terminal_usage_seed,
build_terminal_usage_event_from_seed, build_upsert_usage_record_from_event,
build_usage_queue_worker, settle_usage_if_needed, LifecycleUsageSeed,
StreamTerminalUsagePayloadSeed, SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed,
@@ -74,17 +73,6 @@ pub struct UsageRuntime {
config: UsageRuntimeConfig,
}
struct SyncTerminalUsageTaskInput {
context_seed: TerminalUsageContextSeed,
payload_seed: SyncTerminalUsagePayloadSeed,
}
struct StreamTerminalUsageTaskInput {
context_seed: TerminalUsageContextSeed,
payload_seed: StreamTerminalUsagePayloadSeed,
cancelled: bool,
}
impl Default for UsageRuntime {
fn default() -> Self {
Self::disabled()
@@ -126,7 +114,7 @@ impl UsageRuntime {
Some(worker.spawn())
}
pub fn record_pending<T>(&self, data: &T, seed: &LifecycleUsageSeed)
pub fn record_pending<T>(&self, data: &T, seed: LifecycleUsageSeed)
where
T: UsageRuntimeAccess + Clone + 'static,
{
@@ -134,11 +122,10 @@ impl UsageRuntime {
return;
}
let data = T::clone(data);
let seed = seed.clone();
let request_id = seed.request_id.clone();
spawn_on_usage_background_runtime(boxed_usage_task(async move {
let now_unix_secs = now_unix_secs();
match build_pending_usage_record_offthread(&seed, now_unix_secs).await {
match build_pending_usage_record_offthread(seed, now_unix_secs).await {
Ok(mut record) => {
apply_body_capture_policy_to_record_from_data(&data, &mut record).await;
if let Err(err) = data.upsert_usage_record(record).await {
@@ -183,9 +170,9 @@ impl UsageRuntime {
spawn_on_usage_background_runtime(boxed_usage_task(async move {
let now_unix_secs = now_unix_secs();
match build_streaming_usage_record_offthread(
&seed,
seed,
status_code,
telemetry.as_ref(),
telemetry,
now_unix_secs,
)
.await
@@ -218,8 +205,8 @@ impl UsageRuntime {
pub fn record_sync_terminal<T>(
&self,
data: &T,
context_seed: &TerminalUsageContextSeed,
payload_seed: &SyncTerminalUsagePayloadSeed,
context_seed: TerminalUsageContextSeed,
payload_seed: SyncTerminalUsagePayloadSeed,
) where
T: UsageRuntimeAccess + Clone + 'static,
{
@@ -229,12 +216,8 @@ impl UsageRuntime {
let runtime = self.clone();
let data = T::clone(data);
let request_id = context_seed.request_id.clone();
let input = Box::new(SyncTerminalUsageTaskInput {
context_seed: context_seed.clone(),
payload_seed: payload_seed.clone(),
});
spawn_on_usage_background_runtime(boxed_usage_task(async move {
match build_sync_terminal_usage_event_offthread(input).await {
match build_sync_terminal_usage_event_offthread(context_seed, payload_seed).await {
Ok(mut event) => {
apply_body_capture_policy_from_data(&data, &mut event).await;
if let Err(err) = data.enrich_usage_event(&mut event).await {
@@ -264,8 +247,8 @@ impl UsageRuntime {
pub fn record_stream_terminal<T>(
&self,
data: &T,
context_seed: &TerminalUsageContextSeed,
payload_seed: &StreamTerminalUsagePayloadSeed,
context_seed: TerminalUsageContextSeed,
payload_seed: StreamTerminalUsagePayloadSeed,
cancelled: bool,
) where
T: UsageRuntimeAccess + Clone + 'static,
@@ -276,13 +259,10 @@ impl UsageRuntime {
let runtime = self.clone();
let data = T::clone(data);
let request_id = context_seed.request_id.clone();
let input = Box::new(StreamTerminalUsageTaskInput {
context_seed: context_seed.clone(),
payload_seed: payload_seed.clone(),
cancelled,
});
spawn_on_usage_background_runtime(boxed_usage_task(async move {
match build_stream_terminal_usage_event_offthread(input).await {
match build_stream_terminal_usage_event_offthread(context_seed, payload_seed, cancelled)
.await
{
Ok(mut event) => {
apply_body_capture_policy_from_data(&data, &mut event).await;
if let Err(err) = data.enrich_usage_event(&mut event).await {
@@ -413,28 +393,27 @@ impl UsageRuntime {
}
async fn build_pending_usage_record_offthread(
seed: &LifecycleUsageSeed,
seed: LifecycleUsageSeed,
now_unix_secs: u64,
) -> Result<UpsertUsageRecord, DataLayerError> {
let seed = seed.clone();
tokio::task::spawn_blocking(move || build_pending_usage_record_from_seed(&seed, now_unix_secs))
.await
.map_err(join_error_to_data_layer)?
tokio::task::spawn_blocking(move || {
crate::write::build_pending_usage_record_from_owned_seed(seed, now_unix_secs)
})
.await
.map_err(join_error_to_data_layer)?
}
async fn build_streaming_usage_record_offthread(
seed: &LifecycleUsageSeed,
seed: LifecycleUsageSeed,
status_code: u16,
telemetry: Option<&ExecutionTelemetry>,
telemetry: Option<ExecutionTelemetry>,
now_unix_secs: u64,
) -> Result<UpsertUsageRecord, DataLayerError> {
let seed = seed.clone();
let telemetry = telemetry.cloned();
tokio::task::spawn_blocking(move || {
build_streaming_usage_record_from_seed(
&seed,
crate::write::build_streaming_usage_record_from_owned_seed(
seed,
status_code,
telemetry.as_ref(),
telemetry,
now_unix_secs,
)
})
@@ -443,12 +422,13 @@ async fn build_streaming_usage_record_offthread(
}
async fn build_sync_terminal_usage_event_offthread(
input: Box<SyncTerminalUsageTaskInput>,
context_seed: TerminalUsageContextSeed,
payload_seed: SyncTerminalUsagePayloadSeed,
) -> Result<UsageEvent, DataLayerError> {
tokio::task::spawn_blocking(move || {
build_terminal_usage_event_from_seed(build_sync_terminal_usage_seed(
input.context_seed,
input.payload_seed,
context_seed,
payload_seed,
))
})
.await
@@ -456,13 +436,15 @@ async fn build_sync_terminal_usage_event_offthread(
}
async fn build_stream_terminal_usage_event_offthread(
input: Box<StreamTerminalUsageTaskInput>,
context_seed: TerminalUsageContextSeed,
payload_seed: StreamTerminalUsagePayloadSeed,
cancelled: bool,
) -> Result<UsageEvent, DataLayerError> {
tokio::task::spawn_blocking(move || {
build_terminal_usage_event_from_seed(build_stream_terminal_usage_seed(
input.context_seed,
input.payload_seed,
input.cancelled,
context_seed,
payload_seed,
cancelled,
))
})
.await

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

File diff suppressed because it is too large Load Diff