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

feat(security): harden gateway boundaries and usage policies

Browse Source 
Consolidate subscription usage policy enforcement, privacy-safe persistence, and gateway security hardening into one reviewable change.

Includes bounded HTTP and execution envelopes, header and protocol guards, DNS and relay validation, authentication and secret projection hardening, secure backup/install paths, and regression coverage.
This commit is contained in:
elky
2026-09-04 03:45:52 +08:00
parent ddcbeb3ae9
commit 579f2c7cc1
1019 changed files with 190437 additions and 26080 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-usage/runtime/src/body_capture.rs
+216 -6
View File
@@ -118,15 +118,34 @@ impl UsageBodyCaptureEngine {
}
pub fn apply_to_event(self, event: &mut UsageEvent) {
self.apply_to_payload(UsageBodyCapturePayloadMut::from_event(event));
let force_disabled = is_sensitive_oauth_exchange(
event.data.api_format.as_deref(),
event.data.endpoint_api_format.as_deref(),
);
self.apply_to_payload(
UsageBodyCapturePayloadMut::from_event(event),
force_disabled,
);
}
pub fn apply_to_record(self, record: &mut UpsertUsageRecord) {
self.apply_to_payload(UsageBodyCapturePayloadMut::from_record(record));
let force_disabled = is_sensitive_oauth_exchange(
record.api_format.as_deref(),
record.endpoint_api_format.as_deref(),
);
self.apply_to_payload(
UsageBodyCapturePayloadMut::from_record(record),
force_disabled,
);
}
fn apply_to_payload(self, payload: UsageBodyCapturePayloadMut<'_>) {
if matches!(self.policy.record_level, UsageRequestRecordLevel::Basic) {
fn apply_to_payload(self, payload: UsageBodyCapturePayloadMut<'_>, force_disabled: bool) {
if force_disabled || matches!(self.policy.record_level, UsageRequestRecordLevel::Basic) {
let reason = if force_disabled {
"sensitive_oauth_exchange"
} else {
"request_record_level_basic"
};
disable_usage_body_capture_field(
UsageBodyField::RequestBody,
"request",
@@ -134,6 +153,7 @@ impl UsageBodyCaptureEngine {
payload.request_body_ref,
payload.request_body_state,
payload.request_metadata,
reason,
);
disable_usage_body_capture_field(
UsageBodyField::ProviderRequestBody,
@@ -142,6 +162,7 @@ impl UsageBodyCaptureEngine {
payload.provider_request_body_ref,
payload.provider_request_body_state,
payload.request_metadata,
reason,
);
disable_usage_body_capture_field(
UsageBodyField::ResponseBody,
@@ -150,6 +171,7 @@ impl UsageBodyCaptureEngine {
payload.response_body_ref,
payload.response_body_state,
payload.request_metadata,
reason,
);
disable_usage_body_capture_field(
UsageBodyField::ClientResponseBody,
@@ -158,6 +180,7 @@ impl UsageBodyCaptureEngine {
payload.client_response_body_ref,
payload.client_response_body_state,
payload.request_metadata,
reason,
);
return;
}
@@ -201,6 +224,22 @@ impl UsageBodyCaptureEngine {
}
}
fn is_sensitive_oauth_exchange(
api_format: Option<&str>,
endpoint_api_format: Option<&str>,
) -> bool {
[api_format, endpoint_api_format]
.into_iter()
.flatten()
.map(str::trim)
.any(|format| {
format.eq_ignore_ascii_case("oauth:exchange")
|| format.eq_ignore_ascii_case("provider_oauth:exchange")
|| format.eq_ignore_ascii_case("provider_oauth:local_refresh")
|| format.eq_ignore_ascii_case("vertex_ai:service_account_token")
})
}
pub fn apply_usage_body_capture_policy_to_event(
policy: UsageBodyCapturePolicy,
event: &mut UsageEvent,
@@ -222,6 +261,7 @@ fn disable_usage_body_capture_field(
body_ref: &mut Option<String>,
state: &mut Option<UsageBodyCaptureState>,
request_metadata: &mut Option<Value>,
reason: &'static str,
) {
*body = None;
*body_ref = None;
@@ -233,7 +273,7 @@ fn disable_usage_body_capture_field(
Some(UsageBodyCaptureState::Disabled),
None,
None,
Some("request_record_level_basic"),
Some(reason),
);
}
@@ -712,19 +752,189 @@ fn usage_value_kind(value: &Value) -> &'static str {
#[cfg(test)]
mod tests {
use super::{
apply_usage_body_capture_policy_to_event, apply_usage_body_capture_policy_to_record,
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};
use serde_json::{json, Map, Value};
use crate::{
UsageBodyCapturePolicy, UsageEvent, UsageEventData, UsageEventType, UsageRequestRecordLevel,
};
fn sensitive_oauth_event(api_format: &str) -> UsageEvent {
UsageEvent::new(
UsageEventType::Completed,
"oauth-sensitive-request",
UsageEventData {
provider_name: "oauth".to_string(),
model: "oauth-exchange".to_string(),
api_format: Some(api_format.to_string()),
endpoint_api_format: Some(api_format.to_string()),
request_body: Some(json!({"client_secret":"request-secret"})),
request_body_ref: Some("usage://oauth/request".to_string()),
provider_request_body: Some(json!({"refresh_token":"refresh-secret"})),
provider_request_body_ref: Some("usage://oauth/provider-request".to_string()),
response_body: Some(json!({
"access_token":"access-secret",
"refresh_token":"rotated-refresh-secret"
})),
response_body_ref: Some("usage://oauth/response".to_string()),
client_response_body: Some(json!({"access_token":"client-access-secret"})),
client_response_body_ref: Some("usage://oauth/client-response".to_string()),
..UsageEventData::default()
},
)
}
fn assert_sensitive_bodies_disabled(
request_body: &Option<Value>,
request_body_ref: &Option<String>,
request_body_state: Option<UsageBodyCaptureState>,
provider_request_body: &Option<Value>,
provider_request_body_ref: &Option<String>,
provider_request_body_state: Option<UsageBodyCaptureState>,
response_body: &Option<Value>,
response_body_ref: &Option<String>,
response_body_state: Option<UsageBodyCaptureState>,
client_response_body: &Option<Value>,
client_response_body_ref: &Option<String>,
client_response_body_state: Option<UsageBodyCaptureState>,
request_metadata: &Option<Value>,
) {
assert!(request_body.is_none());
assert!(request_body_ref.is_none());
assert_eq!(request_body_state, Some(UsageBodyCaptureState::Disabled));
assert!(provider_request_body.is_none());
assert!(provider_request_body_ref.is_none());
assert_eq!(
provider_request_body_state,
Some(UsageBodyCaptureState::Disabled)
);
assert!(response_body.is_none());
assert!(response_body_ref.is_none());
assert_eq!(response_body_state, Some(UsageBodyCaptureState::Disabled));
assert!(client_response_body.is_none());
assert!(client_response_body_ref.is_none());
assert_eq!(
client_response_body_state,
Some(UsageBodyCaptureState::Disabled)
);
let body_capture = request_metadata
.as_ref()
.and_then(|metadata| metadata.get("body_capture"))
.and_then(Value::as_object)
.expect("body capture metadata should exist");
for field in ["request", "provider_request", "response", "client_response"] {
assert_eq!(
body_capture
.get(field)
.and_then(|entry| entry.get("reason"))
.and_then(Value::as_str),
Some("sensitive_oauth_exchange")
);
}
}
#[test]
fn build_plan_body_capture_metadata_returns_none_without_base64_body() {
assert!(build_plan_body_capture_metadata(None).is_none());
}
#[test]
fn full_policy_never_captures_sensitive_oauth_event_bodies() {
let mut event = sensitive_oauth_event("oauth:exchange");
apply_usage_body_capture_policy_to_event(
UsageBodyCapturePolicy {
record_level: UsageRequestRecordLevel::Full,
},
&mut event,
);
assert_sensitive_bodies_disabled(
&event.data.request_body,
&event.data.request_body_ref,
event.data.request_body_state,
&event.data.provider_request_body,
&event.data.provider_request_body_ref,
event.data.provider_request_body_state,
&event.data.response_body,
&event.data.response_body_ref,
event.data.response_body_state,
&event.data.client_response_body,
&event.data.client_response_body_ref,
event.data.client_response_body_state,
&event.data.request_metadata,
);
}
#[test]
fn full_policy_never_captures_sensitive_oauth_record_bodies() {
let event = sensitive_oauth_event("provider_oauth:exchange");
let mut record = crate::record::build_upsert_usage_record_from_event(&event)
.expect("usage record should build");
apply_usage_body_capture_policy_to_record(
UsageBodyCapturePolicy {
record_level: UsageRequestRecordLevel::Full,
},
&mut record,
);
assert_sensitive_bodies_disabled(
&record.request_body,
&record.request_body_ref,
record.request_body_state,
&record.provider_request_body,
&record.provider_request_body_ref,
record.provider_request_body_state,
&record.response_body,
&record.response_body_ref,
record.response_body_state,
&record.client_response_body,
&record.client_response_body_ref,
record.client_response_body_state,
&record.request_metadata,
);
}
#[test]
fn full_policy_never_captures_refresh_or_service_account_token_bodies() {
for api_format in [
"provider_oauth:local_refresh",
"vertex_ai:service_account_token",
] {
let mut event = sensitive_oauth_event(api_format);
apply_usage_body_capture_policy_to_event(
UsageBodyCapturePolicy {
record_level: UsageRequestRecordLevel::Full,
},
&mut event,
);
assert_sensitive_bodies_disabled(
&event.data.request_body,
&event.data.request_body_ref,
event.data.request_body_state,
&event.data.provider_request_body,
&event.data.provider_request_body_ref,
event.data.provider_request_body_state,
&event.data.response_body,
&event.data.response_body_ref,
event.data.response_body_state,
&event.data.client_response_body,
&event.data.client_response_body_ref,
event.data.client_response_body_state,
&event.data.request_metadata,
);
}
}
#[test]
fn trim_owned_non_empty_string_preserves_clean_values_and_drops_blank_ones() {
assert_eq!(
crates/aether-usage/runtime/src/lib.rs
+10 -7
View File
@@ -24,17 +24,18 @@ pub use event::{now_ms, UsageEvent, UsageEventData, UsageEventType, USAGE_EVENT_
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,
decode_internal_report_body_base64, 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,
MAX_INTERNAL_REPORT_BODY_BYTES, 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,
@@ -46,7 +47,9 @@ pub use runtime::{
DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES,
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
};
pub use settlement::{settle_usage_if_needed, UsageSettlementWriter};
pub use settlement::{
reconcile_usage_policy_cost_for_event, settle_usage_if_needed, UsageSettlementWriter,
};
pub use standardized_usage::StandardizedUsage;
pub use usage_mapper::{map_usage, map_usage_from_response, UsageMapper};
pub use worker::{
crates/aether-usage/runtime/src/record.rs
+2 -1
View File
@@ -608,7 +608,8 @@ mod tests {
assert_eq!(
record.request_metadata,
Some(serde_json::json!({
"billing_snapshot": { "status": "complete" }
"billing_snapshot": { "status": "complete" },
"billing_snapshot_status": "complete"
}))
);
}
crates/aether-usage/runtime/src/report.rs
+129 -9
View File
@@ -1,12 +1,23 @@
use std::collections::BTreeMap;
use aether_contracts::{ExecutionStreamTerminalSummary, ExecutionTelemetry};
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
use aether_data_contracts::repository::{
gemini_file_mappings::{
GEMINI_FILE_MAPPING_MAX_DISPLAY_NAME_CHARS, GEMINI_FILE_MAPPING_MAX_FILE_NAME_CHARS,
GEMINI_FILE_MAPPING_MAX_MIME_TYPE_CHARS,
},
usage::UsageBodyCaptureState,
};
use base64::Engine as _;
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub const GEMINI_FILE_MAPPING_TTL_SECONDS: u64 = 60 * 60 * 48;
/// Maximum decoded size accepted for a body carried in an internal usage
/// report. Execution transports already cap normal response bodies at this
/// size; keeping the same ceiling here prevents a base64 field from causing a
/// second, unchecked allocation while preserving large image responses.
pub const MAX_INTERNAL_REPORT_BODY_BYTES: usize = 64 * 1024 * 1024;
const GEMINI_FILE_MAPPING_CACHE_PREFIX: &str = "gemini_files:key";
pub const STREAM_MISSING_TERMINAL_EVENT_CATEGORY: &str = "stream_missing_terminal_event";
pub const STREAM_TERMINAL_ERROR_CATEGORY: &str = "stream_terminal_error";
@@ -62,6 +73,40 @@ pub struct GatewayStreamReportRequest {
pub telemetry: Option<ExecutionTelemetry>,
}
/// Decode an internal report body only after checking the decoded-size bound.
/// The report payload itself is JSON, so the base64 text may be larger than the
/// raw body by roughly one third. Checking the encoded length first avoids
/// asking the base64 engine to allocate for an attacker-controlled oversized
/// value.
pub fn decode_internal_report_body_base64(body_base64: &str) -> Result<Vec<u8>, String> {
if body_base64.is_empty() {
return Ok(Vec::new());
}
let max_encoded_len = MAX_INTERNAL_REPORT_BODY_BYTES
.checked_add(2)
.and_then(|value| value.checked_div(3))
.and_then(|value| value.checked_mul(4))
.unwrap_or(usize::MAX);
if body_base64.len() > max_encoded_len {
return Err(format!(
"internal report body exceeds {} decoded bytes",
MAX_INTERNAL_REPORT_BODY_BYTES
));
}
let bytes = base64::engine::general_purpose::STANDARD
.decode(body_base64)
.map_err(|error| error.to_string())?;
if bytes.len() > MAX_INTERNAL_REPORT_BODY_BYTES {
return Err(format!(
"internal report body exceeds {} decoded bytes",
MAX_INTERNAL_REPORT_BODY_BYTES
));
}
Ok(bytes)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct InternalFinalizeRoute {
pub public_path: &'static str,
@@ -192,7 +237,12 @@ pub fn normalize_gemini_file_name(file_name: &str) -> Option<String> {
if file_name.is_empty() {
return None;
}
if file_name.starts_with("files/") {
let prefix_chars = usize::from(!file_name.starts_with("files/")) * "files/".len();
let allowed_input_chars = GEMINI_FILE_MAPPING_MAX_FILE_NAME_CHARS.checked_sub(prefix_chars)?;
if file_name.chars().nth(allowed_input_chars).is_some() {
return None;
}
if prefix_chars == 0 {
Some(file_name.to_string())
} else {
Some(format!("files/{file_name}"))
@@ -537,9 +587,7 @@ fn stream_capture_terminal_state_from_base64(
body_state: Option<UsageBodyCaptureState>,
) -> Option<StreamCapturedTerminalState> {
let body_base64 = body_base64?;
let bytes = base64::engine::general_purpose::STANDARD
.decode(body_base64)
.ok()?;
let bytes = decode_internal_report_body_base64(body_base64).ok()?;
let state = if let Ok(value) = serde_json::from_slice::<Value>(&bytes) {
stream_capture_terminal_state(&value)
} else {
@@ -766,6 +814,12 @@ fn maybe_push_gemini_file_mapping_entry(
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|value| {
value
.chars()
.nth(GEMINI_FILE_MAPPING_MAX_DISPLAY_NAME_CHARS)
.is_none()
})
.map(ToOwned::to_owned),
mime_type: object
.get("mimeType")
@@ -773,6 +827,12 @@ fn maybe_push_gemini_file_mapping_entry(
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.filter(|value| {
value
.chars()
.nth(GEMINI_FILE_MAPPING_MAX_MIME_TYPE_CHARS)
.is_none()
})
.map(ToOwned::to_owned),
});
}
@@ -789,9 +849,7 @@ fn extract_sync_report_body_json(payload: &GatewaySyncReportRequest) -> Option<V
}
let body_base64 = payload.body_base64.as_deref()?;
let bytes = base64::engine::general_purpose::STANDARD
.decode(body_base64)
.ok()?;
let bytes = decode_internal_report_body_base64(body_base64).ok()?;
serde_json::from_slice(&bytes).ok()
}
@@ -808,7 +866,13 @@ mod tests {
use std::collections::BTreeMap;
use aether_contracts::ExecutionStreamTerminalSummary;
use aether_data_contracts::repository::usage::UsageBodyCaptureState;
use aether_data_contracts::repository::{
gemini_file_mappings::{
GEMINI_FILE_MAPPING_MAX_DISPLAY_NAME_CHARS, GEMINI_FILE_MAPPING_MAX_FILE_NAME_CHARS,
GEMINI_FILE_MAPPING_MAX_MIME_TYPE_CHARS,
},
usage::UsageBodyCaptureState,
};
use base64::Engine as _;
use serde_json::json;
@@ -1152,6 +1216,16 @@ mod tests {
Some("files/abc123".to_string())
);
assert_eq!(normalize_gemini_file_name(" "), None);
assert!(normalize_gemini_file_name(
&"x".repeat(GEMINI_FILE_MAPPING_MAX_FILE_NAME_CHARS - "files/".len())
)
.is_some());
assert_eq!(
normalize_gemini_file_name(
&"x".repeat(GEMINI_FILE_MAPPING_MAX_FILE_NAME_CHARS - "files/".len() + 1)
),
None
);
}
#[test]
@@ -1244,6 +1318,31 @@ mod tests {
assert_eq!(entries[0].file_name, "files/base64-file");
}
#[test]
fn gemini_mapping_extraction_drops_oversized_persisted_metadata() {
let payload = GatewaySyncReportRequest {
trace_id: "trace-oversized-metadata".to_string(),
report_kind: "gemini_files_store_mapping".to_string(),
report_context: None,
status_code: 200,
headers: BTreeMap::new(),
body_json: Some(json!({
"name": "safe-file",
"displayName": "d".repeat(GEMINI_FILE_MAPPING_MAX_DISPLAY_NAME_CHARS + 1),
"mimeType": "m".repeat(GEMINI_FILE_MAPPING_MAX_MIME_TYPE_CHARS + 1),
})),
client_body_json: None,
body_base64: None,
telemetry: None,
};
let entries = extract_gemini_file_mapping_entries(&payload);
assert_eq!(entries.len(), 1);
assert_eq!(entries[0].file_name, "files/safe-file");
assert_eq!(entries[0].display_name, None);
assert_eq!(entries[0].mime_type, None);
}
#[test]
fn reads_report_request_id_from_context() {
assert_eq!(
@@ -1310,4 +1409,25 @@ mod tests {
"openai_chat_stream_error"
));
}
#[test]
fn bounded_report_body_decode_rejects_oversized_encoded_values_before_allocation() {
let encoded =
"A".repeat(((super::MAX_INTERNAL_REPORT_BODY_BYTES + 2) / 3 * 4).saturating_add(4));
let error = super::decode_internal_report_body_base64(&encoded)
.expect_err("oversized report body must be rejected before decoding");
assert!(error.contains("internal report body exceeds"));
}
#[test]
fn bounded_report_body_decode_preserves_empty_and_valid_payloads() {
assert!(super::decode_internal_report_body_base64("")
.expect("empty report body should decode")
.is_empty());
let encoded = base64::engine::general_purpose::STANDARD.encode(b"report-body");
assert_eq!(
super::decode_internal_report_body_base64(&encoded).expect("valid body should decode"),
b"report-body"
);
}
}
crates/aether-usage/runtime/src/request_metadata.rs
+35 -97
View File
@@ -134,6 +134,7 @@ pub(crate) fn retain_first_byte_request_metadata(value: Option<Value>) -> Option
| "upstream_is_stream"
| "client_session_affinity"
| "api_key_is_standalone"
| "plan_usage_reservation_token"
| "websocket_mode"
| "websocket_transport"
| "usage_available"
@@ -357,6 +358,8 @@ fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<St
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, "plan_usage_reservation_token");
copy_bool(source, target, "plan_usage_reservation_deferred");
copy_bool(source, target, WEBSOCKET_MODE_METADATA_KEY);
copy_non_empty_string(source, target, WEBSOCKET_TRANSPORT_METADATA_KEY);
copy_bool(source, target, USAGE_AVAILABLE_METADATA_KEY);
@@ -417,6 +420,8 @@ fn move_allowed_metadata_fields(mut source: Map<String, Value>, target: &mut Map
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, "plan_usage_reservation_token");
remove_bool(&mut source, target, "plan_usage_reservation_deferred");
remove_bool(&mut source, target, WEBSOCKET_MODE_METADATA_KEY);
remove_non_empty_string(&mut source, target, WEBSOCKET_TRANSPORT_METADATA_KEY);
remove_bool(&mut source, target, USAGE_AVAILABLE_METADATA_KEY);
@@ -739,8 +744,7 @@ mod tests {
build_usage_request_metadata_seed, merge_usage_request_metadata,
merge_usage_request_metadata_owned, refresh_provider_response_body_metadata,
retain_first_byte_request_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,
sanitize_usage_request_metadata_ref,
};
fn sample_plan() -> ExecutionPlan {
@@ -767,38 +771,8 @@ mod tests {
}
}
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",
@@ -844,8 +818,8 @@ mod tests {
"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(),
"stage_timings_ms": {"planning": 12},
"db_timings_ms": {"query": "SELECT credential"},
"original_headers": {"authorization": "Bearer secret"},
"original_request_body": {"messages": []},
"provider_request_headers": {"authorization": "Bearer secret"},
@@ -858,7 +832,7 @@ mod tests {
json!({
"trace_id": "trace-1",
"client_ip": "203.0.113.8",
"user_agent": "Claude-Code/1.0",
"client_family": "claude_code",
"client_requested_stream": false,
"upstream_is_stream": true,
"api_key_is_standalone": true,
@@ -884,9 +858,7 @@ mod tests {
"routing_candidate_skip_reason": "provider_request_body_build_failed",
"routing_failure_diagnostic": {
"kind": "request_body_build",
"path": "$.reasoning.summary",
"message": "invalid reasoning summary",
"safe_to_show": true
"path": "$.reasoning.summary"
},
"rate_multiplier": 1.25,
"is_free_tier": false,
@@ -894,9 +866,7 @@ mod tests {
"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
"price_per_request": 0.02
})
);
}
@@ -922,8 +892,7 @@ mod tests {
"client_ip": "203.0.113.8",
"request_path": "/v1/chat/completions",
"request_path_and_query": "/v1/chat/completions",
"upstream_is_stream": true,
"proxy": {"mode": "manual", "node_id": "proxy-1"}
"upstream_is_stream": true
})
);
}
@@ -946,6 +915,20 @@ mod tests {
);
}
#[test]
fn sanitizes_plan_usage_reservation_deferred_as_a_boolean() {
let metadata = sanitize_usage_request_metadata(Some(json!({
"plan_usage_reservation_deferred": true,
})))
.expect("deferred marker should remain");
assert_eq!(metadata, json!({"plan_usage_reservation_deferred": true}));
assert!(sanitize_usage_request_metadata(Some(json!({
"plan_usage_reservation_deferred": "true",
})))
.is_none());
}
#[test]
fn sanitizes_request_path_query_metadata() {
let metadata = sanitize_usage_request_metadata(Some(json!({
@@ -966,35 +949,19 @@ mod tests {
}
#[test]
fn sanitizes_large_allowed_metadata_values_to_bounded_representations() {
let metadata = sanitize_usage_request_metadata(Some(json!({
fn rejects_oversized_tokens_and_unknown_nested_objects() {
assert!(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",
}))
);
.is_none());
}
#[test]
fn sanitizes_request_metadata_preserves_tls_fingerprint() {
let metadata = sanitize_usage_request_metadata(Some(json!({
fn sanitizes_request_metadata_drops_tls_fingerprint() {
assert!(sanitize_usage_request_metadata(Some(json!({
"tls_fingerprint": {
"incoming": {
"source": "forwarded_header",
@@ -1011,25 +978,7 @@ mod tests {
"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
}
}
})
);
.is_none());
}
#[test]
@@ -1052,6 +1001,7 @@ mod tests {
"global_model_id": "global-model-1",
"global_model_name": "gpt-5",
"client_ip": "203.0.113.8",
"client_family": "claude_code",
"user_agent": "Claude-Code/1.0",
"billing_snapshot": {"status": "complete"},
"stage_timings_ms": {
@@ -1088,21 +1038,9 @@ mod tests {
"global_model_id": "global-model-1",
"global_model_name": "gpt-5",
"client_ip": "203.0.113.8",
"user_agent": "Claude-Code/1.0",
"client_family": "claude_code",
"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}
}
}
"billing_snapshot_status": "complete"
})
);
}
crates/aether-usage/runtime/src/runtime.rs
+19 -7
View File
@@ -27,9 +27,10 @@ use crate::worker::{
use crate::{
apply_usage_body_capture_policy_to_event, build_stream_terminal_usage_seed,
build_sync_terminal_usage_seed, build_terminal_usage_event_from_seed,
build_upsert_usage_record_from_event, settle_usage_if_needed, LifecycleUsageSeed,
StreamTerminalUsagePayloadSeed, SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed,
UsageEvent, UsageQueue, UsageRecordWriter, UsageRuntimeConfig, UsageSettlementWriter,
build_upsert_usage_record_from_event, reconcile_usage_policy_cost_for_event,
settle_usage_if_needed, LifecycleUsageSeed, StreamTerminalUsagePayloadSeed,
SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed, UsageEvent, UsageQueue,
UsageRecordWriter, UsageRuntimeConfig, UsageSettlementWriter,
};
#[async_trait]
@@ -39,8 +40,8 @@ pub trait UsageBillingEventEnricher: Send + Sync {
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum UsageRequestRecordLevel {
Basic,
#[default]
Basic,
Full,
}
@@ -55,7 +56,7 @@ pub struct UsageBodyCapturePolicy {
impl Default for UsageBodyCapturePolicy {
fn default() -> Self {
Self {
record_level: UsageRequestRecordLevel::Full,
record_level: UsageRequestRecordLevel::Basic,
}
}
}
@@ -4130,9 +4131,9 @@ impl UsageRuntime {
event_name = "usage_body_capture_policy_read_failed",
log_type = "event",
request_id = %event.request_id,
fallback = "default",
fallback = "basic",
error = %err,
"usage runtime failed to read body capture policy; keeping default capture"
"usage runtime failed to read body capture policy; disabling body capture"
);
apply_usage_body_capture_policy_to_event(UsageBodyCapturePolicy::default(), event);
}
@@ -4765,6 +4766,17 @@ impl UsageRuntime {
where
T: UsageRuntimeAccess,
{
if let Err(err) = reconcile_usage_policy_cost_for_event(data, event).await {
warn!(
event_name = "usage_event_cost_reconciliation_failed",
log_type = "event",
usage_event_type = ?event.event_type,
request_id = %event.request_id,
error = %err,
"usage runtime failed to reconcile plan cost before direct usage upsert"
);
return false;
}
match build_upsert_usage_record_from_event(event) {
Ok(record) => match catch_usage_writer_panic(
"direct usage upsert",
crates/aether-usage/runtime/src/settlement.rs
+388 -17
View File
@@ -1,36 +1,144 @@
use std::sync::{Arc, OnceLock};
use aether_data_contracts::repository::settlement::{StoredUsageSettlement, UsageSettlementInput};
use aether_data_contracts::repository::billing::nonnegative_usd_to_usage_policy_cost_units;
use aether_data_contracts::repository::settlement::{
ReconcileUsagePolicyCostInput, StoredUsagePolicyCostReservation, StoredUsageSettlement,
UsagePolicyCostReservationState, UsageSettlementInput,
};
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use aether_data_contracts::repository::usage::PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY;
use aether_data_contracts::{DataLayerError, DataLayerError::InvalidInput};
use async_trait::async_trait;
use crate::event::{UsageEvent, UsageEventType};
use crate::keyed_lock::KeyedAsyncLockPool;
#[async_trait]
pub trait UsageSettlementWriter: Send + Sync {
fn has_usage_settlement_writer(&self) -> bool;
async fn reconcile_usage_policy_cost(
&self,
_input: ReconcileUsagePolicyCostInput,
) -> Result<Option<StoredUsagePolicyCostReservation>, DataLayerError> {
Ok(None)
}
async fn settle_usage(
&self,
input: UsageSettlementInput,
) -> Result<Option<StoredUsageSettlement>, DataLayerError>;
}
pub async fn reconcile_usage_policy_cost_for_event(
writer: &dyn UsageSettlementWriter,
event: &UsageEvent,
) -> Result<(), DataLayerError> {
if !writer.has_usage_settlement_writer() {
return Ok(());
}
let terminal_state = match event.event_type {
UsageEventType::Completed => UsagePolicyCostReservationState::Finalized,
UsageEventType::Failed | UsageEventType::Cancelled => {
UsagePolicyCostReservationState::Released
}
UsageEventType::Pending | UsageEventType::Streaming => return Ok(()),
};
if plan_usage_reservation_reconciliation_is_deferred(event.data.request_metadata.as_ref()) {
return Ok(());
}
let Some(subject_id) = event.data.user_id.as_deref().and_then(non_empty_trimmed) else {
return Ok(());
};
let Some(reservation_token) = event_usage_policy_reservation_token(event) else {
return Ok(());
};
let actual_cost_units = if terminal_state == UsagePolicyCostReservationState::Finalized {
let actual_cost_usd = event.data.actual_total_cost_usd.ok_or_else(|| {
InvalidInput(
"completed usage event with a plan reservation token is missing actual cost"
.to_string(),
)
})?;
nonnegative_usd_to_usage_policy_cost_units(finite_cost(actual_cost_usd)?.max(0.0))
.ok_or_else(|| {
InvalidInput("usage policy settlement cost exceeds the supported range".to_string())
})?
} else {
0
};
let _ = writer
.reconcile_usage_policy_cost(ReconcileUsagePolicyCostInput {
request_id: event.request_id.clone(),
subject_id: subject_id.to_string(),
reservation_token: reservation_token.to_string(),
actual_cost_units,
terminal_state,
finalized_at_unix_secs: event.timestamp_ms / 1_000,
})
.await?;
Ok(())
}
pub async fn settle_usage_if_needed(
writer: &dyn UsageSettlementWriter,
usage: &StoredRequestUsageAudit,
) -> Result<(), DataLayerError> {
if !writer.has_usage_settlement_writer() || usage.billing_status != "pending" {
if !writer.has_usage_settlement_writer() {
return Ok(());
}
if !matches!(usage.status.as_str(), "completed" | "failed") {
if !matches!(usage.status.as_str(), "completed" | "failed" | "cancelled") {
return Ok(());
}
let finalized_at_unix_secs = usage
.finalized_at_unix_secs
.or(Some(usage.updated_at_unix_secs));
let settlement_key = usage_settlement_lock_key_for_usage(usage);
let settlement_lock = usage_settlement_lock(&settlement_key);
let _guard = settlement_lock.lock().await;
// Cost reservations are tied to a server-issued per-request token. Legacy usage rows do not
// have that token, so they must continue through wallet settlement without touching a cost
// reservation selected only by the client-visible request id.
if !plan_usage_reservation_reconciliation_is_deferred(usage.request_metadata.as_ref()) {
if let (Some(subject_id), Some(reservation_token)) = (
usage.user_id.as_deref().and_then(non_empty_trimmed),
usage_policy_reservation_token(usage),
) {
let (terminal_state, actual_cost_units) = if usage.status == "completed" {
(
UsagePolicyCostReservationState::Finalized,
nonnegative_usd_to_usage_policy_cost_units(
finite_cost(usage.actual_total_cost_usd)?.max(0.0),
)
.ok_or_else(|| {
InvalidInput(
"usage policy settlement cost exceeds the supported range".to_string(),
)
})?,
)
} else {
(UsagePolicyCostReservationState::Released, 0)
};
let _ = writer
.reconcile_usage_policy_cost(ReconcileUsagePolicyCostInput {
request_id: usage.request_id.clone(),
subject_id: subject_id.to_string(),
reservation_token: reservation_token.to_string(),
actual_cost_units,
terminal_state,
finalized_at_unix_secs: finalized_at_unix_secs
.unwrap_or(usage.updated_at_unix_secs),
})
.await?;
}
}
if usage.status == "cancelled" || usage.billing_status != "pending" {
return Ok(());
}
let input = UsageSettlementInput {
request_id: usage.request_id.clone(),
user_id: usage.user_id.clone(),
@@ -43,31 +151,36 @@ pub async fn settle_usage_if_needed(
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 plan_usage_reservation_reconciliation_is_deferred(metadata: Option<&serde_json::Value>) -> bool {
metadata
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get(PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY))
.and_then(serde_json::Value::as_bool)
.unwrap_or(false)
}
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) {
fn usage_settlement_lock_key_for_usage(usage: &StoredRequestUsageAudit) -> String {
if usage_api_key_is_standalone(usage) {
if let Some(api_key_id) = usage.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) {
if let Some(user_id) = usage.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) {
if let Some(api_key_id) = usage.api_key_id.as_deref().and_then(non_empty_trimmed) {
return format!("api-key:{api_key_id}");
}
format!("request:{}", input.request_id.trim())
format!("request:{}", usage.request_id.trim())
}
fn non_empty_trimmed(value: &str) -> Option<&str> {
@@ -84,6 +197,27 @@ fn usage_api_key_is_standalone(usage: &StoredRequestUsageAudit) -> bool {
.unwrap_or(false)
}
fn usage_policy_reservation_token(usage: &StoredRequestUsageAudit) -> Option<&str> {
usage
.request_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("plan_usage_reservation_token"))
.and_then(serde_json::Value::as_str)
.and_then(non_empty_trimmed)
}
fn event_usage_policy_reservation_token(event: &UsageEvent) -> Option<&str> {
event
.data
.request_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("plan_usage_reservation_token"))
.and_then(serde_json::Value::as_str)
.and_then(non_empty_trimmed)
}
fn finite_cost(value: f64) -> Result<f64, DataLayerError> {
if value.is_finite() {
Ok(value)
@@ -100,16 +234,24 @@ mod tests {
use std::sync::Mutex;
use std::time::Duration;
use super::{settle_usage_if_needed, UsageSettlementWriter};
use aether_data_contracts::repository::settlement::UsageSettlementInput;
use super::{
reconcile_usage_policy_cost_for_event, settle_usage_if_needed, UsageSettlementWriter,
};
use aether_data_contracts::repository::settlement::{
ReconcileUsagePolicyCostInput, StoredUsagePolicyCostReservation,
UsagePolicyCostReservationState, UsageSettlementInput,
};
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
use async_trait::async_trait;
use serde_json::json;
use crate::{UsageEvent, UsageEventData, UsageEventType};
#[derive(Default)]
struct TestSettlementWriter {
has_writer: bool,
inputs: Mutex<Vec<UsageSettlementInput>>,
reconciliations: Mutex<Vec<ReconcileUsagePolicyCostInput>>,
}
#[derive(Default)]
@@ -125,6 +267,18 @@ mod tests {
self.has_writer
}
async fn reconcile_usage_policy_cost(
&self,
input: ReconcileUsagePolicyCostInput,
) -> Result<Option<StoredUsagePolicyCostReservation>, aether_data_contracts::DataLayerError>
{
self.reconciliations
.lock()
.expect("reconciliation inputs lock")
.push(input);
Ok(None)
}
async fn settle_usage(
&self,
input: UsageSettlementInput,
@@ -166,7 +320,7 @@ mod tests {
}
fn sample_usage() -> StoredRequestUsageAudit {
StoredRequestUsageAudit::new(
let mut usage = StoredRequestUsageAudit::new(
"usage-1".to_string(),
"req-1".to_string(),
Some("user-1".to_string()),
@@ -204,7 +358,11 @@ mod tests {
200,
None,
)
.expect("usage should build")
.expect("usage should build");
usage.request_metadata = Some(json!({
"plan_usage_reservation_token": "token-1"
}));
usage
}
#[tokio::test]
@@ -228,10 +386,22 @@ mod tests {
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);
drop(inputs);
let reconciliations = writer
.reconciliations
.lock()
.expect("reconciliation inputs lock");
assert_eq!(reconciliations.len(), 1);
assert_eq!(reconciliations[0].actual_cost_units, 75_000_000);
assert_eq!(reconciliations[0].reservation_token, "token-1");
assert_eq!(
reconciliations[0].terminal_state,
UsagePolicyCostReservationState::Finalized
);
}
#[tokio::test]
async fn skips_pending_cancelled_usage() {
async fn releases_pending_cancelled_usage_without_wallet_settlement() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
@@ -246,6 +416,207 @@ mod tests {
let inputs = writer.inputs.lock().expect("settlement inputs lock");
assert!(inputs.is_empty());
drop(inputs);
let reconciliations = writer
.reconciliations
.lock()
.expect("reconciliation inputs lock");
assert_eq!(reconciliations.len(), 1);
assert_eq!(reconciliations[0].actual_cost_units, 0);
assert_eq!(
reconciliations[0].terminal_state,
UsagePolicyCostReservationState::Released
);
}
#[tokio::test]
async fn releases_failed_usage_before_void_settlement() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let mut usage = sample_usage();
usage.status = "failed".to_string();
settle_usage_if_needed(&writer, &usage)
.await
.expect("failed usage should settle");
assert_eq!(
writer.inputs.lock().expect("settlement inputs lock").len(),
1
);
let reconciliations = writer
.reconciliations
.lock()
.expect("reconciliation inputs lock");
assert_eq!(
reconciliations[0].terminal_state,
UsagePolicyCostReservationState::Released
);
assert_eq!(reconciliations[0].actual_cost_units, 0);
}
#[tokio::test]
async fn skips_cost_reconciliation_for_legacy_usage_without_token() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let mut usage = sample_usage();
usage.request_metadata = None;
settle_usage_if_needed(&writer, &usage)
.await
.expect("legacy usage should still settle its wallet charge");
assert_eq!(
writer
.reconciliations
.lock()
.expect("reconciliation lock")
.len(),
0
);
assert_eq!(
writer.inputs.lock().expect("settlement inputs lock").len(),
1
);
}
#[tokio::test]
async fn ignores_blank_reservation_token_metadata() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let mut usage = sample_usage();
usage.request_metadata = Some(json!({
"plan_usage_reservation_token": " "
}));
settle_usage_if_needed(&writer, &usage)
.await
.expect("blank token should be treated as legacy usage");
assert!(writer
.reconciliations
.lock()
.expect("reconciliation lock")
.is_empty());
assert_eq!(
writer.inputs.lock().expect("settlement inputs lock").len(),
1
);
}
#[tokio::test]
async fn event_reconciliation_requires_enriched_cost_and_preserves_server_token() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let mut event = UsageEvent::new(
UsageEventType::Completed,
"shared-trace",
UsageEventData {
user_id: Some("user-1".to_string()),
provider_name: "openai".to_string(),
model: "gpt-5".to_string(),
request_metadata: Some(json!({
"plan_usage_reservation_token": "server-token"
})),
..UsageEventData::default()
},
);
assert!(matches!(
reconcile_usage_policy_cost_for_event(&writer, &event).await,
Err(aether_data_contracts::DataLayerError::InvalidInput(_))
));
assert!(writer
.reconciliations
.lock()
.expect("reconciliations lock")
.is_empty());
event.data.actual_total_cost_usd = Some(1.25);
reconcile_usage_policy_cost_for_event(&writer, &event)
.await
.expect("enriched terminal event should reconcile");
let reconciliations = writer.reconciliations.lock().expect("reconciliations lock");
assert_eq!(reconciliations.len(), 1);
assert_eq!(reconciliations[0].request_id, "shared-trace");
assert_eq!(reconciliations[0].reservation_token, "server-token");
assert_eq!(reconciliations[0].actual_cost_units, 125_000_000);
}
#[tokio::test]
async fn deferred_event_keeps_cost_reservation_without_requiring_actual_cost() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let event = UsageEvent::new(
UsageEventType::Completed,
"possibly-sent-request",
UsageEventData {
user_id: Some("user-1".to_string()),
provider_name: "openai".to_string(),
model: "gpt-5".to_string(),
request_metadata: Some(json!({
"plan_usage_reservation_token": "server-token",
"plan_usage_reservation_deferred": true
})),
..UsageEventData::default()
},
);
reconcile_usage_policy_cost_for_event(&writer, &event)
.await
.expect("deferred reconciliation should not require unknown actual cost");
assert!(writer
.reconciliations
.lock()
.expect("reconciliations lock")
.is_empty());
}
#[tokio::test]
async fn deferred_stored_usage_skips_cost_reconcile_but_still_settles_wallet() {
let writer = TestSettlementWriter {
has_writer: true,
..Default::default()
};
let mut usage = sample_usage();
usage.request_metadata = Some(json!({
"plan_usage_reservation_token": "server-token",
"plan_usage_reservation_deferred": true
}));
settle_usage_if_needed(&writer, &usage)
.await
.expect("wallet settlement should continue");
assert!(writer
.reconciliations
.lock()
.expect("reconciliations lock")
.is_empty());
assert_eq!(
writer.inputs.lock().expect("settlement inputs lock").len(),
1
);
}
#[test]
fn deferred_metadata_requires_a_boolean_true() {
assert!(super::plan_usage_reservation_reconciliation_is_deferred(
Some(&json!({"plan_usage_reservation_deferred": true}))
));
assert!(!super::plan_usage_reservation_reconciliation_is_deferred(
Some(&json!({"plan_usage_reservation_deferred": "true"}))
));
}
#[tokio::test]
crates/aether-usage/runtime/src/worker.rs
+51 -3
View File
@@ -15,8 +15,9 @@ use crate::runtime::{
UsageBillingEventEnricher, UsageRuntimeAccess, UsageWorkerRecordConcurrencyGate,
};
use crate::{
build_upsert_usage_record_from_event, settle_usage_if_needed, UsageEvent, UsageEventType,
UsageQueue, UsageRuntimeConfig, UsageSettlementWriter,
build_upsert_usage_record_from_event, reconcile_usage_policy_cost_for_event,
settle_usage_if_needed, UsageEvent, UsageEventType, UsageQueue, UsageRuntimeConfig,
UsageSettlementWriter,
};
const USAGE_WORKER_DB_PRESSURE_DEFER_MS: u64 = 10;
@@ -676,6 +677,7 @@ pub async fn write_event_record<T>(data: &T, event: &UsageEvent) -> Result<(), D
where
T: UsageRecordWriter + UsageSettlementWriter + Send + Sync,
{
reconcile_usage_policy_cost_for_event(data, event).await?;
let record = build_upsert_usage_record_from_event(event)?;
if let Some(stored) = data.upsert_usage_record(record).await? {
settle_usage_if_needed(data, &stored).await?;
@@ -728,7 +730,8 @@ mod tests {
use std::time::Duration;
use aether_data_contracts::repository::settlement::{
StoredUsageSettlement, UsageSettlementInput,
ReconcileUsagePolicyCostInput, StoredUsagePolicyCostReservation, StoredUsageSettlement,
UsageSettlementInput,
};
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::DataLayerError;
@@ -755,6 +758,7 @@ mod tests {
struct TestUsageStore {
records: Mutex<Vec<UpsertUsageRecord>>,
settlements: Mutex<Vec<UsageSettlementInput>>,
reconciliations: Mutex<Vec<ReconcileUsagePolicyCostInput>>,
enrich_calls: Mutex<Vec<String>>,
manual_proxy_counter_calls: AtomicUsize,
}
@@ -958,6 +962,17 @@ mod tests {
true
}
async fn reconcile_usage_policy_cost(
&self,
input: ReconcileUsagePolicyCostInput,
) -> Result<Option<StoredUsagePolicyCostReservation>, DataLayerError> {
self.reconciliations
.lock()
.expect("reconciliations lock")
.push(input);
Ok(None)
}
async fn settle_usage(
&self,
input: UsageSettlementInput,
@@ -1147,6 +1162,39 @@ mod tests {
assert_eq!(settlements[0].request_id, "req-worker-123");
}
#[tokio::test]
async fn same_request_id_terminal_events_reconcile_each_reservation_token_before_upsert() {
let store = TestUsageStore::default();
let mut first = sample_event();
first.request_id = "shared-client-trace".to_string();
first.data.actual_total_cost_usd = Some(0.25);
first.data.request_metadata = Some(serde_json::json!({
"plan_usage_reservation_token": "server-token-a"
}));
let mut second = first.clone();
second.data.actual_total_cost_usd = Some(0.75);
second.data.request_metadata = Some(serde_json::json!({
"plan_usage_reservation_token": "server-token-b"
}));
write_event_record(&store, &first)
.await
.expect("first terminal event");
write_event_record(&store, &second)
.await
.expect("second terminal event");
let reconciliations = store.reconciliations.lock().expect("reconciliations lock");
assert_eq!(reconciliations.len(), 2);
assert_eq!(reconciliations[0].request_id, "shared-client-trace");
assert_eq!(reconciliations[0].reservation_token, "server-token-a");
assert_eq!(reconciliations[0].actual_cost_units, 25_000_000);
assert_eq!(reconciliations[1].request_id, "shared-client-trace");
assert_eq!(reconciliations[1].reservation_token, "server-token-b");
assert_eq!(reconciliations[1].actual_cost_units, 75_000_000);
assert_eq!(store.records.lock().expect("records lock").len(), 2);
}
#[tokio::test]
async fn replayable_usage_write_does_not_duplicate_transport_owned_proxy_counter() {
let store = TestUsageStore::default();
crates/aether-usage/runtime/src/write.rs
+71 -107
View File
@@ -8,7 +8,6 @@ use aether_data_contracts::repository::usage::{
WEBSOCKET_MODE_METADATA_KEY, WEBSOCKET_TRANSPORT_METADATA_KEY,
};
use aether_data_contracts::DataLayerError;
use base64::Engine as _;
use serde_json::{json, Map, Value};
use crate::body_capture::{
@@ -24,11 +23,11 @@ use crate::request_metadata::{
sanitize_usage_request_metadata_ref,
};
use crate::{
map_usage_from_response, stream_capture_terminal_state, GatewayStreamReportRequest,
GatewaySyncReportRequest, StandardizedUsage, StreamCapturedTerminalState, UsageEvent,
UsageEventData, UsageEventType, STREAM_MISSING_TERMINAL_EVENT_CATEGORY,
STREAM_MISSING_TERMINAL_EVENT_MESSAGE, STREAM_TERMINAL_ERROR_CATEGORY,
STREAM_TERMINAL_ERROR_MESSAGE,
decode_internal_report_body_base64, map_usage_from_response, stream_capture_terminal_state,
GatewayStreamReportRequest, GatewaySyncReportRequest, StandardizedUsage,
StreamCapturedTerminalState, UsageEvent, UsageEventData, UsageEventType,
STREAM_MISSING_TERMINAL_EVENT_CATEGORY, STREAM_MISSING_TERMINAL_EVENT_MESSAGE,
STREAM_TERMINAL_ERROR_CATEGORY, STREAM_TERMINAL_ERROR_MESSAGE,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -2130,6 +2129,12 @@ fn build_runtime_request_metadata_seed_from_parts(
Value::String(websocket_transport),
);
}
if let Some(reservation_token) = context_string(context, "plan_usage_reservation_token") {
metadata.insert(
"plan_usage_reservation_token".to_string(),
Value::String(reservation_token),
);
}
if let Some(usage_available) = context_bool(context, USAGE_AVAILABLE_METADATA_KEY) {
metadata.insert(
USAGE_AVAILABLE_METADATA_KEY.to_string(),
@@ -2416,12 +2421,8 @@ fn sanitize_usage_event_capture_fields(mut data: UsageEventData) -> UsageEventDa
data
}
fn sanitize_usage_event_capture_fields_trusted(mut data: UsageEventData) -> UsageEventData {
data.request_headers = capture_usage_header_capture(data.request_headers);
data.provider_request_headers = capture_usage_header_capture(data.provider_request_headers);
data.response_headers = capture_usage_header_capture(data.response_headers);
data.client_response_headers = capture_usage_header_capture(data.client_response_headers);
data
fn sanitize_usage_event_capture_fields_trusted(data: UsageEventData) -> UsageEventData {
sanitize_usage_event_capture_fields(data)
}
fn sanitize_usage_event_data(mut data: UsageEventData) -> UsageEventData {
@@ -2430,10 +2431,6 @@ fn sanitize_usage_event_data(mut data: UsageEventData) -> UsageEventData {
data
}
fn capture_usage_header_capture(value: Option<Value>) -> Option<Value> {
value.map(capture_usage_storage_value)
}
fn sanitize_usage_header_capture(value: Option<Value>) -> Option<Value> {
mask_sensitive_headers_in_json_value(value).map(capture_usage_storage_value)
}
@@ -2662,29 +2659,29 @@ fn headers_to_json(headers: &BTreeMap<String, String>) -> Option<Value> {
))))
}
/// 默认敏感请求头清单。与
/// `apps/aether-gateway/src/handlers/admin/system/shared/configs.rs` 中
/// `sensitive_headers` 系统配置默认值保持一致。
const DEFAULT_SENSITIVE_HEADERS: &[&str] = &[
"authorization",
"x-api-key",
"api-key",
"x-goog-api-key",
"cookie",
"set-cookie",
"proxy-authorization",
const REDACTED_USAGE_VALUE: &str = "[redacted]";
/// Only headers whose values are protocol metadata are persisted verbatim.
/// Unknown headers are treated as credentials because providers commonly use
/// custom `X-*` names for authentication.
const SAFE_USAGE_HEADER_VALUE_NAMES: &[&str] = &[
"accept",
"accept-encoding",
"content-encoding",
"content-length",
"content-type",
"transfer-encoding",
"x-request-id",
"x-trace-id",
];
/// 判断 header 名是否属于敏感字段（大小写不敏感）。
fn is_sensitive_header(name: &str) -> bool {
let trimmed = name.trim();
DEFAULT_SENSITIVE_HEADERS
!SAFE_USAGE_HEADER_VALUE_NAMES
.iter()
.any(|candidate| trimmed.eq_ignore_ascii_case(candidate))
}
/// 对单个 header value 进行脱敏：保留前 4 + 后 4 字符，中间替换为 `****`。
/// 长度小于等于 8 时整体替换为 `****`。
fn mask_header_value(name: &str, value: &str) -> String {
if !is_sensitive_header(name) {
return value.to_string();
@@ -2692,28 +2689,16 @@ fn mask_header_value(name: &str, value: &str) -> String {
mask_sensitive_header_value(value)
}
fn mask_sensitive_header_value(value: &str) -> String {
if value.len() <= 8 {
return "****".to_string();
}
let prefix: String = value.chars().take(4).collect();
let suffix: String = value
.chars()
.rev()
.take(4)
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect();
format!("{prefix}****{suffix}")
fn mask_sensitive_header_value(_value: &str) -> String {
REDACTED_USAGE_VALUE.to_string()
}
/// 对 JSON 形式的 headers 做就地脱敏。仅当 value 是 Object 时才会处理；
/// 其它形式的值保持不变。
/// Non-object values cannot be established as a valid header map and are
/// discarded instead of being persisted verbatim.
fn mask_sensitive_headers_in_json_value(value: Option<Value>) -> Option<Value> {
let mut value = value?;
let Value::Object(map) = &mut value else {
return Some(value);
return None;
};
for (key, val) in map.iter_mut() {
if !is_sensitive_header(key) {
@@ -2941,9 +2926,7 @@ fn extract_generic_error_message_from_json(value: &Value) -> Option<String> {
fn decode_body_for_storage(body_base64: Option<&str>) -> Option<Value> {
let body_base64 = body_base64?;
let bytes = base64::engine::general_purpose::STANDARD
.decode(body_base64)
.ok()?;
let bytes = decode_internal_report_body_base64(body_base64).ok()?;
if let Some(error_body) =
aether_ai_formats::api::extract_provider_private_stream_error_body(None, &bytes)
{
@@ -3499,6 +3482,7 @@ mod tests {
parse_sse_body_for_storage, resolve_error_message, trim_owned_non_empty_string,
LifecycleUsageSeed, TerminalUsageSeed, UsageBodyRefsSeed, UsageBodyStatesSeed,
UsageRoutingSeed, UsageTerminalState, MAX_USAGE_CAPTURE_BYTES, MAX_USAGE_CAPTURE_DEPTH,
REDACTED_USAGE_VALUE,
};
use crate::{
build_upsert_usage_record_from_event, GatewayStreamReportRequest, GatewaySyncReportRequest,
@@ -3934,7 +3918,8 @@ mod tests {
.expect("pending usage should keep request metadata");
assert_eq!(metadata.get("api_key_is_standalone"), Some(&json!(true)));
assert_eq!(metadata.get("client_ip"), Some(&json!("203.0.113.8")));
assert_eq!(metadata.get("user_agent"), Some(&json!("Claude-Code/1.0")));
assert_eq!(metadata.get("client_family"), Some(&json!("claude_code")));
assert!(metadata.get("user_agent").is_none());
let body_size = metadata
.get("body_size")
.and_then(Value::as_object)
@@ -5644,14 +5629,14 @@ mod tests {
assert_eq!(
event.data.response_headers,
Some(json!({
"authorization": "Bear****oken",
"authorization": REDACTED_USAGE_VALUE,
"content-type": "application/json"
}))
);
assert_eq!(
event.data.client_response_headers,
Some(json!({
"authorization": "Bear****oken",
"authorization": REDACTED_USAGE_VALUE,
"content-type": "application/json"
}))
);
@@ -6615,26 +6600,26 @@ mod tests {
assert_eq!(
event.data.request_headers,
Some(json!({
"authorization": "Bear****oken",
"authorization": REDACTED_USAGE_VALUE,
"accept": "application/json"
}))
);
assert_eq!(
event.data.provider_request_headers,
Some(json!({
"x-api-key": "sk-p****cret"
"x-api-key": REDACTED_USAGE_VALUE
}))
);
assert_eq!(
event.data.response_headers,
Some(json!({
"set-cookie": "sess****okie"
"set-cookie": REDACTED_USAGE_VALUE
}))
);
assert_eq!(
event.data.client_response_headers,
Some(json!({
"authorization": "Bear****cret"
"authorization": REDACTED_USAGE_VALUE
}))
);
assert_eq!(
@@ -6661,14 +6646,7 @@ mod tests {
.request_metadata
.as_ref()
.and_then(|value| value.get("billing_snapshot")),
Some(&json!({
"truncated": true,
"reason": "usage_request_metadata_limits_exceeded",
"max_depth": 32,
"max_nodes": 4_000,
"max_bytes": 16 * 1024,
"value_kind": "object"
}))
None
);
}
@@ -6714,19 +6692,7 @@ mod tests {
.expect("pending record should build");
assert_eq!(record.candidate_id.as_deref(), Some("cand-1"));
assert_eq!(
record.request_metadata,
Some(json!({
"billing_snapshot": {
"truncated": true,
"reason": "usage_request_metadata_limits_exceeded",
"max_depth": 32,
"max_nodes": 4_000,
"max_bytes": 16 * 1024,
"value_kind": "object"
}
}))
);
assert_eq!(record.request_metadata, None);
}
#[test]
@@ -6769,7 +6735,7 @@ mod tests {
assert_eq!(
data.request_headers,
Some(json!({
"authorization": "Bear****cret",
"authorization": REDACTED_USAGE_VALUE,
"accept": "application/json"
}))
);
@@ -6826,10 +6792,7 @@ mod tests {
metadata.get("end_to_end_first_byte_time_ms"),
Some(&json!(10_120))
);
assert_eq!(
metadata.get("db_timings_ms"),
Some(&json!({"query_count": 2}))
);
assert_eq!(metadata.get("db_timings_ms"), None);
assert_eq!(
metadata.get("trace_id"),
Some(&json!("trace-seed-metadata-1"))
@@ -6846,23 +6809,30 @@ mod tests {
fn masks_known_sensitive_header_values() {
let token = "Bearer eyJhbGciOiJSUzI1NiJ9.payload-here.signature-tail";
let masked = mask_header_value("authorization", token);
assert!(masked.starts_with("Bear"));
assert!(masked.ends_with("tail"));
assert!(masked.contains("****"));
assert!(!masked.contains("payload-here"));
assert_eq!(masked, REDACTED_USAGE_VALUE);
// 大小写不敏感
assert_eq!(
mask_header_value("Authorization", "12345678"),
"****",
"短值整体替换为 ****",
REDACTED_USAGE_VALUE,
);
assert_eq!(
mask_header_value("X-Api-Key", "abcdefghij"),
REDACTED_USAGE_VALUE,
);
assert_eq!(mask_header_value("X-Api-Key", "abcdefghij"), "abcd****ghij",);
// 非敏感头保持原样
// Unknown custom headers are redacted by default.
assert_eq!(
mask_header_value("user-agent", "codex-tui/0.1"),
"codex-tui/0.1",
REDACTED_USAGE_VALUE,
);
assert_eq!(
mask_header_value("x-custom-auth", "tenant-secret"),
REDACTED_USAGE_VALUE,
);
assert_eq!(
mask_header_value("content-type", "application/json"),
"application/json",
);
}
@@ -6886,21 +6856,17 @@ mod tests {
.get("authorization")
.and_then(|v| v.as_str())
.expect("authorization should be string");
assert!(auth.starts_with("Bear"));
assert!(auth.contains("****"));
assert!(!auth.contains("eyJhbGciOiJSUzI1NiJ9"));
assert_eq!(auth, REDACTED_USAGE_VALUE);
let api_key = object
.get("x-api-key")
.and_then(|v| v.as_str())
.expect("x-api-key should be string");
assert!(api_key.starts_with("sk-p"));
assert!(api_key.contains("****"));
assert!(!api_key.contains("1234567890"));
assert_eq!(api_key, REDACTED_USAGE_VALUE);
assert_eq!(
object.get("user-agent").and_then(|v| v.as_str()),
Some("codex-tui/0.1"),
Some(REDACTED_USAGE_VALUE),
);
}
@@ -6924,15 +6890,13 @@ mod tests {
.get("Authorization")
.and_then(|v| v.as_str())
.expect("Authorization should be string");
assert!(auth.contains("****"));
assert!(!auth.contains("eyJhbGciOiJSUzI1NiJ9"));
assert_eq!(auth, REDACTED_USAGE_VALUE);
let cookie = object
.get("Cookie")
.and_then(|v| v.as_str())
.expect("Cookie should be string");
assert!(cookie.contains("****"));
assert!(!cookie.contains("verylongcookievalue"));
assert_eq!(cookie, REDACTED_USAGE_VALUE);
assert_eq!(
object.get("Accept").and_then(|v| v.as_str()),
@@ -6941,12 +6905,12 @@ mod tests {
}
#[test]
fn mask_sensitive_headers_passthrough_for_non_object() {
fn mask_sensitive_headers_discards_non_object() {
// None 输入返回 None
assert!(mask_sensitive_headers_in_json_value(None).is_none());
// 非 object 输入原样返回
// 非 object 不是可验证的 header map，直接丢弃。
let masked = mask_sensitive_headers_in_json_value(Some(json!("not an object")));
assert_eq!(masked, Some(json!("not an object")));
assert_eq!(masked, None);
}
#[test]