refactor: 移除独立 hub/proxy/executor/gateway crate,统一为 gateway tunnel 架构

- 删除 aether-hub、aether-proxy 独立项目及其 Dockerfile/配置
- 删除 crates/aether-executor 和 crates/aether-gateway 全部模块
- 新增 apps/ 目录作为应用入口
- 将 hub 概念重构为 gateway tunnel transport
- 将 executor 重构为 execution runtime
- 新增 tunnel.rs 合约定义和 testkit tunnel/execution_runtime 模块
- 更新 Python 服务层和测试适配新架构命名
This commit is contained in:
fawney19
2026-04-03 14:59:58 +08:00
parent ddf18fed9a
commit 8f26e1a31f
983 changed files with 103098 additions and 105837 deletions

View File

@@ -0,0 +1,57 @@
use aether_data::DataLayerError;
use crate::gateway::gateway_data::DecisionTrace;
use crate::gateway::gateway_data::GatewayDataState;
use crate::gateway::gateway_data::StoredGatewayAuthApiKeySnapshot;
use super::{read_request_usage_audit, RequestUsageAudit};
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub(crate) struct RequestAuditBundle {
pub(crate) request_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) usage: Option<RequestUsageAudit>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) decision_trace: Option<DecisionTrace>,
#[serde(skip_serializing_if = "Option::is_none")]
pub(crate) auth_snapshot: Option<StoredGatewayAuthApiKeySnapshot>,
}
pub(crate) async fn read_request_audit_bundle(
state: &GatewayDataState,
request_id: &str,
attempted_only: bool,
now_unix_secs: u64,
) -> Result<Option<RequestAuditBundle>, DataLayerError> {
let usage = read_request_usage_audit(state, request_id).await?;
let decision_trace = state
.read_decision_trace(request_id, attempted_only)
.await?;
let auth_snapshot = if let Some(usage) = usage.as_ref() {
match (
usage.usage.user_id.as_deref(),
usage.usage.api_key_id.as_deref(),
) {
(Some(user_id), Some(api_key_id)) => {
state
.read_auth_api_key_snapshot(user_id, api_key_id, now_unix_secs)
.await?
}
_ => None,
}
} else {
None
};
if usage.is_none() && decision_trace.is_none() && auth_snapshot.is_none() {
return Ok(None);
}
Ok(Some(RequestAuditBundle {
request_id: request_id.to_string(),
usage,
decision_trace,
auth_snapshot,
}))
}

View File

@@ -0,0 +1,112 @@
use aether_data::DataLayerError;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UsageRuntimeConfig {
pub enabled: bool,
pub stream_key: String,
pub consumer_group: String,
pub dlq_stream_key: String,
pub stream_maxlen: usize,
pub consumer_batch_size: usize,
pub consumer_block_ms: u64,
pub reclaim_idle_ms: u64,
pub reclaim_count: usize,
pub reclaim_interval_ms: u64,
}
impl Default for UsageRuntimeConfig {
fn default() -> Self {
Self {
enabled: false,
stream_key: "usage:events".to_string(),
consumer_group: "usage_consumers".to_string(),
dlq_stream_key: "usage:events:dlq".to_string(),
stream_maxlen: 2_000,
consumer_batch_size: 200,
consumer_block_ms: 500,
reclaim_idle_ms: 30_000,
reclaim_count: 200,
reclaim_interval_ms: 5_000,
}
}
}
impl UsageRuntimeConfig {
pub(crate) fn disabled() -> Self {
Self::default()
}
pub(crate) fn validate(&self) -> Result<(), DataLayerError> {
if !self.enabled {
return Ok(());
}
if self.stream_key.trim().is_empty() {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime stream_key cannot be empty".to_string(),
));
}
if self.consumer_group.trim().is_empty() {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime consumer_group cannot be empty".to_string(),
));
}
if self.dlq_stream_key.trim().is_empty() {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime dlq_stream_key cannot be empty".to_string(),
));
}
if self.stream_maxlen == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime stream_maxlen must be positive".to_string(),
));
}
if self.consumer_batch_size == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime consumer_batch_size must be positive".to_string(),
));
}
if self.consumer_block_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime consumer_block_ms must be positive".to_string(),
));
}
if self.reclaim_idle_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime reclaim_idle_ms must be positive".to_string(),
));
}
if self.reclaim_count == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime reclaim_count must be positive".to_string(),
));
}
if self.reclaim_interval_ms == 0 {
return Err(DataLayerError::InvalidConfiguration(
"usage runtime reclaim_interval_ms must be positive".to_string(),
));
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::UsageRuntimeConfig;
#[test]
fn disabled_config_is_valid() {
assert!(UsageRuntimeConfig::disabled().validate().is_ok());
}
#[test]
fn enabled_config_rejects_empty_stream_key() {
let config = UsageRuntimeConfig {
enabled: true,
stream_key: String::new(),
..UsageRuntimeConfig::default()
};
assert!(config.validate().is_err());
}
}

