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

Refactor usage body capture and stream terminal reporting

Browse Source
This commit is contained in:
fawney19
2026-04-18 17:48:21 +08:00
parent 569242d72f
commit 3363592751
36 changed files with 2673 additions and 512 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
apps/aether-gateway/src/data/state/integrations.rs
View File

@@ -16,8 +16,10 @@ use aether_data_contracts::repository::settlement::{StoredUsageSettlement, Usage
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
use aether_data_contracts::repository::video_tasks::{StoredVideoTask, VideoTaskLookupKey};
use aether_usage_runtime::{
UsageBillingEventEnricher, UsageEvent, UsageRecordWriter, UsageRequestRecordLevel,
UsageRuntimeAccess, UsageSettlementWriter,
UsageBillingEventEnricher, UsageBodyCapturePolicy, UsageEvent, UsageRecordWriter,
UsageRequestRecordLevel, UsageRuntimeAccess, UsageSettlementWriter,
DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES,
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
};
use aether_video_tasks_core::StoredVideoTaskReadSide;
use async_trait::async_trait;
@@ -29,6 +31,8 @@ use crate::provider_transport::ProviderTransportSnapshotSource;
const REQUEST_RECORD_LEVEL_KEY: &str = "request_record_level";
const LEGACY_REQUEST_LOG_LEVEL_KEY: &str = "request_log_level";
const MAX_REQUEST_BODY_SIZE_KEY: &str = "max_request_body_size";
const MAX_RESPONSE_BODY_SIZE_KEY: &str = "max_response_body_size";
fn usage_request_record_level_from_value(value: Option<&Value>) -> UsageRequestRecordLevel {
let Some(value) = value.and_then(Value::as_str).map(str::trim) else {
@@ -47,6 +51,14 @@ fn usage_request_record_level_from_value(value: Option<&Value>) -> UsageRequestR
}
}
fn usage_body_capture_limit_from_value(value: Option<&Value>, default: usize) -> Option<usize> {
match value.and_then(Value::as_u64) {
Some(0) => None,
Some(limit) => usize::try_from(limit).ok().filter(|limit| *limit > 0),
None => Some(default),
}
}
#[async_trait]
impl RequestAuditReader for GatewayDataState {
async fn find_request_usage_audit_by_request_id(
@@ -196,14 +208,28 @@ impl UsageRuntimeAccess for GatewayDataState {
GatewayDataState::usage_worker_runner(self)
}
async fn request_record_level(&self) -> Result<UsageRequestRecordLevel, DataLayerError> {
async fn body_capture_policy(&self) -> Result<UsageBodyCapturePolicy, DataLayerError> {
let value = GatewayDataState::find_system_config_value(self, REQUEST_RECORD_LEVEL_KEY)
.await?
.or(
GatewayDataState::find_system_config_value(self, LEGACY_REQUEST_LOG_LEVEL_KEY)
.await?,
);
Ok(usage_request_record_level_from_value(value.as_ref()))
let max_request_body_size =
GatewayDataState::find_system_config_value(self, MAX_REQUEST_BODY_SIZE_KEY).await?;
let max_response_body_size =
GatewayDataState::find_system_config_value(self, MAX_RESPONSE_BODY_SIZE_KEY).await?;
Ok(UsageBodyCapturePolicy {
record_level: usage_request_record_level_from_value(value.as_ref()),
max_request_body_bytes: usage_body_capture_limit_from_value(
max_request_body_size.as_ref(),
DEFAULT_USAGE_REQUEST_BODY_CAPTURE_LIMIT_BYTES,
),
max_response_body_bytes: usage_body_capture_limit_from_value(
max_response_body_size.as_ref(),
DEFAULT_USAGE_RESPONSE_BODY_CAPTURE_LIMIT_BYTES,
),
})
}
}
@@ -327,4 +353,35 @@ mod tests {
assert_eq!(level, UsageRequestRecordLevel::Full);
}
#[tokio::test]
async fn usage_runtime_access_reads_body_capture_limits_from_system_config() {
let state = GatewayDataState::disabled().with_system_config_values_for_tests([
("max_request_body_size".to_string(), json!(1234)),
("max_response_body_size".to_string(), json!(5678)),
]);
let policy = UsageRuntimeAccess::body_capture_policy(&state)
.await
.expect("body capture policy should read");
assert_eq!(policy.record_level, UsageRequestRecordLevel::Full);
assert_eq!(policy.max_request_body_bytes, Some(1234));
assert_eq!(policy.max_response_body_bytes, Some(5678));
}
#[tokio::test]
async fn usage_runtime_access_treats_zero_body_capture_limit_as_unbounded() {
let state = GatewayDataState::disabled().with_system_config_values_for_tests([
("max_request_body_size".to_string(), json!(0)),
("max_response_body_size".to_string(), json!(0)),
]);
let policy = UsageRuntimeAccess::body_capture_policy(&state)
.await
.expect("body capture policy should read");
assert_eq!(policy.max_request_body_bytes, None);
assert_eq!(policy.max_response_body_bytes, None);
}
}