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

feat(gateway): 重构 usage 数据层、迁移系统与系统导入

Browse Source 
数据库迁移:
- 引入 baseline v2 bootstrap，空库首次启动自动初始化
- 服务启动不再自动执行迁移，需显式 `--migrate` 运行
- 新增 pending migration 检测，schema 落后时拒绝启动

Usage 数据层:
- usage body 存储外部化为独立 blob 表
- 新增 HTTP audit 表拆分存储请求/响应头与 body ref
- 后台清理任务支持 legacy body ref 元数据迁移
- usage runtime 写入迁移到专用 tokio runtime（独立线程池, 8MB 栈）

系统导入/导出:
- 支持用户、API Keys、钱包数据的完整导入
- 兼容 legacy 与 v1.3+ 两种导出格式

其他改进:
- executor outcome 增加 runtime miss 诊断上下文
- 主 tokio runtime 栈大小调整为 8MB
- 前端 provider 管理支持 base URL 配置
- dev.sh 支持 --migrate 参数
This commit is contained in:
fawney19
2026-04-13 14:01:22 +08:00
parent 3698e5a833
commit 5bb08e6aa4
106 changed files with 21736 additions and 1529 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
crates/aether-usage-runtime/src/event.rs
View File

@@ -94,18 +94,42 @@ pub struct UsageEventData {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_body_ref: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_request_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_request_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub provider_request_body_ref: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub response_body_ref: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_response_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_response_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub client_response_body_ref: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub candidate_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub candidate_index: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub key_name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub planner_kind: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub route_family: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub route_kind: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub execution_path: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub local_execution_runtime_miss_reason: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub request_metadata: Option<Value>,
}

48
crates/aether-usage-runtime/src/executor.rs Normal file
View File

@@ -0,0 +1,48 @@
use std::future::Future;
use std::sync::OnceLock;
const USAGE_BACKGROUND_RUNTIME_THREADS: usize = 2;
const USAGE_BACKGROUND_RUNTIME_STACK_BYTES: usize = 8 * 1024 * 1024;
const USAGE_BACKGROUND_RUNTIME_THREAD_NAME: &str = "aether-usage-runtime";
pub(crate) fn spawn_on_usage_background_runtime<F>(task: F) -> tokio::task::JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
usage_background_runtime().handle().spawn(task)
}
fn usage_background_runtime() -> &'static tokio::runtime::Runtime {
static RUNTIME: OnceLock<&'static tokio::runtime::Runtime> = OnceLock::new();
RUNTIME.get_or_init(|| {
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.worker_threads(USAGE_BACKGROUND_RUNTIME_THREADS)
.thread_name(USAGE_BACKGROUND_RUNTIME_THREAD_NAME)
.thread_stack_size(USAGE_BACKGROUND_RUNTIME_STACK_BYTES)
.build()
.expect("usage background runtime should build");
Box::leak(Box::new(runtime))
})
}
#[cfg(test)]
mod tests {
use super::spawn_on_usage_background_runtime;
#[tokio::test]
async fn usage_background_runtime_runs_on_dedicated_named_threads() {
let thread_name = spawn_on_usage_background_runtime(async move {
std::thread::current()
.name()
.unwrap_or_default()
.to_string()
})
.await
.expect("background task should complete");
assert_eq!(thread_name, "aether-usage-runtime");
}
}

14
crates/aether-usage-runtime/src/lib.rs
View File

@@ -1,5 +1,6 @@
pub mod config;
pub mod event;
mod executor;
pub mod queue;
pub mod record;
pub mod report;
@@ -38,9 +39,14 @@ pub use worker::{
UsageQueueWorker, UsageRecordWriter,
};
pub use write::{
build_pending_usage_record, build_stream_terminal_usage_event,
build_stream_terminal_usage_outcome, build_streaming_usage_record,
build_lifecycle_usage_seed, build_pending_usage_record, build_pending_usage_record_from_seed,
build_stream_terminal_usage_event, build_stream_terminal_usage_outcome,
build_stream_terminal_usage_payload_seed, build_stream_terminal_usage_seed,
build_streaming_usage_record, build_streaming_usage_record_from_seed,
build_sync_terminal_usage_event, build_sync_terminal_usage_outcome,
build_terminal_usage_event_from_outcome, build_usage_event_data_seed, TerminalUsageOutcome,
UsageTerminalState,
build_sync_terminal_usage_payload_seed, build_sync_terminal_usage_seed,
build_terminal_usage_context_seed, build_terminal_usage_event_from_outcome,
build_terminal_usage_event_from_seed, build_usage_event_data_seed, LifecycleUsageSeed,
StreamTerminalUsagePayloadSeed, SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed,
TerminalUsageOutcome, TerminalUsageSeed, UsageTerminalState,
};