View File

@@ -0,0 +1,218 @@
use std::collections::BTreeMap;
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data::DataLayerError;
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub(crate) const USAGE_EVENT_VERSION: u8 = 1;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum UsageEventType {
Pending,
Streaming,
Completed,
Failed,
Cancelled,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
pub(crate) struct UsageEventData {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) user_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) api_key_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) username: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) api_key_name: Option<String>,
pub(crate) provider_name: String,
pub(crate) model: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) target_model: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) provider_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) provider_endpoint_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) provider_api_key_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) request_type: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) api_format: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) api_family: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) endpoint_kind: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) endpoint_api_format: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) provider_api_family: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) provider_endpoint_kind: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) has_format_conversion: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) is_stream: Option<bool>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) input_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) output_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) total_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) cache_creation_input_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) cache_read_input_tokens: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) cache_creation_cost_usd: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) cache_read_cost_usd: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) output_price_per_1m: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) total_cost_usd: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) actual_total_cost_usd: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) status_code: Option<u16>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) error_message: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) error_category: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) response_time_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) first_byte_time_ms: Option<u64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) request_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) request_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) provider_request_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) provider_request_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) response_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) response_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) client_response_headers: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) client_response_body: Option<Value>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) request_metadata: Option<Value>,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct UsageEvent {
pub(crate) event_type: UsageEventType,
pub(crate) request_id: String,
pub(crate) timestamp_ms: u64,
pub(crate) data: UsageEventData,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
struct UsageEventEnvelope {
v: u8,
#[serde(rename = "type")]
event_type: UsageEventType,
request_id: String,
timestamp_ms: u64,
data: UsageEventData,
}
impl UsageEvent {
pub(crate) fn new(
event_type: UsageEventType,
request_id: impl Into<String>,
data: UsageEventData,
) -> Self {
Self {
event_type,
request_id: request_id.into(),
timestamp_ms: now_ms(),
data,
}
}
pub(crate) fn to_stream_fields(&self) -> Result<BTreeMap<String, String>, DataLayerError> {
let payload = UsageEventEnvelope {
v: USAGE_EVENT_VERSION,
event_type: self.event_type,
request_id: self.request_id.clone(),
timestamp_ms: self.timestamp_ms,
data: self.data.clone(),
};
let payload = serde_json::to_string(&payload).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"failed to serialize usage event payload: {err}"
))
})?;
Ok(BTreeMap::from([("payload".to_string(), payload)]))
}
pub(crate) fn from_stream_fields(
fields: &BTreeMap<String, String>,
) -> Result<Self, DataLayerError> {
let payload = fields.get("payload").ok_or_else(|| {
DataLayerError::UnexpectedValue(
"usage event stream entry missing payload field".to_string(),
)
})?;
let envelope: UsageEventEnvelope = serde_json::from_str(payload).map_err(|err| {
DataLayerError::UnexpectedValue(format!(
"failed to deserialize usage event payload: {err}"
))
})?;
if envelope.v != USAGE_EVENT_VERSION {
return Err(DataLayerError::UnexpectedValue(format!(
"unsupported usage event version: {}",
envelope.v
)));
}
Ok(Self {
event_type: envelope.event_type,
request_id: envelope.request_id,
timestamp_ms: envelope.timestamp_ms,
data: envelope.data,
})
}
}
pub(crate) fn now_ms() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_millis() as u64
}
#[cfg(test)]
mod tests {
use super::{UsageEvent, UsageEventData, UsageEventType};
#[test]
fn usage_event_round_trips_through_stream_fields() {
let event = UsageEvent::new(
UsageEventType::Completed,
"req-1",
UsageEventData {
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
input_tokens: Some(10),
output_tokens: Some(20),
..UsageEventData::default()
},
);
let fields = event.to_stream_fields().expect("event should serialize");
let parsed = UsageEvent::from_stream_fields(&fields).expect("event should parse");
assert_eq!(parsed.request_id, "req-1");
assert_eq!(parsed.event_type, UsageEventType::Completed);
assert_eq!(parsed.data.total_tokens, None);
assert_eq!(parsed.data.output_tokens, Some(20));
}
}

