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

merge(main): 解决 usage 展示与生命周期同步冲突

Browse Source
This commit is contained in:
MMEXA
2026-07-18 05:38:38 +08:00
parent 3c348dff3a 6c33b8d8fb
commit e31df5989a
111 changed files with 8984 additions and 1715 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-usage/runtime/src/record.rs
+185 -9
View File
@@ -2,7 +2,10 @@ 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,
attach_client_request_body_metadata, attach_provider_request_body_metadata,
clear_client_request_body_metadata, clear_provider_request_body_metadata,
request_body_derived_facts_action, sanitize_usage_request_metadata,
RequestBodyDerivedFactsAction,
};
use crate::{UsageEvent, UsageEventType};
@@ -38,15 +41,42 @@ pub fn build_upsert_usage_record_from_event(
}
};
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()),
// Request-derived facts are captured before body capture policy is applied. Do not let a
// truncation/disabled placeholder clear those facts while converting the queued event into a
// database record. Inline (or ref-loaded) bodies remain authoritative and may clear stale
// metadata when the final upstream request no longer contains a value.
match request_body_derived_facts_action(data.request_body.as_ref(), data.request_body_state) {
RequestBodyDerivedFactsAction::Refresh => {
data.request_metadata = attach_client_request_body_metadata(
data.request_metadata,
data.request_body.as_ref(),
);
}
RequestBodyDerivedFactsAction::Clear => {
data.request_metadata = clear_client_request_body_metadata(data.request_metadata);
}
RequestBodyDerivedFactsAction::Preserve => {}
}
match request_body_derived_facts_action(
data.provider_request_body.as_ref(),
);
data.provider_request_body_state,
) {
RequestBodyDerivedFactsAction::Refresh => {
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(),
);
}
RequestBodyDerivedFactsAction::Clear => {
data.request_metadata = clear_provider_request_body_metadata(data.request_metadata);
}
RequestBodyDerivedFactsAction::Preserve => {}
}
let now_unix_secs = event.timestamp_ms / 1_000;
Ok(UpsertUsageRecord {
@@ -165,6 +195,8 @@ fn empty_to_none(value: Option<String>) -> Option<String> {
#[cfg(test)]
mod tests {
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
use crate::{UsageEvent, UsageEventData, UsageEventType};
use super::build_upsert_usage_record_from_event;
@@ -186,6 +218,9 @@ mod tests {
output_tokens: Some(20),
total_tokens: Some(30),
status_code: Some(200),
request_body: Some(serde_json::json!({
"reasoning": { "effort": "xhigh" }
})),
provider_request_body: Some(serde_json::json!({
"reasoning": { "effort": "max" },
"service_tier": "priority"
@@ -202,6 +237,14 @@ mod tests {
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("requested_reasoning_effort"))
.and_then(serde_json::Value::as_str),
Some("xhigh")
);
assert_eq!(
record
.request_metadata
@@ -229,6 +272,139 @@ mod tests {
assert_eq!(record.finalized_at_unix_secs, Some(1_700_000_000));
}
#[test]
fn truncated_request_placeholders_preserve_pre_capture_request_facts() {
let truncated = serde_json::json!({
"truncated": true,
"reason": "body_capture_limit_exceeded",
"max_bytes": 128,
"source_bytes": 2048,
"value_kind": "object"
});
let record = build_upsert_usage_record_from_event(&UsageEvent {
event_type: UsageEventType::Completed,
request_id: "req-truncated-request-facts".to_string(),
timestamp_ms: 1_700_000_000_000,
data: UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
api_format: Some("openai:responses".to_string()),
endpoint_api_format: Some("openai:responses".to_string()),
request_body: Some(truncated.clone()),
request_body_state: Some(UsageBodyCaptureState::Truncated),
provider_request_body: Some(truncated),
provider_request_body_state: Some(UsageBodyCaptureState::Truncated),
request_metadata: Some(serde_json::json!({
"requested_reasoning_effort": "xhigh",
"provider_reasoning_effort": "max",
"provider_service_tier": "priority"
})),
..UsageEventData::default()
},
})
.expect("record should build");
let metadata = record
.request_metadata
.as_ref()
.expect("derived request facts should remain");
assert_eq!(metadata["requested_reasoning_effort"], "xhigh");
assert_eq!(metadata["provider_reasoning_effort"], "max");
assert_eq!(metadata["provider_service_tier"], "priority");
}
#[test]
fn disabled_reference_and_unavailable_bodies_preserve_pre_capture_request_facts() {
for state in [
UsageBodyCaptureState::Disabled,
UsageBodyCaptureState::Reference,
UsageBodyCaptureState::Unavailable,
] {
let record = build_upsert_usage_record_from_event(&UsageEvent {
event_type: UsageEventType::Completed,
request_id: format!("req-preserved-{state:?}"),
timestamp_ms: 1_700_000_000_000,
data: UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
request_body_state: Some(state),
provider_request_body_state: Some(state),
request_metadata: Some(serde_json::json!({
"requested_reasoning_effort": "xhigh",
"provider_reasoning_effort": "max",
"provider_service_tier": "priority"
})),
..UsageEventData::default()
},
})
.expect("record should build");
let metadata = record
.request_metadata
.as_ref()
.expect("derived request facts should remain");
assert_eq!(metadata["requested_reasoning_effort"], "xhigh");
assert_eq!(metadata["provider_reasoning_effort"], "max");
assert_eq!(metadata["provider_service_tier"], "priority");
}
}
#[test]
fn missing_or_complete_factless_final_bodies_clear_stale_request_facts() {
let cases = [
("typed-missing", Some(UsageBodyCaptureState::None), None),
(
"typed-missing-with-stale-body",
Some(UsageBodyCaptureState::None),
Some(serde_json::json!({
"reasoning_effort": "xhigh",
"service_tier": "priority"
})),
),
("untyped-missing", None, None),
(
"inline-without-facts",
Some(UsageBodyCaptureState::Inline),
Some(serde_json::json!({"model": "gpt-5"})),
),
];
for (name, state, body) in cases {
let record = build_upsert_usage_record_from_event(&UsageEvent {
event_type: UsageEventType::Completed,
request_id: format!("req-clear-{name}"),
timestamp_ms: 1_700_000_000_000,
data: UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
request_body: body.clone(),
request_body_state: state,
provider_request_body: body,
provider_request_body_state: state,
request_metadata: Some(serde_json::json!({
"trace_id": "trace-1",
"requested_reasoning_effort": "xhigh",
"provider_reasoning_effort": "max",
"provider_service_tier": "priority",
"provider_actual_service_tier": "priority"
})),
..UsageEventData::default()
},
})
.expect("record should build");
let metadata = record
.request_metadata
.as_ref()
.expect("audit facts remain");
assert_eq!(metadata["trace_id"], "trace-1");
assert_eq!(metadata["provider_actual_service_tier"], "priority");
assert!(metadata.get("requested_reasoning_effort").is_none());
assert!(metadata.get("provider_reasoning_effort").is_none());
assert!(metadata.get("provider_service_tier").is_none());
}
}
#[test]
fn cancelled_terminal_record_is_void_for_billing() {
let record = build_upsert_usage_record_from_event(&UsageEvent {
crates/aether-usage/runtime/src/request_metadata.rs
+275 -4
View File
@@ -6,9 +6,10 @@ 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,
normalize_provider_service_tier, resolve_provider_cache_ttl_minutes, UsageBodyCaptureState,
PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY,
PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY,
REQUESTED_REASONING_EFFORT_METADATA_KEY,
};
use serde_json::{json, Map, Value};
@@ -17,6 +18,53 @@ 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;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RequestBodyDerivedFactsAction {
Refresh,
Preserve,
Clear,
}
pub(crate) fn request_body_derived_facts_action(
request_body: Option<&Value>,
state: Option<UsageBodyCaptureState>,
) -> RequestBodyDerivedFactsAction {
// A typed `none` marker is produced by the final capture attempt. It must take precedence over
// a body value that may have survived from an earlier candidate. A missing marker (`None`)
// remains compatible with legacy events, where a present body is still authoritative.
if state == Some(UsageBodyCaptureState::None) {
return RequestBodyDerivedFactsAction::Clear;
}
if let Some(request_body) = request_body {
if matches!(
state,
Some(
UsageBodyCaptureState::Truncated
| UsageBodyCaptureState::Disabled
| UsageBodyCaptureState::Unavailable
)
) || request_body.as_object().is_some_and(|body| {
body.get("truncated").and_then(Value::as_bool) == Some(true)
&& body.get("reason").and_then(Value::as_str) == Some("body_capture_limit_exceeded")
}) {
return RequestBodyDerivedFactsAction::Preserve;
}
return RequestBodyDerivedFactsAction::Refresh;
}
match state {
Some(
UsageBodyCaptureState::Inline
| UsageBodyCaptureState::Reference
| UsageBodyCaptureState::Truncated
| UsageBodyCaptureState::Disabled
| UsageBodyCaptureState::Unavailable,
) => RequestBodyDerivedFactsAction::Preserve,
Some(UsageBodyCaptureState::None) | None => RequestBodyDerivedFactsAction::Clear,
}
}
pub(crate) fn build_usage_request_metadata_seed(
_plan: &ExecutionPlan,
context: Option<&Map<String, Value>>,
@@ -76,6 +124,32 @@ pub(crate) fn sanitize_usage_request_metadata_ref(value: Option<&Value>) -> Opti
(!filtered.is_empty()).then_some(Value::Object(filtered))
}
pub(crate) fn attach_client_request_body_metadata(
metadata: Option<Value>,
request_body: Option<&Value>,
) -> Option<Value> {
let request_body_is_object = request_body.and_then(Value::as_object).is_some();
let reasoning_effort = extract_provider_reasoning_effort_from_body(request_body);
if !request_body_is_object && reasoning_effort.is_none() {
return metadata;
}
let mut object = match metadata {
Some(Value::Object(object)) => object,
_ => Map::new(),
};
if request_body_is_object {
object.remove(REQUESTED_REASONING_EFFORT_METADATA_KEY);
}
if let Some(reasoning_effort) = reasoning_effort {
object.insert(
REQUESTED_REASONING_EFFORT_METADATA_KEY.to_string(),
Value::String(reasoning_effort),
);
}
(!object.is_empty()).then_some(Value::Object(object))
}
pub(crate) fn attach_provider_request_body_metadata(
metadata: Option<Value>,
provider_api_format: Option<&str>,
@@ -129,6 +203,31 @@ pub(crate) fn attach_provider_request_body_metadata(
(!object.is_empty()).then_some(Value::Object(object))
}
pub(crate) fn clear_client_request_body_metadata(metadata: Option<Value>) -> Option<Value> {
clear_request_metadata_fields(metadata, &[REQUESTED_REASONING_EFFORT_METADATA_KEY])
}
pub(crate) fn clear_provider_request_body_metadata(metadata: Option<Value>) -> Option<Value> {
clear_request_metadata_fields(
metadata,
&[
PROVIDER_REASONING_EFFORT_METADATA_KEY,
PROVIDER_SERVICE_TIER_METADATA_KEY,
PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY,
],
)
}
fn clear_request_metadata_fields(metadata: Option<Value>, keys: &[&str]) -> Option<Value> {
let Some(Value::Object(mut object)) = metadata else {
return None;
};
for key in keys {
object.remove(*key);
}
(!object.is_empty()).then_some(Value::Object(object))
}
pub(crate) fn attach_provider_response_body_metadata(
metadata: Option<Value>,
provider_response_body: Option<&Value>,
@@ -148,6 +247,47 @@ pub(crate) fn attach_provider_response_body_metadata(
attach_provider_actual_service_tier_metadata(metadata, actual_service_tier.as_deref())
}
/// Refreshes the response-derived tier for a terminal snapshot. Complete response objects are
/// authoritative even when they contain no tier (which clears a stale candidate value). Capture
/// placeholders/absent bodies are not authoritative, so a terminal summary already present in
/// metadata is preserved for those cases.
pub(crate) fn refresh_provider_response_body_metadata(
metadata: Option<Value>,
provider_response_body: Option<&Value>,
) -> Option<Value> {
let is_capture_placeholder = provider_response_body
.and_then(Value::as_object)
.is_some_and(|body| {
body.get("truncated").and_then(Value::as_bool) == Some(true)
&& body.get("reason").and_then(Value::as_str) == Some("body_capture_limit_exceeded")
});
let body_is_complete_object =
provider_response_body.and_then(Value::as_object).is_some() && !is_capture_placeholder;
let actual_service_tier =
extract_provider_actual_service_tier_from_response(provider_response_body)
.and_then(|value| normalize_provider_service_tier(&value));
let Some(actual_service_tier) = actual_service_tier else {
if !body_is_complete_object {
return metadata;
}
let Some(Value::Object(mut object)) = metadata else {
return None;
};
object.remove(PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY);
return (!object.is_empty()).then_some(Value::Object(object));
};
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))
}
pub(crate) fn attach_provider_actual_service_tier_metadata(
metadata: Option<Value>,
actual_service_tier: Option<&str>,
@@ -179,6 +319,7 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
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, REQUESTED_REASONING_EFFORT_METADATA_KEY);
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);
@@ -226,6 +367,7 @@ fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map
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, REQUESTED_REASONING_EFFORT_METADATA_KEY);
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(
@@ -517,10 +659,16 @@ mod tests {
use serde_json::{json, Value};
use std::collections::BTreeMap;
use crate::{
apply_usage_body_capture_policy_to_event, UsageBodyCapturePolicy, UsageEvent,
UsageEventData, UsageEventType, UsageRequestRecordLevel,
};
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,
attach_client_request_body_metadata, 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, refresh_provider_response_body_metadata,
sanitize_usage_request_metadata, sanitize_usage_request_metadata_ref,
MAX_USAGE_REQUEST_METADATA_BYTES, MAX_USAGE_REQUEST_METADATA_DEPTH,
MAX_USAGE_REQUEST_METADATA_NODES,
@@ -830,6 +978,27 @@ mod tests {
assert_eq!(metadata, None);
}
#[test]
fn client_request_body_metadata_preserves_requested_reasoning_mapping_source() {
let updated = attach_client_request_body_metadata(
Some(json!({
"trace_id": "trace-1",
"requested_reasoning_effort": "low"
})),
Some(&json!({
"reasoning": { "effort": "XHigh" }
})),
)
.expect("metadata should remain");
assert_eq!(updated["requested_reasoning_effort"], "xhigh");
let cleared =
attach_client_request_body_metadata(Some(updated), Some(&json!({ "model": "gpt-5" })))
.expect("trace metadata should remain");
assert!(cleared.get("requested_reasoning_effort").is_none());
}
#[test]
fn provider_request_body_metadata_uses_final_provider_body_as_source_of_truth() {
let metadata = Some(json!({
@@ -880,6 +1049,75 @@ mod tests {
);
}
#[test]
fn final_provider_request_tier_survives_basic_body_capture_as_derived_metadata() {
let request_body = json!({
"model": "gpt-5",
"reasoning": { "effort": "xhigh" }
});
let provider_request_body = json!({
"model": "gpt-5",
"reasoning": { "effort": "max" },
"service_tier": "priority"
});
let request_metadata = attach_client_request_body_metadata(None, Some(&request_body));
let request_metadata = attach_provider_request_body_metadata(
request_metadata,
Some("openai:responses"),
Some("gpt-5"),
Some("gpt-5"),
Some(&provider_request_body),
);
let mut event = UsageEvent::new(
UsageEventType::Completed,
"req-final-provider-tier",
UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
request_body: Some(request_body),
provider_request_body: Some(provider_request_body),
request_metadata,
..UsageEventData::default()
},
);
apply_usage_body_capture_policy_to_event(
UsageBodyCapturePolicy {
record_level: UsageRequestRecordLevel::Basic,
max_request_body_bytes: Some(1024),
max_response_body_bytes: Some(1024),
},
&mut event,
);
assert_eq!(event.data.request_body, None);
assert_eq!(event.data.provider_request_body, None);
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("requested_reasoning_effort")),
Some(&json!("xhigh"))
);
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("provider_reasoning_effort")),
Some(&json!("max"))
);
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("provider_service_tier")),
Some(&json!("priority"))
);
}
#[test]
fn provider_response_metadata_preserves_terminal_actual_service_tier() {
let metadata = attach_provider_response_body_metadata(
@@ -902,6 +1140,39 @@ mod tests {
);
}
#[test]
fn terminal_response_refresh_replaces_stale_actual_tier() {
let metadata = refresh_provider_response_body_metadata(
Some(json!({"provider_actual_service_tier": "priority"})),
Some(&json!({"service_tier": "Default"})),
)
.expect("actual tier should remain");
assert_eq!(metadata["provider_actual_service_tier"], "default");
let metadata = refresh_provider_response_body_metadata(
Some(json!({
"trace_id": "trace-1",
"provider_actual_service_tier": "priority"
})),
Some(&json!({"id": "response-without-tier"})),
)
.expect("un-tiered complete response should remain auditable");
assert_eq!(metadata, json!({"trace_id": "trace-1"}));
}
#[test]
fn terminal_response_refresh_preserves_summary_over_capture_placeholder() {
let metadata = refresh_provider_response_body_metadata(
Some(json!({"provider_actual_service_tier": "priority"})),
Some(&json!({
"truncated": true,
"reason": "body_capture_limit_exceeded"
})),
)
.expect("summary should survive placeholder capture");
assert_eq!(metadata["provider_actual_service_tier"], "priority");
}
#[test]
fn terminal_summary_tier_uses_the_same_normalized_metadata_field() {
let metadata = attach_provider_actual_service_tier_metadata(
crates/aether-usage/runtime/src/runtime.rs
+164 -8
View File
@@ -13,7 +13,12 @@ use tokio::sync::mpsc;
use tracing::{info, warn};
use crate::executor::spawn_on_usage_background_runtime;
use crate::request_metadata::attach_provider_response_body_metadata;
use crate::request_metadata::{
attach_client_request_body_metadata, attach_provider_request_body_metadata,
attach_provider_response_body_metadata, clear_client_request_body_metadata,
clear_provider_request_body_metadata, request_body_derived_facts_action,
RequestBodyDerivedFactsAction,
};
use crate::worker::{
build_usage_queue_worker_with_record_gate, UsageWorkerControl, UsageWorkerObservation,
};
@@ -1264,6 +1269,7 @@ impl UsageRuntime {
where
T: UsageRuntimeAccess,
{
preserve_request_facts(event);
preserve_provider_response_facts(event);
match self.cached_body_capture_policy(data).await {
Ok(policy) => apply_usage_body_capture_policy_to_event(policy, event),
@@ -1734,6 +1740,44 @@ impl UsageRuntime {
}
}
fn preserve_request_facts(event: &mut UsageEvent) {
let data = &mut event.data;
match request_body_derived_facts_action(data.request_body.as_ref(), data.request_body_state) {
RequestBodyDerivedFactsAction::Refresh => {
data.request_metadata = attach_client_request_body_metadata(
data.request_metadata.take(),
data.request_body.as_ref(),
);
}
RequestBodyDerivedFactsAction::Clear => {
data.request_metadata =
clear_client_request_body_metadata(data.request_metadata.take());
}
RequestBodyDerivedFactsAction::Preserve => {}
}
match request_body_derived_facts_action(
data.provider_request_body.as_ref(),
data.provider_request_body_state,
) {
RequestBodyDerivedFactsAction::Refresh => {
data.request_metadata = attach_provider_request_body_metadata(
data.request_metadata.take(),
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(),
);
}
RequestBodyDerivedFactsAction::Clear => {
data.request_metadata =
clear_provider_request_body_metadata(data.request_metadata.take());
}
RequestBodyDerivedFactsAction::Preserve => {}
}
}
fn preserve_provider_response_facts(event: &mut UsageEvent) {
let metadata = event.data.request_metadata.take();
event.data.request_metadata =
@@ -2989,7 +3033,9 @@ mod tests {
use aether_data_contracts::repository::settlement::{
StoredUsageSettlement, UsageSettlementInput,
};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::repository::usage::{
StoredRequestUsageAudit, UpsertUsageRecord, UsageBodyCaptureState,
};
use aether_data_contracts::DataLayerError;
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeQueueStore, RuntimeState};
use async_trait::async_trait;
@@ -2998,9 +3044,10 @@ mod tests {
use tokio::time::{sleep, timeout, Duration};
use super::{
preserve_provider_response_facts, LifecycleEventCoalescer, UsageBillingEventEnricher,
UsageBodyCapturePolicy, UsageEnqueueRetryDispatcher, UsageRequestRecordLevel,
UsageRuntimeAccess, UsageWorkerObservation, UsageWorkerSupervisorState,
preserve_provider_response_facts, preserve_request_facts, LifecycleEventCoalescer,
UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageEnqueueRetryDispatcher,
UsageRequestRecordLevel, UsageRuntimeAccess, UsageWorkerObservation,
UsageWorkerSupervisorState,
};
use crate::worker::ManualProxyNodeCounter;
use crate::{
@@ -5816,6 +5863,76 @@ mod tests {
);
}
#[test]
fn preserve_request_facts_ignores_post_capture_truncation_placeholders() {
let truncated = json!({
"truncated": true,
"reason": "body_capture_limit_exceeded"
});
let mut event = UsageEvent::new(
UsageEventType::Completed,
"req-truncated-preserve",
UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
request_body: Some(truncated.clone()),
request_body_state: Some(UsageBodyCaptureState::Truncated),
provider_request_body: Some(truncated),
provider_request_body_state: Some(UsageBodyCaptureState::Truncated),
request_metadata: Some(json!({
"requested_reasoning_effort": "xhigh",
"provider_reasoning_effort": "max",
"provider_service_tier": "priority"
})),
..UsageEventData::default()
},
);
preserve_request_facts(&mut event);
let metadata = event
.data
.request_metadata
.as_ref()
.expect("metadata remains");
assert_eq!(metadata["requested_reasoning_effort"], "xhigh");
assert_eq!(metadata["provider_reasoning_effort"], "max");
assert_eq!(metadata["provider_service_tier"], "priority");
}
#[test]
fn preserve_request_facts_clears_stale_facts_when_final_bodies_are_missing() {
let mut event = UsageEvent::new(
UsageEventType::Completed,
"req-missing-final-body",
UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
request_body_state: Some(UsageBodyCaptureState::None),
provider_request_body_state: Some(UsageBodyCaptureState::None),
request_metadata: Some(json!({
"requested_reasoning_effort": "xhigh",
"provider_reasoning_effort": "max",
"provider_service_tier": "priority",
"provider_actual_service_tier": "priority"
})),
..UsageEventData::default()
},
);
preserve_request_facts(&mut event);
let metadata = event
.data
.request_metadata
.as_ref()
.expect("response audit fact remains");
assert!(metadata.get("requested_reasoning_effort").is_none());
assert!(metadata.get("provider_reasoning_effort").is_none());
assert!(metadata.get("provider_service_tier").is_none());
assert_eq!(metadata["provider_actual_service_tier"], "priority");
}
#[test]
fn basic_request_record_level_strips_body_capture_but_preserves_derived_fields() {
let mut event = UsageEvent::new(
@@ -5826,9 +5943,16 @@ mod tests {
model: "gpt-5".to_string(),
total_tokens: Some(42),
error_message: Some("upstream failed".to_string()),
request_body: Some(json!({"messages":[{"role":"user","content":"hello"}]})),
request_body: Some(json!({
"messages":[{"role":"user","content":"hello"}],
"reasoning": {"effort": "xhigh"}
})),
request_body_ref: Some("usage://request/req-basic-1/request_body".to_string()),
provider_request_body: Some(json!({"model":"gpt-5"})),
provider_request_body: Some(json!({
"model":"gpt-5",
"reasoning": {"effort": "max"},
"service_tier": "priority"
})),
provider_request_body_ref: Some(
"usage://request/req-basic-1/provider_request_body".to_string(),
),
@@ -5841,11 +5965,16 @@ mod tests {
client_response_body_ref: Some(
"usage://request/req-basic-1/client_response_body".to_string(),
),
request_metadata: Some(json!({"provider_service_tier": "priority"})),
request_metadata: Some(json!({
"requested_reasoning_effort": "low",
"provider_reasoning_effort": "medium",
"provider_service_tier": "standard"
})),
..UsageEventData::default()
},
);
preserve_request_facts(&mut event);
preserve_provider_response_facts(&mut event);
apply_usage_body_capture_policy_to_event(
UsageBodyCapturePolicy {
@@ -5865,6 +5994,33 @@ mod tests {
assert!(event.data.response_body_ref.is_none());
assert!(event.data.client_response_body.is_none());
assert!(event.data.client_response_body_ref.is_none());
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("requested_reasoning_effort"))
.and_then(serde_json::Value::as_str),
Some("xhigh")
);
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("provider_reasoning_effort"))
.and_then(serde_json::Value::as_str),
Some("max")
);
assert_eq!(
event
.data
.request_metadata
.as_ref()
.and_then(|metadata| metadata.get("provider_service_tier"))
.and_then(serde_json::Value::as_str),
Some("priority")
);
assert_eq!(
event
.data
crates/aether-usage/runtime/src/write.rs
+18 -8
View File
@@ -13,10 +13,11 @@ use crate::body_capture::{
RuntimeBodyCaptureMetadataInput,
};
use crate::request_metadata::{
attach_provider_actual_service_tier_metadata, attach_provider_request_body_metadata,
attach_provider_response_body_metadata, build_usage_request_metadata_seed,
attach_client_request_body_metadata, attach_provider_actual_service_tier_metadata,
attach_provider_request_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,
refresh_provider_response_body_metadata, sanitize_usage_request_metadata,
sanitize_usage_request_metadata_ref,
};
use crate::{
map_usage_from_response, stream_capture_terminal_state, GatewayStreamReportRequest,
@@ -696,6 +697,8 @@ fn build_terminal_usage_event_from_seed_impl(
} else {
merge_usage_request_metadata(request_metadata, audit_payload)
};
let request_metadata =
attach_client_request_body_metadata(request_metadata, request_body.as_ref());
let request_metadata = attach_provider_request_body_metadata(
request_metadata,
Some(provider_contract.as_str()),
@@ -1021,7 +1024,7 @@ pub fn build_sync_terminal_usage_seed(
status_code,
provider_response_full.as_ref(),
);
let request_metadata = attach_provider_response_body_metadata(
let request_metadata = refresh_provider_response_body_metadata(
context_seed.request_metadata,
provider_response_full.as_ref(),
);
@@ -1199,12 +1202,16 @@ pub fn build_stream_terminal_usage_seed(
missing_observed_finish,
terminal_error_message.is_some(),
);
let request_metadata = attach_provider_actual_service_tier_metadata(
let request_metadata = refresh_provider_response_body_metadata(
context_seed.request_metadata,
provider_response_full.as_ref(),
);
// The parser's terminal summary is authoritative when a response body is truncated or the
// body and summary disagree; attach it after the body refresh so it wins.
let request_metadata = attach_provider_actual_service_tier_metadata(
request_metadata,
provider_actual_service_tier.as_deref(),
);
let request_metadata =
attach_provider_response_body_metadata(request_metadata, provider_response_full.as_ref());
TerminalUsageSeed {
terminal_state,
@@ -2127,7 +2134,10 @@ fn build_runtime_request_metadata_seed_from_parts(
metadata.insert("body_size".to_string(), body_size);
}
(!metadata.is_empty()).then_some(Value::Object(metadata))
attach_client_request_body_metadata(
(!metadata.is_empty()).then_some(Value::Object(metadata)),
context_value_ref(context, "original_request_body"),
)
}
fn build_runtime_body_size_metadata(