mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
refactor: 拆分 gateway 单体为独立 crate，新增 systemd 部署方案
将 gateway 内部的 model-fetch、provider-transport、scheduler-core、 usage-runtime、video-tasks-core 模块提取为独立 crate；重构 gateway 内部模块结构（state/router/cache/data/query 等）；移除大量遗留模块 文件；新增 systemd 二进制部署骨架及相关文档；更新前端 usage 相关 API 和组件。
This commit is contained in:
fawney19
@@ -1,57 +0,0 @@
|
||||
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,
|
||||
}))
|
||||
}
|
||||
@@ -1,112 +1 @@
|
||||
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());
|
||||
}
|
||||
}
|
||||
pub use aether_usage_runtime::UsageRuntimeConfig;
|
||||
|
||||
@@ -1,218 +1,3 @@
|
||||
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));
|
||||
}
|
||||
}
|
||||
pub use aether_usage_runtime::{
|
||||
now_ms, UsageEvent, UsageEventData, UsageEventType, USAGE_EVENT_VERSION,
|
||||
};
|
||||
|
||||
@@ -6,7 +6,9 @@ use axum::Json;
|
||||
use serde::Deserialize;
|
||||
use serde_json::json;
|
||||
|
||||
use crate::gateway::AppState;
|
||||
use crate::AppState;
|
||||
use aether_data::repository::audit::RequestAuditBundle;
|
||||
use aether_data::repository::usage::StoredRequestUsageAudit;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub(crate) struct GetRequestAuditBundleQuery {
|
||||
@@ -16,7 +18,7 @@ pub(crate) struct GetRequestAuditBundleQuery {
|
||||
pub(crate) async fn get_request_usage_audit(
|
||||
State(state): State<AppState>,
|
||||
Path(request_id): Path<String>,
|
||||
) -> Result<Json<super::RequestUsageAudit>, axum::response::Response> {
|
||||
) -> Result<Json<StoredRequestUsageAudit>, axum::response::Response> {
|
||||
let usage = state
|
||||
.read_request_usage_audit(&request_id)
|
||||
.await
|
||||
@@ -40,7 +42,7 @@ 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> {
|
||||
) -> Result<Json<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())
|
||||
|
||||
@@ -1,22 +1,15 @@
|
||||
mod bundle;
|
||||
mod config;
|
||||
mod event;
|
||||
mod http;
|
||||
pub(crate) mod event;
|
||||
pub(crate) mod http;
|
||||
mod queue;
|
||||
mod read;
|
||||
mod reporting;
|
||||
pub(crate) mod reporting;
|
||||
mod runtime;
|
||||
mod worker;
|
||||
mod write;
|
||||
pub(crate) 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,
|
||||
spawn_sync_report, submit_stream_report, submit_sync_report, GatewayStreamReportRequest,
|
||||
GatewaySyncReportRequest,
|
||||
};
|
||||
pub(crate) use runtime::UsageRuntime;
|
||||
pub(crate) use write::build_upsert_usage_record_from_event;
|
||||
|
||||
@@ -1,161 +1 @@
|
||||
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));
|
||||
}
|
||||
}
|
||||
pub use aether_usage_runtime::UsageQueue;
|
||||
|
||||
@@ -1,20 +0,0 @@
|
||||
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 }))
|
||||
}
|
||||
@@ -1,20 +1,14 @@
|
||||
use aether_data::repository::video_tasks::VideoTaskLookupKey;
|
||||
use serde_json::{Map, Value};
|
||||
use aether_usage_runtime::{
|
||||
build_locally_actionable_report_context_from_request_candidate,
|
||||
build_locally_actionable_report_context_from_video_task,
|
||||
};
|
||||
use serde_json::Value;
|
||||
|
||||
use crate::gateway::AppState;
|
||||
use crate::video_tasks::{resolve_video_task_report_lookup, VideoTaskReportLookup};
|
||||
use crate::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) use aether_usage_runtime::report_context_is_locally_actionable;
|
||||
|
||||
pub(crate) async fn resolve_locally_actionable_report_context(
|
||||
state: &AppState,
|
||||
@@ -61,143 +55,38 @@ async fn resolve_locally_actionable_report_context_from_request_candidates(
|
||||
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)
|
||||
build_locally_actionable_report_context_from_request_candidate(context, &existing_candidates[0])
|
||||
}
|
||||
|
||||
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()??
|
||||
let task = match resolve_video_task_report_lookup(context)? {
|
||||
VideoTaskReportLookup::Lookup(lookup) => {
|
||||
state.data.find_video_task(lookup).await.ok()??
|
||||
}
|
||||
VideoTaskReportLookup::TaskIdOrExternal { task_id, user_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()
|
||||
build_locally_actionable_report_context_from_video_task(context, &task)
|
||||
}
|
||||
|
||||
@@ -1,58 +1,63 @@
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use aether_contracts::ExecutionError;
|
||||
use aether_contracts::ExecutionTelemetry;
|
||||
use base64::Engine as _;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::Value;
|
||||
use tracing::{debug, warn};
|
||||
use uuid::Uuid;
|
||||
|
||||
use crate::gateway::scheduler::{
|
||||
use crate::scheduler::{
|
||||
current_unix_secs, execution_error_details, record_report_request_candidate_status,
|
||||
};
|
||||
use crate::gateway::{AppState, GatewayError};
|
||||
use crate::{AppState, GatewayError};
|
||||
|
||||
mod context;
|
||||
use context::{report_context_is_locally_actionable, resolve_locally_actionable_report_context};
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub(crate) struct GatewaySyncReportRequest {
|
||||
pub(crate) trace_id: String,
|
||||
pub(crate) report_kind: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) report_context: Option<serde_json::Value>,
|
||||
pub(crate) status_code: u16,
|
||||
pub(crate) headers: BTreeMap<String, String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) body_json: Option<serde_json::Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) client_body_json: Option<serde_json::Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) body_base64: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) telemetry: Option<ExecutionTelemetry>,
|
||||
use aether_usage_runtime::{
|
||||
extract_gemini_file_mapping_entries, gemini_file_mapping_cache_key,
|
||||
is_local_ai_stream_report_kind, is_local_ai_sync_report_kind, normalize_gemini_file_name,
|
||||
report_request_id, should_handle_local_stream_report, should_handle_local_sync_report,
|
||||
sync_report_represents_failure, GEMINI_FILE_MAPPING_TTL_SECONDS,
|
||||
};
|
||||
pub(crate) use aether_usage_runtime::{GatewayStreamReportRequest, GatewaySyncReportRequest};
|
||||
|
||||
fn log_local_report_handled(
|
||||
trace_id: &str,
|
||||
report_kind: &str,
|
||||
report_scope: &'static str,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
) {
|
||||
debug!(
|
||||
event_name = "execution_report_handled_locally",
|
||||
log_type = "debug",
|
||||
debug_context = "redacted",
|
||||
trace_id = %trace_id,
|
||||
report_scope,
|
||||
report_kind = %report_kind,
|
||||
report_request_id = report_request_id(report_context),
|
||||
has_report_context = report_context.is_some(),
|
||||
"gateway handled execution report locally"
|
||||
);
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
pub(crate) struct GatewayStreamReportRequest {
|
||||
pub(crate) trace_id: String,
|
||||
pub(crate) report_kind: String,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) report_context: Option<serde_json::Value>,
|
||||
pub(crate) status_code: u16,
|
||||
pub(crate) headers: BTreeMap<String, String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) provider_body_base64: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) client_body_base64: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub(crate) telemetry: Option<ExecutionTelemetry>,
|
||||
fn log_dropped_report(
|
||||
trace_id: &str,
|
||||
report_kind: &str,
|
||||
report_scope: &'static str,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
) {
|
||||
warn!(
|
||||
event_name = "execution_report_dropped",
|
||||
log_type = "ops",
|
||||
status = "dropped",
|
||||
trace_id = %trace_id,
|
||||
report_scope,
|
||||
report_kind = %report_kind,
|
||||
report_request_id = report_request_id(report_context),
|
||||
has_report_context = report_context.is_some(),
|
||||
"gateway dropped execution report because local handling context was not actionable"
|
||||
);
|
||||
}
|
||||
|
||||
const GEMINI_FILE_MAPPING_TTL_SECONDS: u64 = 60 * 60 * 48;
|
||||
const GEMINI_FILE_MAPPING_CACHE_PREFIX: &str = "gemini_files:key";
|
||||
|
||||
pub(crate) async fn submit_sync_report(
|
||||
state: &AppState,
|
||||
trace_id: &str,
|
||||
@@ -63,32 +68,40 @@ pub(crate) async fn submit_sync_report(
|
||||
{
|
||||
let mut local_payload = payload.clone();
|
||||
local_payload.report_context = Some(report_context);
|
||||
if should_handle_local_sync_report(state, &local_payload) {
|
||||
if should_handle_local_sync_report(
|
||||
local_payload.report_context.as_ref(),
|
||||
local_payload.report_kind.as_str(),
|
||||
) {
|
||||
handle_local_sync_report(state, &local_payload).await;
|
||||
debug!(
|
||||
trace_id = %trace_id,
|
||||
report_kind = %local_payload.report_kind,
|
||||
"gateway handled sync execution report locally"
|
||||
log_local_report_handled(
|
||||
trace_id,
|
||||
&local_payload.report_kind,
|
||||
"sync",
|
||||
local_payload.report_context.as_ref(),
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
if should_handle_local_sync_report(state, &payload) {
|
||||
if should_handle_local_sync_report(
|
||||
payload.report_context.as_ref(),
|
||||
payload.report_kind.as_str(),
|
||||
) {
|
||||
handle_local_sync_report(state, &payload).await;
|
||||
debug!(
|
||||
trace_id = %trace_id,
|
||||
report_kind = %payload.report_kind,
|
||||
"gateway handled sync execution report locally"
|
||||
log_local_report_handled(
|
||||
trace_id,
|
||||
&payload.report_kind,
|
||||
"sync",
|
||||
payload.report_context.as_ref(),
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
report_kind = %payload.report_kind,
|
||||
report_context = ?payload.report_context,
|
||||
"gateway dropped sync execution report because local handling context was not actionable"
|
||||
log_dropped_report(
|
||||
trace_id,
|
||||
&payload.report_kind,
|
||||
"sync",
|
||||
payload.report_context.as_ref(),
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
@@ -100,7 +113,14 @@ pub(crate) fn spawn_sync_report(
|
||||
) {
|
||||
tokio::spawn(async move {
|
||||
if let Err(err) = submit_sync_report(&state, &trace_id, payload).await {
|
||||
warn!(trace_id = %trace_id, error = ?err, "gateway failed to submit sync execution report");
|
||||
warn!(
|
||||
event_name = "execution_report_submit_failed",
|
||||
log_type = "ops",
|
||||
trace_id = %trace_id,
|
||||
report_scope = "sync",
|
||||
error = ?err,
|
||||
"gateway failed to submit sync execution report"
|
||||
);
|
||||
}
|
||||
});
|
||||
}
|
||||
@@ -123,57 +143,50 @@ pub(crate) async fn submit_stream_report(
|
||||
client_body_base64: payload.client_body_base64.clone(),
|
||||
telemetry: payload.telemetry.clone(),
|
||||
};
|
||||
if should_handle_local_stream_report(state, &local_payload) {
|
||||
if should_handle_local_stream_report(
|
||||
local_payload.report_context.as_ref(),
|
||||
local_payload.report_kind.as_str(),
|
||||
) {
|
||||
handle_local_stream_report(state, &local_payload).await;
|
||||
debug!(
|
||||
trace_id = %trace_id,
|
||||
report_kind = %local_payload.report_kind,
|
||||
"gateway handled stream execution report locally"
|
||||
log_local_report_handled(
|
||||
trace_id,
|
||||
&local_payload.report_kind,
|
||||
"stream",
|
||||
local_payload.report_context.as_ref(),
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
|
||||
if should_handle_local_stream_report(state, &payload) {
|
||||
if should_handle_local_stream_report(
|
||||
payload.report_context.as_ref(),
|
||||
payload.report_kind.as_str(),
|
||||
) {
|
||||
handle_local_stream_report(state, &payload).await;
|
||||
debug!(
|
||||
trace_id = %trace_id,
|
||||
report_kind = %payload.report_kind,
|
||||
"gateway handled stream execution report locally"
|
||||
log_local_report_handled(
|
||||
trace_id,
|
||||
&payload.report_kind,
|
||||
"stream",
|
||||
payload.report_context.as_ref(),
|
||||
);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
warn!(
|
||||
trace_id = %trace_id,
|
||||
report_kind = %payload.report_kind,
|
||||
report_context = ?payload.report_context,
|
||||
"gateway dropped stream execution report because local handling context was not actionable"
|
||||
log_dropped_report(
|
||||
trace_id,
|
||||
&payload.report_kind,
|
||||
"stream",
|
||||
payload.report_context.as_ref(),
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn should_handle_local_sync_report(state: &AppState, payload: &GatewaySyncReportRequest) -> bool {
|
||||
let _ = state;
|
||||
report_context_is_locally_actionable(payload.report_context.as_ref())
|
||||
&& is_local_ai_sync_report_kind(payload.report_kind.as_str())
|
||||
}
|
||||
|
||||
fn should_handle_local_stream_report(
|
||||
state: &AppState,
|
||||
payload: &GatewayStreamReportRequest,
|
||||
) -> bool {
|
||||
let _ = state;
|
||||
report_context_is_locally_actionable(payload.report_context.as_ref())
|
||||
&& is_local_ai_stream_report_kind(payload.report_kind.as_str())
|
||||
}
|
||||
|
||||
async fn handle_local_sync_report(state: &AppState, payload: &GatewaySyncReportRequest) {
|
||||
apply_local_gemini_file_mapping_side_effect(state, payload).await;
|
||||
let terminal_unix_secs = current_unix_secs();
|
||||
let (error_type, error_message) =
|
||||
execution_error_details(None::<&ExecutionError>, payload.body_json.as_ref());
|
||||
let status = if sync_report_represents_failure(payload, error_type.as_ref()) {
|
||||
let status = if sync_report_represents_failure(payload, error_type.as_deref()) {
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Failed
|
||||
} else {
|
||||
aether_data::repository::candidates::RequestCandidateStatus::Success
|
||||
@@ -216,24 +229,6 @@ async fn handle_local_stream_report(state: &AppState, payload: &GatewayStreamRep
|
||||
.await;
|
||||
}
|
||||
|
||||
fn sync_report_represents_failure(
|
||||
payload: &GatewaySyncReportRequest,
|
||||
error_type: Option<&String>,
|
||||
) -> bool {
|
||||
if payload.report_kind == "openai_video_delete_sync_success" && payload.status_code == 404 {
|
||||
return false;
|
||||
}
|
||||
|
||||
payload.status_code >= 400
|
||||
|| payload.report_kind.contains("error")
|
||||
|| error_type.is_some()
|
||||
|| payload
|
||||
.body_json
|
||||
.as_ref()
|
||||
.and_then(|body| body.get("error"))
|
||||
.is_some()
|
||||
}
|
||||
|
||||
async fn apply_local_gemini_file_mapping_side_effect(
|
||||
state: &AppState,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
@@ -248,14 +243,14 @@ async fn apply_local_gemini_file_mapping_side_effect(
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|context| context.get("file_key_id"))
|
||||
.and_then(Value::as_str)
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
let user_id = payload
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|context| context.get("user_id"))
|
||||
.and_then(Value::as_str)
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty());
|
||||
let Some(key_id) = key_id else {
|
||||
@@ -274,7 +269,10 @@ async fn apply_local_gemini_file_mapping_side_effect(
|
||||
.await
|
||||
{
|
||||
warn!(
|
||||
event_name = "gemini_file_mapping_store_failed",
|
||||
log_type = "ops",
|
||||
report_kind = %payload.report_kind,
|
||||
report_request_id = report_request_id(payload.report_context.as_ref()),
|
||||
file_name = %entry.file_name,
|
||||
error = ?err,
|
||||
"gateway failed to persist gemini file mapping locally"
|
||||
@@ -287,15 +285,18 @@ async fn apply_local_gemini_file_mapping_side_effect(
|
||||
.report_context
|
||||
.as_ref()
|
||||
.and_then(|context| context.get("file_name"))
|
||||
.and_then(Value::as_str)
|
||||
.and_then(normalize_file_name);
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.and_then(normalize_gemini_file_name);
|
||||
let Some(file_name) = file_name else {
|
||||
return;
|
||||
};
|
||||
|
||||
if let Err(err) = delete_local_gemini_file_mapping(state, file_name.as_str()).await {
|
||||
warn!(
|
||||
event_name = "gemini_file_mapping_delete_failed",
|
||||
log_type = "ops",
|
||||
report_kind = %payload.report_kind,
|
||||
report_request_id = report_request_id(payload.report_context.as_ref()),
|
||||
file_name = %file_name,
|
||||
error = ?err,
|
||||
"gateway failed to delete gemini file mapping locally"
|
||||
@@ -306,118 +307,6 @@ async fn apply_local_gemini_file_mapping_side_effect(
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct LocalGeminiFileMappingEntry {
|
||||
file_name: String,
|
||||
display_name: Option<String>,
|
||||
mime_type: Option<String>,
|
||||
}
|
||||
|
||||
fn extract_gemini_file_mapping_entries(
|
||||
payload: &GatewaySyncReportRequest,
|
||||
) -> Vec<LocalGeminiFileMappingEntry> {
|
||||
let Some(body) = extract_sync_report_body_json(payload) else {
|
||||
return Vec::new();
|
||||
};
|
||||
let Some(object) = body.as_object() else {
|
||||
return Vec::new();
|
||||
};
|
||||
|
||||
let mut entries = Vec::new();
|
||||
maybe_push_local_gemini_file_mapping_entry(&mut entries, object);
|
||||
|
||||
if let Some(file_object) = object.get("file").and_then(Value::as_object) {
|
||||
maybe_push_local_gemini_file_mapping_entry(&mut entries, file_object);
|
||||
}
|
||||
|
||||
if let Some(files) = object.get("files").and_then(Value::as_array) {
|
||||
for item in files {
|
||||
if let Some(file_object) = item.as_object() {
|
||||
maybe_push_local_gemini_file_mapping_entry(&mut entries, file_object);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
entries
|
||||
}
|
||||
|
||||
fn maybe_push_local_gemini_file_mapping_entry(
|
||||
entries: &mut Vec<LocalGeminiFileMappingEntry>,
|
||||
object: &serde_json::Map<String, Value>,
|
||||
) {
|
||||
let file_name = object
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.and_then(normalize_file_name);
|
||||
let Some(file_name) = file_name else {
|
||||
return;
|
||||
};
|
||||
|
||||
if entries.iter().any(|entry| entry.file_name == file_name) {
|
||||
return;
|
||||
}
|
||||
|
||||
entries.push(LocalGeminiFileMappingEntry {
|
||||
file_name,
|
||||
display_name: object
|
||||
.get("displayName")
|
||||
.or_else(|| object.get("display_name"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
mime_type: object
|
||||
.get("mimeType")
|
||||
.or_else(|| object.get("mime_type"))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned),
|
||||
});
|
||||
}
|
||||
|
||||
fn extract_sync_report_body_json(payload: &GatewaySyncReportRequest) -> Option<Value> {
|
||||
if let Some(body_json) = payload.body_json.as_ref() {
|
||||
return Some(body_json.clone());
|
||||
}
|
||||
if let Some(client_body_json) = payload.client_body_json.as_ref() {
|
||||
return Some(client_body_json.clone());
|
||||
}
|
||||
if !content_type_starts_with(&payload.headers, "application/json") {
|
||||
return None;
|
||||
}
|
||||
|
||||
let body_base64 = payload.body_base64.as_deref()?;
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(body_base64)
|
||||
.ok()?;
|
||||
serde_json::from_slice(&bytes).ok()
|
||||
}
|
||||
|
||||
fn content_type_starts_with(headers: &BTreeMap<String, String>, expected_prefix: &str) -> bool {
|
||||
headers
|
||||
.iter()
|
||||
.find(|(key, _)| key.eq_ignore_ascii_case("content-type"))
|
||||
.map(|(_, value)| value.trim().to_ascii_lowercase())
|
||||
.is_some_and(|value| value.starts_with(expected_prefix))
|
||||
}
|
||||
|
||||
fn normalize_file_name(file_name: &str) -> Option<String> {
|
||||
let file_name = file_name.trim();
|
||||
if file_name.is_empty() {
|
||||
return None;
|
||||
}
|
||||
if file_name.starts_with("files/") {
|
||||
Some(file_name.to_string())
|
||||
} else {
|
||||
Some(format!("files/{file_name}"))
|
||||
}
|
||||
}
|
||||
|
||||
fn gemini_file_mapping_cache_key(file_name: &str) -> String {
|
||||
format!("{GEMINI_FILE_MAPPING_CACHE_PREFIX}:{file_name}")
|
||||
}
|
||||
|
||||
pub(crate) async fn store_local_gemini_file_mapping(
|
||||
state: &AppState,
|
||||
file_name: &str,
|
||||
@@ -426,7 +315,7 @@ pub(crate) async fn store_local_gemini_file_mapping(
|
||||
display_name: Option<&str>,
|
||||
mime_type: Option<&str>,
|
||||
) -> Result<(), GatewayError> {
|
||||
let Some(file_name) = normalize_file_name(file_name) else {
|
||||
let Some(file_name) = normalize_gemini_file_name(file_name) else {
|
||||
return Ok(());
|
||||
};
|
||||
let expires_at_unix_secs = current_unix_secs().saturating_add(GEMINI_FILE_MAPPING_TTL_SECONDS);
|
||||
@@ -459,7 +348,7 @@ async fn delete_local_gemini_file_mapping(
|
||||
state: &AppState,
|
||||
file_name: &str,
|
||||
) -> Result<(), GatewayError> {
|
||||
let Some(file_name) = normalize_file_name(file_name) else {
|
||||
let Some(file_name) = normalize_gemini_file_name(file_name) else {
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
@@ -472,51 +361,6 @@ async fn delete_local_gemini_file_mapping(
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn is_local_ai_sync_report_kind(report_kind: &str) -> bool {
|
||||
matches!(
|
||||
report_kind,
|
||||
"openai_chat_sync_success"
|
||||
| "claude_chat_sync_success"
|
||||
| "gemini_chat_sync_success"
|
||||
| "openai_chat_sync_error"
|
||||
| "claude_chat_sync_error"
|
||||
| "gemini_chat_sync_error"
|
||||
| "openai_cli_sync_success"
|
||||
| "claude_cli_sync_success"
|
||||
| "gemini_cli_sync_success"
|
||||
| "openai_cli_sync_error"
|
||||
| "openai_compact_sync_error"
|
||||
| "claude_cli_sync_error"
|
||||
| "gemini_cli_sync_error"
|
||||
| "openai_video_create_sync_success"
|
||||
| "openai_video_remix_sync_success"
|
||||
| "gemini_video_create_sync_success"
|
||||
| "openai_video_delete_sync_success"
|
||||
| "openai_video_cancel_sync_success"
|
||||
| "gemini_video_cancel_sync_success"
|
||||
| "openai_video_create_sync_error"
|
||||
| "openai_video_remix_sync_error"
|
||||
| "gemini_video_create_sync_error"
|
||||
| "openai_video_delete_sync_error"
|
||||
| "openai_video_cancel_sync_error"
|
||||
| "gemini_video_cancel_sync_error"
|
||||
| "gemini_files_store_mapping"
|
||||
| "gemini_files_delete_mapping"
|
||||
)
|
||||
}
|
||||
|
||||
fn is_local_ai_stream_report_kind(report_kind: &str) -> bool {
|
||||
matches!(
|
||||
report_kind,
|
||||
"openai_chat_stream_success"
|
||||
| "claude_chat_stream_success"
|
||||
| "gemini_chat_stream_success"
|
||||
| "openai_cli_stream_success"
|
||||
| "claude_cli_stream_success"
|
||||
| "gemini_cli_stream_success"
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
@@ -539,8 +383,8 @@ mod tests {
|
||||
resolve_locally_actionable_report_context, submit_stream_report, submit_sync_report,
|
||||
GatewayStreamReportRequest, GatewaySyncReportRequest,
|
||||
};
|
||||
use crate::gateway::gateway_data::GatewayDataState;
|
||||
use crate::gateway::AppState;
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::AppState;
|
||||
|
||||
fn sample_request_candidate(id: &str, request_id: &str) -> StoredRequestCandidate {
|
||||
StoredRequestCandidate::new(
|
||||
@@ -1,211 +1 @@
|
||||
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()
|
||||
}
|
||||
pub(crate) use aether_usage_runtime::UsageRuntime;
|
||||
|
||||
@@ -1,152 +1,4 @@
|
||||
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())
|
||||
}
|
||||
pub(crate) use aether_usage_runtime::{build_usage_queue_worker, write_event_record};
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
@@ -155,8 +7,8 @@ mod tests {
|
||||
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};
|
||||
use crate::data::GatewayDataState;
|
||||
use crate::usage::event::{UsageEvent, UsageEventData, UsageEventType};
|
||||
|
||||
#[tokio::test]
|
||||
async fn worker_writes_usage_record_from_terminal_event() {
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user