View File

@@ -0,0 +1,69 @@
use std::time::{SystemTime, UNIX_EPOCH};
use axum::extract::{Path, Query, State};
use axum::response::IntoResponse;
use axum::Json;
use serde::Deserialize;
use serde_json::json;
use crate::gateway::AppState;
#[derive(Debug, Deserialize)]
pub(crate) struct GetRequestAuditBundleQuery {
pub(crate) attempted_only: Option<bool>,
}
pub(crate) async fn get_request_usage_audit(
State(state): State<AppState>,
Path(request_id): Path<String>,
) -> Result<Json<super::RequestUsageAudit>, axum::response::Response> {
let usage = state
.read_request_usage_audit(&request_id)
.await
.map_err(IntoResponse::into_response)?;
match usage {
Some(usage) => Ok(Json(usage)),
None => Err((
axum::http::StatusCode::NOT_FOUND,
Json(json!({
"error": {
"message": "Request usage not found",
}
})),
)
.into_response()),
}
}
pub(crate) async fn get_request_audit_bundle(
State(state): State<AppState>,
Path(request_id): Path<String>,
Query(query): Query<GetRequestAuditBundleQuery>,
) -> Result<Json<super::RequestAuditBundle>, axum::response::Response> {
let attempted_only = query.attempted_only.unwrap_or(false);
let bundle = state
.read_request_audit_bundle(&request_id, attempted_only, current_unix_secs())
.await
.map_err(IntoResponse::into_response)?;
match bundle {
Some(bundle) => Ok(Json(bundle)),
None => Err((
axum::http::StatusCode::NOT_FOUND,
Json(json!({
"error": {
"message": "Request audit bundle not found",
}
})),
)
.into_response()),
}
}
fn current_unix_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}

View File

@@ -0,0 +1,22 @@
mod bundle;
mod config;
mod event;
mod http;
mod queue;
mod read;
mod reporting;
mod runtime;
mod worker;
mod write;
pub(crate) use bundle::{read_request_audit_bundle, RequestAuditBundle};
pub use config::UsageRuntimeConfig;
pub(crate) use event::{UsageEvent, UsageEventData, UsageEventType};
pub(crate) use http::{get_request_audit_bundle, get_request_usage_audit};
pub(crate) use read::{read_request_usage_audit, RequestUsageAudit};
pub(crate) use reporting::{
spawn_sync_report, store_local_gemini_file_mapping, submit_stream_report, submit_sync_report,
GatewayStreamReportRequest, GatewaySyncReportRequest,
};
pub(crate) use runtime::UsageRuntime;
pub(crate) use write::build_upsert_usage_record_from_event;

View File