64
crates/aether-usage-runtime/src/record.rs
View File

@@ -4,6 +4,27 @@ use aether_data_contracts::DataLayerError;
use crate::request_metadata::sanitize_usage_request_metadata;
use crate::{UsageEvent, UsageEventType};
fn metadata_string(metadata: Option<&serde_json::Value>, key: &str) -> Option<String> {
metadata
.and_then(serde_json::Value::as_object)
.and_then(|object| object.get(key))
.and_then(serde_json::Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
}
fn metadata_u64(metadata: Option<&serde_json::Value>, key: &str) -> Option<u64> {
metadata
.and_then(serde_json::Value::as_object)
.and_then(|object| object.get(key))
.and_then(|value| {
value
.as_u64()
.or_else(|| value.as_i64().and_then(|number| u64::try_from(number).ok()))
})
}
pub fn build_upsert_usage_record_from_event(
event: &UsageEvent,
) -> Result<UpsertUsageRecord, DataLayerError> {
@@ -53,12 +74,51 @@ pub fn build_upsert_usage_record_from_event(
billing_status: billing_status.to_string(),
request_headers: data.request_headers,
request_body: data.request_body,
request_body_ref: empty_to_none(data.request_body_ref)
.or_else(|| metadata_string(data.request_metadata.as_ref(), "request_body_ref")),
provider_request_headers: data.provider_request_headers,
provider_request_body: data.provider_request_body,
provider_request_body_ref: empty_to_none(data.provider_request_body_ref).or_else(|| {
metadata_string(data.request_metadata.as_ref(), "provider_request_body_ref")
}),
response_headers: data.response_headers,
response_body: data.response_body,
response_body_ref: empty_to_none(data.response_body_ref)
.or_else(|| metadata_string(data.request_metadata.as_ref(), "response_body_ref")),
client_response_headers: data.client_response_headers,
client_response_body: data.client_response_body,
client_response_body_ref: empty_to_none(data.client_response_body_ref).or_else(|| {
metadata_string(data.request_metadata.as_ref(), "client_response_body_ref")
}),
candidate_id: data
.candidate_id
.or_else(|| metadata_string(data.request_metadata.as_ref(), "candidate_id")),
candidate_index: data
.candidate_index
.or_else(|| metadata_u64(data.request_metadata.as_ref(), "candidate_index")),
key_name: data
.key_name
.or_else(|| metadata_string(data.request_metadata.as_ref(), "key_name")),
planner_kind: data
.planner_kind
.or_else(|| metadata_string(data.request_metadata.as_ref(), "planner_kind")),
route_family: data
.route_family
.or_else(|| metadata_string(data.request_metadata.as_ref(), "route_family")),
route_kind: data
.route_kind
.or_else(|| metadata_string(data.request_metadata.as_ref(), "route_kind")),
execution_path: data
.execution_path
.or_else(|| metadata_string(data.request_metadata.as_ref(), "execution_path")),
local_execution_runtime_miss_reason: data.local_execution_runtime_miss_reason.or_else(
|| {
metadata_string(
data.request_metadata.as_ref(),
"local_execution_runtime_miss_reason",
)
},
),
request_metadata: sanitize_usage_request_metadata(data.request_metadata),
finalized_at_unix_secs: Some(now_unix_secs),
created_at_unix_ms: Some(now_unix_secs),
@@ -138,11 +198,11 @@ mod tests {
})
.expect("record should build");
assert_eq!(record.candidate_id.as_deref(), Some("cand-2"));
assert_eq!(record.key_name.as_deref(), Some("upstream-primary"));
assert_eq!(
record.request_metadata,
Some(serde_json::json!({
"candidate_id": "cand-2",
"key_name": "upstream-primary",
"billing_snapshot": { "status": "complete" }
}))
);

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

@@ -1,43 +1,34 @@
use aether_contracts::ExecutionPlan;
use serde_json::{Map, Value};
use serde_json::{json, Map, Value};
const MAX_USAGE_REQUEST_METADATA_DEPTH: usize = 32;
const MAX_USAGE_REQUEST_METADATA_NODES: usize = 4_000;
const MAX_USAGE_REQUEST_METADATA_BYTES: usize = 16 * 1024;
const MAX_USAGE_REQUEST_METADATA_STRING_BYTES: usize = 1_024;
pub(crate) fn build_usage_request_metadata_seed(
plan: &ExecutionPlan,
_plan: &ExecutionPlan,
context: Option<&Map<String, Value>>,
) -> Option<Value> {
let mut metadata = context.cloned().unwrap_or_default();
if !has_non_empty_string(&metadata, "candidate_id") {
if let Some(candidate_id) = plan
.candidate_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())
{
metadata.insert(
"candidate_id".to_string(),
Value::String(candidate_id.to_string()),
);
}
let mut metadata = Map::new();
if let Some(context) = context {
copy_allowed_metadata_fields(context, &mut metadata);
}
sanitize_usage_request_metadata(Some(Value::Object(metadata)))
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
pub(crate) fn merge_usage_request_metadata(
base: Option<Value>,
override_value: Option<Value>,
) -> Option<Value> {
let merged = match (base, override_value) {
(Some(Value::Object(mut base)), Some(Value::Object(override_object))) => {
for (key, value) in override_object {
base.insert(key, value);
}
Some(Value::Object(base))
}
(Some(base), None) => Some(base),
(_, Some(override_value)) => Some(override_value),
(None, None) => None,
};
sanitize_usage_request_metadata(merged)
let mut metadata = Map::new();
if let Some(Value::Object(base)) = base.as_ref() {
copy_allowed_metadata_fields(base, &mut metadata);
}
if let Some(Value::Object(override_object)) = override_value.as_ref() {
copy_allowed_metadata_fields(override_object, &mut metadata);
}
(!metadata.is_empty()).then_some(Value::Object(metadata))
}
pub(crate) fn sanitize_usage_request_metadata(value: Option<Value>) -> Option<Value> {
@@ -46,26 +37,29 @@ pub(crate) fn sanitize_usage_request_metadata(value: Option<Value>) -> Option<Va
};
let mut filtered = Map::new();
copy_non_empty_string(&object, &mut filtered, "candidate_id");
copy_number(&object, &mut filtered, "candidate_index");
copy_non_empty_string(&object, &mut filtered, "key_name");
copy_non_empty_string(&object, &mut filtered, "trace_id");
copy_non_null_value(&object, &mut filtered, "billing_snapshot");
copy_non_null_value(&object, &mut filtered, "dimensions");
copy_non_null_value(&object, &mut filtered, "billing_rule_snapshot");
copy_non_null_value(&object, &mut filtered, "scheduling_audit");
copy_number(&object, &mut filtered, "rate_multiplier");
copy_bool(&object, &mut filtered, "is_free_tier");
copy_allowed_metadata_fields(&object, &mut filtered);
(!filtered.is_empty()).then_some(Value::Object(filtered))
}
fn has_non_empty_string(object: &Map<String, Value>, key: &str) -> bool {
object
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty())
fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<String, Value>) {
copy_non_empty_string(source, target, "trace_id");
copy_number(source, target, "provider_request_body_base64_bytes");
copy_number(source, target, "provider_response_body_base64_bytes");
copy_number(source, target, "client_response_body_base64_bytes");
copy_non_null_value(source, target, "billing_snapshot");
copy_non_empty_string(source, target, "billing_snapshot_schema_version");
copy_non_empty_string(source, target, "billing_snapshot_status");
copy_non_null_value(source, target, "dimensions");
copy_non_null_value(source, target, "billing_rule_snapshot");
copy_non_null_value(source, target, "scheduling_audit");
copy_number(source, target, "rate_multiplier");
copy_bool(source, target, "is_free_tier");
copy_number(source, target, "input_price_per_1m");
copy_number(source, target, "output_price_per_1m");
copy_number(source, target, "cache_creation_price_per_1m");
copy_number(source, target, "cache_read_price_per_1m");
copy_number(source, target, "price_per_request");
}
fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
@@ -77,7 +71,10 @@ fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, V
else {
return;
};
target.insert(key.to_string(), Value::String(value.to_string()));
target.insert(
key.to_string(),
Value::String(truncate_usage_request_metadata_string(value)),
);
}
fn copy_number(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
@@ -98,18 +95,130 @@ fn copy_non_null_value(source: &Map<String, Value>, target: &mut Map<String, Val
let Some(value) = source.get(key).filter(|value| !value.is_null()) else {
return;
};
target.insert(key.to_string(), value.clone());
target.insert(
key.to_string(),
sanitize_usage_request_metadata_value(value),
);
}
fn sanitize_usage_request_metadata_value(value: &Value) -> Value {
match value {
Value::String(text) => Value::String(truncate_usage_request_metadata_string(text)),
_ if usage_request_metadata_within_limits(value) => value.clone(),
_ => truncated_usage_request_metadata_value(value),
}
}
fn truncate_usage_request_metadata_string(value: &str) -> String {
const TRUNCATED_SUFFIX: &str = "...[truncated]";
if value.len() <= MAX_USAGE_REQUEST_METADATA_STRING_BYTES {
return value.to_string();
}
let target_bytes =
MAX_USAGE_REQUEST_METADATA_STRING_BYTES.saturating_sub(TRUNCATED_SUFFIX.len());
let mut end = 0usize;
for (idx, ch) in value.char_indices() {
let next = idx + ch.len_utf8();
if next > target_bytes {
break;
}
end = next;
}
if end == 0 {
return TRUNCATED_SUFFIX.to_string();
}
format!("{}{TRUNCATED_SUFFIX}", &value[..end])
}
fn truncated_usage_request_metadata_value(value: &Value) -> Value {
json!({
"truncated": true,
"reason": "usage_request_metadata_limits_exceeded",
"max_depth": MAX_USAGE_REQUEST_METADATA_DEPTH,
"max_nodes": MAX_USAGE_REQUEST_METADATA_NODES,
"max_bytes": MAX_USAGE_REQUEST_METADATA_BYTES,
"value_kind": usage_request_metadata_value_kind(value),
})
}
fn usage_request_metadata_within_limits(value: &Value) -> bool {
let mut nodes = 0usize;
let mut estimated_bytes = 0usize;
let mut stack = vec![(value, 1usize)];
while let Some((current, depth)) = stack.pop() {
nodes = nodes.saturating_add(1);
estimated_bytes =
estimated_bytes.saturating_add(usage_request_metadata_value_size_hint(current));
if depth > MAX_USAGE_REQUEST_METADATA_DEPTH
|| nodes > MAX_USAGE_REQUEST_METADATA_NODES
|| estimated_bytes > MAX_USAGE_REQUEST_METADATA_BYTES
{
return false;
}
match current {
Value::Array(items) => {
estimated_bytes = estimated_bytes.saturating_add(items.len().saturating_mul(2));
for item in items.iter().rev() {
stack.push((item, depth + 1));
}
}
Value::Object(object) => {
estimated_bytes = estimated_bytes
.saturating_add(object.len().saturating_mul(3))
.saturating_add(
object
.keys()
.map(|key| key.len().saturating_add(2))
.sum::<usize>(),
);
for item in object.values() {
stack.push((item, depth + 1));
}
}
_ => {}
}
}
true
}
fn usage_request_metadata_value_kind(value: &Value) -> &'static str {
match value {
Value::Null => "null",
Value::Bool(_) => "bool",
Value::Number(_) => "number",
Value::String(_) => "string",
Value::Array(_) => "array",
Value::Object(_) => "object",
}
}
fn usage_request_metadata_value_size_hint(value: &Value) -> usize {
match value {
Value::Null => 4,
Value::Bool(false) => 5,
Value::Bool(true) => 4,
Value::Number(number) => number.to_string().len(),
Value::String(text) => text.len().saturating_add(2),
Value::Array(_) | Value::Object(_) => 2,
}
}
#[cfg(test)]
mod tests {
use aether_contracts::{ExecutionPlan, RequestBody};
use serde_json::json;
use serde_json::{json, Value};
use std::collections::BTreeMap;
use super::{
build_usage_request_metadata_seed, merge_usage_request_metadata,
sanitize_usage_request_metadata,
sanitize_usage_request_metadata, MAX_USAGE_REQUEST_METADATA_BYTES,
MAX_USAGE_REQUEST_METADATA_DEPTH, MAX_USAGE_REQUEST_METADATA_NODES,
};
fn sample_plan() -> ExecutionPlan {
@@ -143,14 +252,22 @@ mod tests {
"provider_id": "provider-1",
"provider_name": "OpenAI",
"model": "gpt-5",
"candidate_id": "cand-1",
"candidate_index": 2,
"key_name": "upstream-primary",
"trace_id": "trace-1",
"provider_request_body_base64_bytes": 512,
"provider_response_body_base64_bytes": 1024,
"client_response_body_base64_bytes": 2048,
"billing_snapshot": {"status": "complete"},
"billing_snapshot_schema_version": "2.0",
"billing_snapshot_status": "complete",
"dimensions": {"total_input_context": 10},
"rate_multiplier": 1.25,
"is_free_tier": false,
"input_price_per_1m": 3.0,
"output_price_per_1m": 15.0,
"cache_creation_price_per_1m": 3.75,
"cache_read_price_per_1m": 0.3,
"price_per_request": 0.02,
"original_headers": {"authorization": "Bearer secret"},
"original_request_body": {"messages": []},
"provider_request_headers": {"authorization": "Bearer secret"},
@@ -161,27 +278,60 @@ mod tests {
assert_eq!(
metadata,
json!({
"candidate_id": "cand-1",
"candidate_index": 2,
"key_name": "upstream-primary",
"trace_id": "trace-1",
"provider_request_body_base64_bytes": 512,
"provider_response_body_base64_bytes": 1024,
"client_response_body_base64_bytes": 2048,
"billing_snapshot": {"status": "complete"},
"billing_snapshot_schema_version": "2.0",
"billing_snapshot_status": "complete",
"dimensions": {"total_input_context": 10},
"rate_multiplier": 1.25,
"is_free_tier": false
"is_free_tier": false,
"input_price_per_1m": 3.0,
"output_price_per_1m": 15.0,
"cache_creation_price_per_1m": 3.75,
"cache_read_price_per_1m": 0.3,
"price_per_request": 0.02
})
);
}
#[test]
fn builds_seed_from_context_and_plan_candidate_id() {
fn sanitizes_large_allowed_metadata_values_to_bounded_representations() {
let metadata = sanitize_usage_request_metadata(Some(json!({
"trace_id": "t".repeat(2_048),
"billing_snapshot": {
"payload": "x".repeat(32 * 1024)
}
})))
.expect("metadata should remain");
assert!(metadata
.get("trace_id")
.and_then(Value::as_str)
.is_some_and(|value| value.ends_with("...[truncated]")));
assert_eq!(
metadata.get("billing_snapshot"),
Some(&json!({
"truncated": true,
"reason": "usage_request_metadata_limits_exceeded",
"max_depth": MAX_USAGE_REQUEST_METADATA_DEPTH,
"max_nodes": MAX_USAGE_REQUEST_METADATA_NODES,
"max_bytes": MAX_USAGE_REQUEST_METADATA_BYTES,
"value_kind": "object",
}))
);
}
#[test]
fn builds_seed_from_context_and_allowlisted_metadata_only() {
let metadata = build_usage_request_metadata_seed(
&sample_plan(),
Some(
json!({
"request_id": "req-1",
"candidate_index": 0,
"key_name": "upstream-primary",
"provider_id": "provider-1",
"billing_snapshot": {"status": "complete"}
})
@@ -194,9 +344,6 @@ mod tests {
assert_eq!(
metadata,
json!({
"candidate_id": "cand-1",
"candidate_index": 0,
"key_name": "upstream-primary",
"billing_snapshot": {"status": "complete"}
})
);
@@ -206,24 +353,14 @@ mod tests {
fn merges_and_filters_request_metadata() {
let metadata = merge_usage_request_metadata(
Some(json!({
"candidate_id": "cand-1",
"request_id": "req-1"
})),
Some(json!({
"candidate_index": 0,
"key_name": "upstream-primary",
"provider_name": "OpenAI"
})),
)
.expect("metadata should remain");
assert_eq!(
metadata,
json!({
"candidate_id": "cand-1",
"candidate_index": 0,
"key_name": "upstream-primary"
})
);
assert_eq!(metadata, None);
}
}

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

@@ -1,17 +1,22 @@
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
use aether_contracts::ExecutionTelemetry;
use aether_data::redis::RedisStreamRunner;
use aether_data_contracts::repository::usage::UpsertUsageRecord;
use aether_data_contracts::DataLayerError;
use async_trait::async_trait;
use tracing::warn;
use crate::executor::spawn_on_usage_background_runtime;
use crate::{
build_pending_usage_record, build_stream_terminal_usage_outcome, build_streaming_usage_record,
build_sync_terminal_usage_outcome, build_terminal_usage_event_from_outcome,
build_upsert_usage_record_from_event, build_usage_queue_worker, settle_usage_if_needed,
GatewayStreamReportRequest, GatewaySyncReportRequest, UsageEvent, UsageQueue,
UsageRecordWriter, UsageRuntimeConfig, UsageSettlementWriter, UsageTerminalState,
build_pending_usage_record_from_seed, build_stream_terminal_usage_seed,
build_streaming_usage_record_from_seed, build_sync_terminal_usage_seed,
build_terminal_usage_event_from_seed, build_upsert_usage_record_from_event,
build_usage_queue_worker, settle_usage_if_needed, LifecycleUsageSeed,
StreamTerminalUsagePayloadSeed, SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed,
UsageEvent, UsageQueue, UsageRecordWriter, UsageRuntimeConfig, UsageSettlementWriter,
};
#[async_trait]
@@ -32,6 +37,17 @@ pub struct UsageRuntime {
config: UsageRuntimeConfig,
}
struct SyncTerminalUsageTaskInput {
context_seed: TerminalUsageContextSeed,
payload_seed: SyncTerminalUsagePayloadSeed,
}
struct StreamTerminalUsageTaskInput {
context_seed: TerminalUsageContextSeed,
payload_seed: StreamTerminalUsagePayloadSeed,
cancelled: bool,
}
impl Default for UsageRuntime {
fn default() -> Self {
Self::disabled()
@@ -73,169 +89,197 @@ impl UsageRuntime {
Some(worker.spawn())
}
pub async fn record_pending<T>(
&self,
data: &T,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
) where
T: UsageRuntimeAccess,
pub fn record_pending<T>(&self, data: &T, seed: &LifecycleUsageSeed)
where
T: UsageRuntimeAccess + Clone + 'static,
{
if !self.is_enabled() {
return;
}
let now_unix_secs = now_unix_secs();
match build_pending_usage_record(plan, report_context, now_unix_secs) {
Ok(record) => {
if let Err(err) = data.upsert_usage_record(record).await {
let data = T::clone(data);
let seed = seed.clone();
let request_id = seed.request_id.clone();
spawn_on_usage_background_runtime(boxed_usage_task(async move {
let now_unix_secs = now_unix_secs();
match build_pending_usage_record_offthread(&seed, now_unix_secs).await {
Ok(record) => {
if let Err(err) = data.upsert_usage_record(record).await {
warn!(
event_name = "usage_pending_record_failed",
log_type = "event",
request_id = %request_id,
error = %err,
"usage runtime failed to record sync pending usage"
);
}
}
Err(err) => {
warn!(
event_name = "usage_pending_record_failed",
event_name = "usage_pending_build_failed",
log_type = "event",
request_id = %plan.request_id,
request_id = %request_id,
error = %err,
"usage runtime failed to record sync pending usage"
);
"usage runtime failed to build sync pending usage"
)
}
}
Err(err) => {
warn!(
event_name = "usage_pending_build_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to build sync pending usage"
)
}
}
}));
}
pub async fn record_stream_started<T>(
pub fn record_stream_started<T>(
&self,
data: &T,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
seed: &LifecycleUsageSeed,
status_code: u16,
headers: &std::collections::BTreeMap<String, String>,
telemetry: Option<&ExecutionTelemetry>,
) where
T: UsageRuntimeAccess,
T: UsageRuntimeAccess + Clone + 'static,
{
if !self.is_enabled() {
return;
}
let now_unix_secs = now_unix_secs();
match build_streaming_usage_record(
plan,
report_context,
status_code,
headers,
telemetry,
now_unix_secs,
) {
Ok(record) => {
if let Err(err) = data.upsert_usage_record(record).await {
let data = T::clone(data);
let seed = seed.clone();
let telemetry = telemetry.cloned();
let request_id = seed.request_id.clone();
spawn_on_usage_background_runtime(boxed_usage_task(async move {
let now_unix_secs = now_unix_secs();
match build_streaming_usage_record_offthread(
&seed,
status_code,
telemetry.as_ref(),
now_unix_secs,
)
.await
{
Ok(record) => {
if let Err(err) = data.upsert_usage_record(record).await {
warn!(
event_name = "usage_stream_record_failed",
log_type = "event",
request_id = %request_id,
error = %err,
"usage runtime failed to record stream usage"
);
}
}
Err(err) => {
warn!(
event_name = "usage_stream_record_failed",
event_name = "usage_stream_build_failed",
log_type = "event",
request_id = %plan.request_id,
request_id = %request_id,
error = %err,
"usage runtime failed to record stream usage"
);
"usage runtime failed to build stream usage"
)
}
}
Err(err) => {
warn!(
event_name = "usage_stream_build_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to build stream usage"
)
}
}
}));
}
pub async fn record_sync_terminal<T>(
pub fn record_sync_terminal<T>(
&self,
data: &T,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
payload: &GatewaySyncReportRequest,
context_seed: &TerminalUsageContextSeed,
payload_seed: &SyncTerminalUsagePayloadSeed,
) where
T: UsageRuntimeAccess,
T: UsageRuntimeAccess + Clone + 'static,
{
if !self.is_enabled() {
return;
}
match build_terminal_usage_event_from_outcome(build_sync_terminal_usage_outcome(
plan,
report_context,
payload,
)) {
Ok(mut event) => {
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_sync_terminal_billing_enrichment_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to enrich sync usage event with billing"
);
let runtime = self.clone();
let data = T::clone(data);
let request_id = context_seed.request_id.clone();
let input = Box::new(SyncTerminalUsageTaskInput {
context_seed: context_seed.clone(),
payload_seed: payload_seed.clone(),
});
spawn_on_usage_background_runtime(boxed_usage_task(async move {
match build_sync_terminal_usage_event_offthread(input).await {
Ok(mut event) => {
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_sync_terminal_billing_enrichment_failed",
log_type = "event",
request_id = %request_id,
error = %err,
"usage runtime failed to enrich sync usage event with billing"
);
}
runtime.enqueue_or_write_terminal(&data, event).await
}
Err(err) => {
warn!(
event_name = "usage_sync_terminal_build_failed",
log_type = "event",
request_id = %request_id,
error = %err,
"usage runtime failed to build sync terminal usage event"
)
}
self.enqueue_or_write_terminal(data, event).await
}
Err(err) => {
warn!(
event_name = "usage_sync_terminal_build_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to build sync terminal usage event"
)
}
}
}));
}
pub async fn record_stream_terminal<T>(
pub fn record_stream_terminal<T>(
&self,
data: &T,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
payload: &GatewayStreamReportRequest,
context_seed: &TerminalUsageContextSeed,
payload_seed: &StreamTerminalUsagePayloadSeed,
cancelled: bool,
) where
T: UsageRuntimeAccess,
T: UsageRuntimeAccess + Clone + 'static,
{
if !self.is_enabled() {
return;
}
let mut outcome = build_stream_terminal_usage_outcome(plan, report_context, payload);
if cancelled {
outcome.terminal_state = UsageTerminalState::Cancelled;
}
match build_terminal_usage_event_from_outcome(outcome) {
Ok(mut event) => {
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_stream_terminal_billing_enrichment_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to enrich stream usage event with billing"
);
let runtime = self.clone();
let data = T::clone(data);
let request_id = context_seed.request_id.clone();
let input = Box::new(StreamTerminalUsageTaskInput {
context_seed: context_seed.clone(),
payload_seed: payload_seed.clone(),
cancelled,
});
spawn_on_usage_background_runtime(boxed_usage_task(async move {
match build_stream_terminal_usage_event_offthread(input).await {
Ok(mut event) => {
if let Err(err) = data.enrich_usage_event(&mut event).await {
warn!(
event_name = "usage_stream_terminal_billing_enrichment_failed",
log_type = "event",
request_id = %request_id,
error = %err,
"usage runtime failed to enrich stream usage event with billing"
);
}
runtime.enqueue_or_write_terminal(&data, event).await
}
Err(err) => {
warn!(
event_name = "usage_stream_terminal_build_failed",
log_type = "event",
request_id = %request_id,
error = %err,
"usage runtime failed to build stream terminal usage event"
)
}
self.enqueue_or_write_terminal(data, event).await
}
Err(err) => {
warn!(
event_name = "usage_stream_terminal_build_failed",
log_type = "event",
request_id = %plan.request_id,
error = %err,
"usage runtime failed to build stream terminal usage event"
)
}
}));
}
pub fn submit_terminal_event<T>(&self, data: &T, event: UsageEvent)
where
T: UsageRuntimeAccess + Clone + 'static,
{
if !self.is_enabled() {
return;
}
let runtime = self.clone();
let data = T::clone(data);
spawn_on_usage_background_runtime(boxed_usage_task(async move {
runtime.record_terminal_event(&data, event).await;
}));
}
pub async fn record_terminal_event<T>(&self, data: &T, mut event: UsageEvent)
@@ -326,6 +370,74 @@ impl UsageRuntime {
}
}
async fn build_pending_usage_record_offthread(
seed: &LifecycleUsageSeed,
now_unix_secs: u64,
) -> Result<UpsertUsageRecord, DataLayerError> {
let seed = seed.clone();
tokio::task::spawn_blocking(move || build_pending_usage_record_from_seed(&seed, now_unix_secs))
.await
.map_err(join_error_to_data_layer)?
}
async fn build_streaming_usage_record_offthread(
seed: &LifecycleUsageSeed,
status_code: u16,
telemetry: Option<&ExecutionTelemetry>,
now_unix_secs: u64,
) -> Result<UpsertUsageRecord, DataLayerError> {
let seed = seed.clone();
let telemetry = telemetry.cloned();
tokio::task::spawn_blocking(move || {
build_streaming_usage_record_from_seed(
&seed,
status_code,
telemetry.as_ref(),
now_unix_secs,
)
})
.await
.map_err(join_error_to_data_layer)?
}
async fn build_sync_terminal_usage_event_offthread(
input: Box<SyncTerminalUsageTaskInput>,
) -> Result<UsageEvent, DataLayerError> {
tokio::task::spawn_blocking(move || {
build_terminal_usage_event_from_seed(build_sync_terminal_usage_seed(
input.context_seed,
input.payload_seed,
))
})
.await
.map_err(join_error_to_data_layer)?
}
async fn build_stream_terminal_usage_event_offthread(
input: Box<StreamTerminalUsageTaskInput>,
) -> Result<UsageEvent, DataLayerError> {
tokio::task::spawn_blocking(move || {
build_terminal_usage_event_from_seed(build_stream_terminal_usage_seed(
input.context_seed,
input.payload_seed,
input.cancelled,
))
})
.await
.map_err(join_error_to_data_layer)?
}
fn join_error_to_data_layer(err: tokio::task::JoinError) -> DataLayerError {
DataLayerError::UnexpectedValue(format!("usage builder task join failed: {err}"))
}
fn boxed_usage_task<F>(task: F) -> Pin<Box<dyn Future<Output = ()> + Send>>
where
F: Future<Output = ()> + Send + 'static,
{
Box::pin(task)
}
fn now_unix_secs() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)

3
crates/aether-usage-runtime/src/worker.rs
View File

@@ -7,6 +7,7 @@ use aether_data_contracts::DataLayerError;
use async_trait::async_trait;
use tracing::warn;
use crate::executor::spawn_on_usage_background_runtime;
use crate::{
build_upsert_usage_record_from_event, settle_usage_if_needed, UsageEvent, UsageQueue,
UsageRuntimeConfig, UsageSettlementWriter,
@@ -69,7 +70,7 @@ impl UsageQueueWorker {
}
pub fn spawn(self) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move { self.run_forever().await })
spawn_on_usage_background_runtime(async move { self.run_forever().await })
}
async fn run_forever(self) {

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

File diff suppressed because it is too large Load Diff