@@ -0,0 +1,161 @@
use serde_json::json;
use aether_data::redis::{
RedisConsumerGroup, RedisConsumerName, RedisStreamEntry, RedisStreamName,
RedisStreamReclaimConfig, RedisStreamRunner, RedisStreamRunnerConfig,
};
use aether_data::DataLayerError;
use super::{UsageEvent, UsageRuntimeConfig};
#[derive(Debug, Clone)]
pub(crate) struct UsageQueue {
runner: RedisStreamRunner,
config: UsageRuntimeConfig,
stream: RedisStreamName,
group: RedisConsumerGroup,
dlq_stream: RedisStreamName,
}
impl UsageQueue {
pub(crate) fn new(
runner: RedisStreamRunner,
config: UsageRuntimeConfig,
) -> Result<Self, DataLayerError> {
config.validate()?;
let tuned_runner = RedisStreamRunner::new(
runner.client().clone(),
runner.keyspace().clone(),
usage_stream_runner_config(&config),
)?;
Ok(Self {
runner: tuned_runner,
stream: RedisStreamName(config.stream_key.clone()),
group: RedisConsumerGroup(config.consumer_group.clone()),
dlq_stream: RedisStreamName(config.dlq_stream_key.clone()),
config,
})
}
pub(crate) async fn ensure_consumer_group(&self) -> Result<(), DataLayerError> {
self.runner
.ensure_consumer_group(&self.stream, &self.group, "0-0")
.await
}
pub(crate) async fn enqueue(&self, event: &UsageEvent) -> Result<String, DataLayerError> {
let fields = event.to_stream_fields()?;
self.runner
.append_fields_with_maxlen(&self.stream, &fields, Some(self.config.stream_maxlen))
.await
}
pub(crate) async fn read_group(
&self,
consumer: &RedisConsumerName,
) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
self.runner
.read_group(&self.stream, &self.group, consumer)
.await
}
pub(crate) async fn claim_stale(
&self,
consumer: &RedisConsumerName,
start_id: &str,
) -> Result<Vec<RedisStreamEntry>, DataLayerError> {
Ok(self
.runner
.claim_stale(
&self.stream,
&self.group,
consumer,
start_id,
RedisStreamReclaimConfig {
min_idle_ms: self.config.reclaim_idle_ms,
count: self.config.reclaim_count,
},
)
.await?
.entries)
}
pub(crate) async fn ack_and_delete(&self, ids: &[String]) -> Result<(), DataLayerError> {
self.runner.ack(&self.stream, &self.group, ids).await?;
self.runner.delete(&self.stream, ids).await?;
Ok(())
}
pub(crate) async fn push_dead_letter(
&self,
entry: &RedisStreamEntry,
error: &str,
) -> Result<String, DataLayerError> {
self.runner
.append_json(
&self.dlq_stream,
"payload",
&json!({
"entry_id": entry.id,
"fields": entry.fields,
"error": error,
}),
)
.await
}
}
fn usage_stream_runner_config(config: &UsageRuntimeConfig) -> RedisStreamRunnerConfig {
let read_block_ms = config.consumer_block_ms.max(1);
let command_timeout_ms = read_block_ms.saturating_add(2_000).max(5_000);
RedisStreamRunnerConfig {
command_timeout_ms: Some(command_timeout_ms),
read_block_ms: Some(read_block_ms),
read_count: config.consumer_batch_size.max(1),
}
}
#[cfg(test)]
mod tests {
use super::{usage_stream_runner_config, UsageQueue};
use crate::gateway::usage::UsageRuntimeConfig;
use aether_data::redis::{RedisClientConfig, RedisClientFactory, RedisStreamRunner};
fn sample_runner() -> RedisStreamRunner {
let config = RedisClientConfig {
url: "redis://127.0.0.1/0".to_string(),
key_prefix: Some("aether".to_string()),
};
let client = RedisClientFactory::new(config.clone())
.expect("factory should build")
.connect_lazy()
.expect("client should build");
RedisStreamRunner::new(
client,
config.keyspace(),
aether_data::redis::RedisStreamRunnerConfig::default(),
)
.expect("runner should build")
}
#[test]
fn usage_queue_applies_runtime_block_and_batch_settings() {
let config = UsageRuntimeConfig {
enabled: true,
consumer_block_ms: 750,
consumer_batch_size: 123,
..UsageRuntimeConfig::default()
};
let queue = UsageQueue::new(sample_runner(), config)
.expect("usage queue should build from runtime config");
assert_eq!(
queue.runner.config(),
usage_stream_runner_config(&queue.config)
);
assert_eq!(queue.runner.config().read_block_ms, Some(750));
assert_eq!(queue.runner.config().read_count, 123);
assert_eq!(queue.runner.config().command_timeout_ms, Some(5_000));
}
}

View File

@@ -0,0 +1,20 @@
use aether_data::repository::usage::StoredRequestUsageAudit;
use aether_data::DataLayerError;
use crate::gateway::gateway_data::GatewayDataState;
#[derive(Debug, Clone, PartialEq, serde::Serialize)]
pub(crate) struct RequestUsageAudit {
#[serde(flatten)]
pub(crate) usage: StoredRequestUsageAudit,
}
pub(crate) async fn read_request_usage_audit(
state: &GatewayDataState,
request_id: &str,
) -> Result<Option<RequestUsageAudit>, DataLayerError> {
Ok(state
.find_request_usage_by_request_id(request_id)
.await?
.map(|usage| RequestUsageAudit { usage }))
}

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,203 @@
use aether_data::repository::video_tasks::VideoTaskLookupKey;
use serde_json::{Map, Value};
use crate::gateway::AppState;
pub(crate) fn report_context_is_locally_actionable(report_context: Option<&Value>) -> bool {
let Some(context) = report_context else {
return false;
};
has_non_empty_str(context, "request_id")
&& (has_non_empty_str(context, "candidate_id")
|| has_u64(context, "candidate_index")
|| has_non_empty_str(context, "provider_id")
|| has_non_empty_str(context, "endpoint_id")
|| has_non_empty_str(context, "key_id"))
}
pub(crate) async fn resolve_locally_actionable_report_context(
state: &AppState,
report_context: Option<&Value>,
) -> Option<Value> {
let context = report_context?.clone();
if report_context_is_locally_actionable(Some(&context)) {
return Some(context);
}
if let Some(resolved) =
resolve_locally_actionable_report_context_from_request_candidates(state, &context).await
{
return Some(resolved);
}
let context = resolve_locally_actionable_report_context_from_video_task(state, &context)
.await
.unwrap_or(context);
if let Some(resolved) =
resolve_locally_actionable_report_context_from_request_candidates(state, &context).await
{
return Some(resolved);
}
report_context_is_locally_actionable(Some(&context)).then_some(context)
}
async fn resolve_locally_actionable_report_context_from_request_candidates(
state: &AppState,
context: &Value,
) -> Option<Value> {
let request_id = context
.get("request_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
let existing_candidates = state
.read_request_candidates_by_request_id(request_id)
.await
.ok()?;
if existing_candidates.len() != 1 {
return None;
}
let mut object = context.as_object()?.clone();
let candidate = &existing_candidates[0];
insert_missing_string_value(&mut object, "candidate_id", Some(&candidate.id));
if !object.contains_key("candidate_index") {
object.insert(
"candidate_index".to_string(),
Value::Number(candidate.candidate_index.into()),
);
}
insert_missing_optional_string_value(
&mut object,
"provider_id",
candidate.provider_id.as_deref(),
);
insert_missing_optional_string_value(
&mut object,
"endpoint_id",
candidate.endpoint_id.as_deref(),
);
insert_missing_optional_string_value(&mut object, "key_id", candidate.key_id.as_deref());
insert_missing_optional_string_value(&mut object, "user_id", candidate.user_id.as_deref());
insert_missing_optional_string_value(
&mut object,
"api_key_id",
candidate.api_key_id.as_deref(),
);
let resolved = Value::Object(object);
report_context_is_locally_actionable(Some(&resolved)).then_some(resolved)
}
async fn resolve_locally_actionable_report_context_from_video_task(
state: &AppState,
context: &Value,
) -> Option<Value> {
let local_task_id = context
.get("local_task_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let local_short_id = context
.get("local_short_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let task_id = context
.get("task_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let user_id = context
.get("user_id")
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty());
let task = if let Some(task_id) = local_task_id {
state
.data
.find_video_task(VideoTaskLookupKey::Id(task_id))
.await
.ok()??
} else if let Some(short_id) = local_short_id {
state
.data
.find_video_task(VideoTaskLookupKey::ShortId(short_id))
.await
.ok()??
} else if let Some(task_id) = task_id {
if let Some(task) = state
.data
.find_video_task(VideoTaskLookupKey::Id(task_id))
.await
.ok()?
{
task
} else {
let user_id = user_id?;
state
.data
.find_video_task(VideoTaskLookupKey::UserExternal {
user_id,
external_task_id: task_id,
})
.await
.ok()??
}
} else {
return None;
};
let mut object = context.as_object()?.clone();
insert_missing_string_value(&mut object, "request_id", Some(task.request_id.as_str()));
insert_missing_optional_string_value(&mut object, "provider_id", task.provider_id.as_deref());
insert_missing_optional_string_value(&mut object, "endpoint_id", task.endpoint_id.as_deref());
insert_missing_optional_string_value(&mut object, "key_id", task.key_id.as_deref());
insert_missing_optional_string_value(&mut object, "user_id", task.user_id.as_deref());
insert_missing_optional_string_value(&mut object, "api_key_id", task.api_key_id.as_deref());
insert_missing_optional_string_value(
&mut object,
"client_api_format",
task.client_api_format.as_deref(),
);
insert_missing_optional_string_value(
&mut object,
"provider_api_format",
task.provider_api_format.as_deref(),
);
Some(Value::Object(object))
}
fn insert_missing_string_value(object: &mut Map<String, Value>, key: &str, value: Option<&str>) {
if object.contains_key(key) {
return;
}
let Some(value) = value.map(str::trim).filter(|value| !value.is_empty()) else {
return;
};
object.insert(key.to_string(), Value::String(value.to_string()));
}
fn insert_missing_optional_string_value(
object: &mut Map<String, Value>,
key: &str,
value: Option<&str>,
) {
insert_missing_string_value(object, key, value);
}
fn has_non_empty_str(value: &Value, key: &str) -> bool {
value
.get(key)
.and_then(Value::as_str)
.map(str::trim)
.is_some_and(|value| !value.is_empty())
}
fn has_u64(value: &Value, key: &str) -> bool {
value.get(key).and_then(Value::as_u64).is_some()
}

View File

@@ -0,0 +1,211 @@
use std::sync::Arc;
use aether_contracts::ExecutionPlan;
use aether_data::DataLayerError;
use tracing::warn;
use super::config::UsageRuntimeConfig;
use super::reporting::{GatewayStreamReportRequest, GatewaySyncReportRequest};
use super::worker::UsageQueueWorker;
use super::write::{
build_pending_usage_record, build_stream_finalized_execution_outcome,
build_streaming_usage_record, build_sync_finalized_execution_outcome,
build_terminal_usage_event_from_outcome,
};
use crate::gateway::billing_runtime::enrich_usage_event_with_billing;
use crate::gateway::gateway_data::GatewayDataState;
use crate::gateway::wallet_runtime::settle_usage_if_needed;
use crate::gateway::FinalizedExecutionState;
#[derive(Debug, Clone)]
pub(crate) struct UsageRuntime {
config: UsageRuntimeConfig,
}
impl Default for UsageRuntime {
fn default() -> Self {
Self::disabled()
}
}
impl UsageRuntime {
pub(crate) fn disabled() -> Self {
Self {
config: UsageRuntimeConfig::disabled(),
}
}
pub(crate) fn new(config: UsageRuntimeConfig) -> Result<Self, DataLayerError> {
config.validate()?;
Ok(Self { config })
}
pub(crate) fn is_enabled(&self) -> bool {
self.config.enabled
}
pub(crate) fn can_spawn_worker(&self, data: &GatewayDataState) -> bool {
self.is_enabled() && data.has_usage_writer() && data.has_usage_worker_runner()
}
pub(crate) fn spawn_worker(
&self,
data: Arc<GatewayDataState>,
) -> Option<tokio::task::JoinHandle<()>> {
if !self.can_spawn_worker(&data) {
return None;
}
let runner = data.usage_worker_runner()?;
let worker = UsageQueueWorker::new(runner, data, self.config.clone()).ok()?;
Some(worker.spawn())
}
pub(crate) async fn record_pending(
&self,
data: &GatewayDataState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
) {
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).await {
warn!(error = %err, request_id = %plan.request_id, "usage runtime failed to record sync pending usage");
}
}
Err(err) => {
warn!(error = %err, request_id = %plan.request_id, "usage runtime failed to build sync pending usage")
}
}
}
pub(crate) async fn record_stream_started(
&self,
data: &GatewayDataState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
status_code: u16,
headers: &std::collections::BTreeMap<String, String>,
telemetry: Option<&aether_contracts::ExecutionTelemetry>,
) {
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).await {
warn!(error = %err, request_id = %plan.request_id, "usage runtime failed to record stream usage");
}
}
Err(err) => {
warn!(error = %err, request_id = %plan.request_id, "usage runtime failed to build stream usage")
}
}
}
pub(crate) async fn record_sync_terminal(
&self,
data: &GatewayDataState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
payload: &GatewaySyncReportRequest,
) {
if !self.is_enabled() {
return;
}
match build_terminal_usage_event_from_outcome(build_sync_finalized_execution_outcome(
plan,
report_context,
payload,
)) {
Ok(mut event) => {
if let Err(err) = enrich_usage_event_with_billing(data, &mut event).await {
warn!(error = %err, request_id = %plan.request_id, "usage runtime failed to enrich sync usage event with billing");
}
self.enqueue_or_write_terminal(data, event).await
}
Err(err) => {
warn!(error = %err, request_id = %plan.request_id, "usage runtime failed to build sync terminal usage event")
}
}
}
pub(crate) async fn record_stream_terminal(
&self,
data: &GatewayDataState,
plan: &ExecutionPlan,
report_context: Option<&serde_json::Value>,
payload: &GatewayStreamReportRequest,
cancelled: bool,
) {
if !self.is_enabled() {
return;
}
let mut outcome = build_stream_finalized_execution_outcome(plan, report_context, payload);
if cancelled {
outcome.terminal_state = FinalizedExecutionState::Cancelled;
}
match build_terminal_usage_event_from_outcome(outcome) {
Ok(mut event) => {
if let Err(err) = enrich_usage_event_with_billing(data, &mut event).await {
warn!(error = %err, request_id = %plan.request_id, "usage runtime failed to enrich stream usage event with billing");
}
self.enqueue_or_write_terminal(data, event).await
}
Err(err) => {
warn!(error = %err, request_id = %plan.request_id, "usage runtime failed to build stream terminal usage event")
}
}
}
async fn enqueue_or_write_terminal(&self, data: &GatewayDataState, event: super::UsageEvent) {
if let Some(runner) = data.usage_worker_runner() {
match super::queue::UsageQueue::new(runner, self.config.clone()) {
Ok(queue) => match queue.enqueue(&event).await {
Ok(_) => return,
Err(err) => {
warn!(error = %err, request_id = %event.request_id, "usage runtime failed to enqueue terminal usage event; falling back to direct write")
}
},
Err(err) => {
warn!(error = %err, request_id = %event.request_id, "usage runtime failed to build queue; falling back to direct write")
}
}
}
match super::write::build_upsert_usage_record_from_event(&event) {
Ok(record) => match data.upsert_usage(record).await {
Ok(Some(stored)) => {
if let Err(err) = settle_usage_if_needed(data, &stored).await {
warn!(error = %err, request_id = %event.request_id, "usage runtime failed to settle terminal usage directly");
}
}
Ok(None) => {}
Err(err) => {
warn!(error = %err, request_id = %event.request_id, "usage runtime failed to upsert terminal usage directly");
}
},
Err(err) => {
warn!(error = %err, request_id = %event.request_id, "usage runtime failed to build terminal usage upsert")
}
}
}
}
fn now_unix_secs() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs()
}

View File

@@ -0,0 +1,204 @@
use std::sync::Arc;
use std::time::Duration;
use aether_data::redis::{RedisConsumerName, RedisStreamEntry, RedisStreamRunner};
use tracing::warn;
use super::config::UsageRuntimeConfig;
use super::event::UsageEvent;
use super::queue::UsageQueue;
use super::write::build_upsert_usage_record_from_event;
use crate::gateway::gateway_data::GatewayDataState;
use crate::gateway::wallet_runtime::settle_usage_if_needed;
pub(crate) struct UsageQueueWorker {
queue: UsageQueue,
data: Arc<GatewayDataState>,
consumer: RedisConsumerName,
config: UsageRuntimeConfig,
}
impl UsageQueueWorker {
pub(crate) fn new(
runner: RedisStreamRunner,
data: Arc<GatewayDataState>,
config: UsageRuntimeConfig,
) -> Result<Self, aether_data::DataLayerError> {
let queue = UsageQueue::new(runner, config.clone())?;
let consumer = RedisConsumerName(consumer_name());
Ok(Self {
queue,
data,
consumer,
config,
})
}
pub(crate) fn spawn(self) -> tokio::task::JoinHandle<()> {
tokio::spawn(async move { self.run_forever().await })
}
async fn run_forever(self) {
if let Err(err) = self.queue.ensure_consumer_group().await {
warn!(error = %err, "usage worker failed to ensure consumer group");
return;
}
let mut reclaim_interval =
tokio::time::interval(Duration::from_millis(self.config.reclaim_interval_ms));
reclaim_interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
reclaim_interval.tick().await;
loop {
tokio::select! {
_ = reclaim_interval.tick() => {
match self.queue.claim_stale(&self.consumer, "0-0").await {
Ok(entries) => {
if let Err(err) = self.process_entries(entries).await {
warn!(error = %err, "usage worker failed while reclaiming stale entries");
}
}
Err(err) => warn!(error = %err, "usage worker failed to reclaim stale entries"),
}
}
result = self.queue.read_group(&self.consumer) => {
match result {
Ok(entries) => {
if let Err(err) = self.process_entries(entries).await {
warn!(error = %err, "usage worker failed to process queue entries");
tokio::time::sleep(Duration::from_millis(250)).await;
}
}
Err(err) => {
warn!(error = %err, "usage worker failed to read queue");
tokio::time::sleep(Duration::from_millis(500)).await;
}
}
}
}
}
}
async fn process_entries(
&self,
entries: Vec<RedisStreamEntry>,
) -> Result<(), aether_data::DataLayerError> {
if entries.is_empty() {
return Ok(());
}
let mut ack_ids = Vec::new();
for entry in entries {
match self.process_entry(&entry).await {
Ok(should_ack) => {
if should_ack {
ack_ids.push(entry.id.clone());
}
}
Err(err) => {
if !ack_ids.is_empty() {
let _ = self.queue.ack_and_delete(&ack_ids).await;
}
return Err(err);
}
}
}
if !ack_ids.is_empty() {
self.queue.ack_and_delete(&ack_ids).await?;
}
Ok(())
}
async fn process_entry(
&self,
entry: &RedisStreamEntry,
) -> Result<bool, aether_data::DataLayerError> {
let event = match UsageEvent::from_stream_fields(&entry.fields) {
Ok(event) => event,
Err(err) => {
self.queue.push_dead_letter(entry, &err.to_string()).await?;
return Ok(true);
}
};
write_event_record(self.data.as_ref(), &event).await?;
Ok(true)
}
}
async fn write_event_record(
data: &GatewayDataState,
event: &UsageEvent,
) -> Result<(), aether_data::DataLayerError> {
let record = build_upsert_usage_record_from_event(event)?;
if let Some(stored) = data.upsert_usage(record).await? {
settle_usage_if_needed(data, &stored).await?;
}
Ok(())
}
fn consumer_name() -> String {
let host = std::env::var("HOSTNAME")
.ok()
.map(|value| value.trim().to_string())
.filter(|value| !value.is_empty())
.unwrap_or_else(|| "aether-gateway".to_string());
format!("{host}:{}", std::process::id())
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use aether_data::repository::usage::{InMemoryUsageReadRepository, UsageReadRepository};
use super::write_event_record;
use crate::gateway::gateway_data::GatewayDataState;
use crate::gateway::usage::{UsageEvent, UsageEventData, UsageEventType};
#[tokio::test]
async fn worker_writes_usage_record_from_terminal_event() {
let repository = Arc::new(InMemoryUsageReadRepository::default());
let data = GatewayDataState::with_usage_repository_for_tests(repository.clone());
let event = UsageEvent::new(
UsageEventType::Completed,
"req-worker-123".to_string(),
UsageEventData {
user_id: Some("user-worker-123".to_string()),
api_key_id: Some("api-key-worker-123".to_string()),
provider_name: "openai".to_string(),
provider_id: Some("provider-worker-123".to_string()),
provider_endpoint_id: Some("endpoint-worker-123".to_string()),
provider_api_key_id: Some("provider-key-worker-123".to_string()),
model: "gpt-5".to_string(),
api_format: Some("openai:chat".to_string()),
endpoint_api_format: Some("openai:chat".to_string()),
is_stream: Some(false),
status_code: Some(200),
input_tokens: Some(4),
output_tokens: Some(6),
total_tokens: Some(10),
response_time_ms: Some(52),
..UsageEventData::default()
},
);
write_event_record(&data, &event)
.await
.expect("worker should write usage record");
let stored = repository
.find_by_request_id("req-worker-123")
.await
.expect("usage lookup should succeed")
.expect("usage record should exist");
assert_eq!(stored.status, "completed");
assert_eq!(stored.billing_status, "pending");
assert_eq!(stored.total_tokens, 10);
assert_eq!(stored.response_time_ms, Some(52));
assert_eq!(stored.user_id.as_deref(), Some("user-worker-123"));
}
}

File diff suppressed because it is too large Load Diff