mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
feat(gateway): 重构 usage 数据层、迁移系统与系统导入
数据库迁移: - 引入 baseline v2 bootstrap，空库首次启动自动初始化 - 服务启动不再自动执行迁移，需显式 `--migrate` 运行 - 新增 pending migration 检测，schema 落后时拒绝启动 Usage 数据层: - usage body 存储外部化为独立 blob 表 - 新增 HTTP audit 表拆分存储请求/响应头与 body ref - 后台清理任务支持 legacy body ref 元数据迁移 - usage runtime 写入迁移到专用 tokio runtime（独立线程池, 8MB 栈） 系统导入/导出: - 支持用户、API Keys、钱包数据的完整导入 - 兼容 legacy 与 v1.3+ 两种导出格式 其他改进: - executor outcome 增加 runtime miss 诊断上下文 - 主 tokio runtime 栈大小调整为 8MB - 前端 provider 管理支持 base URL 配置 - dev.sh 支持 --migrate 参数
This commit is contained in:
fawney19
@@ -1029,7 +1029,8 @@ pub fn build_admin_stats_cost_savings_response(
|
||||
let mut estimated_full_cost: f64 = usage
|
||||
.iter()
|
||||
.map(|item| {
|
||||
item.output_price_per_1m.unwrap_or(0.0) * item.cache_read_input_tokens as f64
|
||||
item.settlement_output_price_per_1m().unwrap_or(0.0)
|
||||
* item.cache_read_input_tokens as f64
|
||||
/ 1_000_000.0
|
||||
})
|
||||
.sum();
|
||||
@@ -1295,6 +1296,7 @@ pub fn build_model_leaderboard_items(
|
||||
pub fn build_user_leaderboard_items(
|
||||
items: &[StoredRequestUsageAudit],
|
||||
users: &std::collections::BTreeMap<String, AdminStatsUserMetadata>,
|
||||
auth_user_reader_available: bool,
|
||||
include_inactive: bool,
|
||||
exclude_admin: bool,
|
||||
) -> Vec<AdminStatsLeaderboardItem> {
|
||||
@@ -1325,12 +1327,16 @@ pub fn build_user_leaderboard_items(
|
||||
if exclude_admin {
|
||||
continue;
|
||||
}
|
||||
item.username
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| user_id.to_string())
|
||||
if auth_user_reader_available {
|
||||
user_id.to_string()
|
||||
} else {
|
||||
item.username
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
.unwrap_or_else(|| user_id.to_string())
|
||||
}
|
||||
};
|
||||
|
||||
let entry =
|
||||
@@ -1359,6 +1365,7 @@ pub fn build_user_leaderboard_items(
|
||||
pub fn build_api_key_leaderboard_items(
|
||||
items: &[StoredRequestUsageAudit],
|
||||
snapshots: Option<&[StoredAuthApiKeySnapshot]>,
|
||||
api_key_names: &std::collections::BTreeMap<String, String>,
|
||||
include_inactive: bool,
|
||||
exclude_admin: bool,
|
||||
) -> Vec<AdminStatsLeaderboardItem> {
|
||||
@@ -1391,17 +1398,18 @@ pub fn build_api_key_leaderboard_items(
|
||||
if exclude_admin && snapshot.user_role.eq_ignore_ascii_case("admin") {
|
||||
continue;
|
||||
}
|
||||
snapshot
|
||||
.api_key_name
|
||||
.clone()
|
||||
.or_else(|| item.api_key_name.clone())
|
||||
api_key_names
|
||||
.get(api_key_id)
|
||||
.cloned()
|
||||
.unwrap_or_else(|| api_key_id.to_string())
|
||||
} else {
|
||||
if snapshots_available {
|
||||
continue;
|
||||
}
|
||||
item.api_key_name
|
||||
.clone()
|
||||
api_key_names
|
||||
.get(api_key_id)
|
||||
.cloned()
|
||||
.or_else(|| item.api_key_name.clone())
|
||||
.unwrap_or_else(|| api_key_id.to_string())
|
||||
};
|
||||
|
||||
@@ -1428,6 +1436,166 @@ pub fn build_api_key_leaderboard_items(
|
||||
grouped.into_values().collect()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{
|
||||
build_api_key_leaderboard_items, build_user_leaderboard_items, AdminStatsUserMetadata,
|
||||
};
|
||||
use aether_data::repository::auth::StoredAuthApiKeySnapshot;
|
||||
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
||||
|
||||
fn sample_usage(api_key_name: Option<&str>) -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
Some("user-1".to_string()),
|
||||
Some("key-1".to_string()),
|
||||
Some("alice".to_string()),
|
||||
api_key_name.map(str::to_string),
|
||||
"OpenAI".to_string(),
|
||||
"gpt-5".to_string(),
|
||||
None,
|
||||
Some("provider-1".to_string()),
|
||||
Some("endpoint-1".to_string()),
|
||||
Some("provider-key-1".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
false,
|
||||
false,
|
||||
10,
|
||||
20,
|
||||
30,
|
||||
0.3,
|
||||
0.3,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(400),
|
||||
Some(120),
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
1_711_000_000,
|
||||
1_711_000_001,
|
||||
Some(1_711_000_002),
|
||||
)
|
||||
.expect("usage should build")
|
||||
}
|
||||
|
||||
fn sample_api_key_snapshot(api_key_name: Option<&str>) -> StoredAuthApiKeySnapshot {
|
||||
StoredAuthApiKeySnapshot {
|
||||
user_id: "user-1".to_string(),
|
||||
username: "alice".to_string(),
|
||||
email: Some("alice@example.com".to_string()),
|
||||
user_role: "user".to_string(),
|
||||
user_auth_source: "local".to_string(),
|
||||
user_is_active: true,
|
||||
user_is_deleted: false,
|
||||
user_rate_limit: None,
|
||||
user_allowed_providers: None,
|
||||
user_allowed_api_formats: None,
|
||||
user_allowed_models: None,
|
||||
api_key_id: "key-1".to_string(),
|
||||
api_key_name: api_key_name.map(str::to_string),
|
||||
api_key_is_active: true,
|
||||
api_key_is_locked: false,
|
||||
api_key_is_standalone: false,
|
||||
api_key_rate_limit: None,
|
||||
api_key_concurrent_limit: None,
|
||||
api_key_expires_at_unix_secs: None,
|
||||
api_key_allowed_providers: None,
|
||||
api_key_allowed_api_formats: None,
|
||||
api_key_allowed_models: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn api_key_leaderboard_prefers_resolved_names_over_legacy_usage_names_when_snapshots_exist() {
|
||||
let leaderboard = build_api_key_leaderboard_items(
|
||||
&[sample_usage(Some("legacy-default"))],
|
||||
Some(&[sample_api_key_snapshot(None)]),
|
||||
&BTreeMap::from([("key-1".to_string(), "fresh-default".to_string())]),
|
||||
false,
|
||||
false,
|
||||
);
|
||||
|
||||
assert_eq!(leaderboard.len(), 1);
|
||||
assert_eq!(leaderboard[0].id, "key-1");
|
||||
assert_eq!(leaderboard[0].name, "fresh-default");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn api_key_leaderboard_keeps_legacy_usage_name_fallback_without_snapshot_reader() {
|
||||
let leaderboard = build_api_key_leaderboard_items(
|
||||
&[sample_usage(Some("legacy-default"))],
|
||||
None,
|
||||
&BTreeMap::new(),
|
||||
false,
|
||||
false,
|
||||
);
|
||||
|
||||
assert_eq!(leaderboard.len(), 1);
|
||||
assert_eq!(leaderboard[0].name, "legacy-default");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_leaderboard_prefers_resolved_names_over_legacy_usage_names_when_reader_exists() {
|
||||
let leaderboard = build_user_leaderboard_items(
|
||||
&[sample_usage(Some("legacy-default"))],
|
||||
&BTreeMap::from([(
|
||||
"user-1".to_string(),
|
||||
AdminStatsUserMetadata {
|
||||
name: "fresh-alice".to_string(),
|
||||
role: "user".to_string(),
|
||||
is_active: true,
|
||||
is_deleted: false,
|
||||
},
|
||||
)]),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
);
|
||||
|
||||
assert_eq!(leaderboard.len(), 1);
|
||||
assert_eq!(leaderboard[0].id, "user-1");
|
||||
assert_eq!(leaderboard[0].name, "fresh-alice");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_leaderboard_does_not_fallback_to_legacy_usage_name_when_reader_exists() {
|
||||
let leaderboard = build_user_leaderboard_items(
|
||||
&[sample_usage(Some("legacy-default"))],
|
||||
&BTreeMap::new(),
|
||||
true,
|
||||
false,
|
||||
false,
|
||||
);
|
||||
|
||||
assert_eq!(leaderboard.len(), 1);
|
||||
assert_eq!(leaderboard[0].name, "user-1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn user_leaderboard_keeps_legacy_usage_name_fallback_without_reader() {
|
||||
let leaderboard = build_user_leaderboard_items(
|
||||
&[sample_usage(Some("legacy-default"))],
|
||||
&BTreeMap::new(),
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
);
|
||||
|
||||
assert_eq!(leaderboard.len(), 1);
|
||||
assert_eq!(leaderboard[0].name, "alice");
|
||||
}
|
||||
}
|
||||
|
||||
pub fn compare_leaderboard_items(
|
||||
metric: AdminStatsLeaderboardMetric,
|
||||
order: AdminStatsSortOrder,
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -144,7 +144,7 @@ pub fn build_gateway_control_plan_request(
|
||||
pub fn augment_sync_report_context(
|
||||
report_context: Option<serde_json::Value>,
|
||||
provider_request_headers: &BTreeMap<String, String>,
|
||||
provider_request_body: &serde_json::Value,
|
||||
_provider_request_body: &serde_json::Value,
|
||||
) -> serde_json::Result<Option<serde_json::Value>> {
|
||||
let mut report_context = match report_context {
|
||||
Some(serde_json::Value::Object(map)) => map,
|
||||
@@ -156,10 +156,6 @@ pub fn augment_sync_report_context(
|
||||
"provider_request_headers".to_string(),
|
||||
serde_json::to_value(provider_request_headers)?,
|
||||
);
|
||||
report_context.insert(
|
||||
"provider_request_body".to_string(),
|
||||
provider_request_body.clone(),
|
||||
);
|
||||
|
||||
Ok(Some(serde_json::Value::Object(report_context)))
|
||||
}
|
||||
@@ -233,7 +229,7 @@ mod tests {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn augment_sync_report_context_attaches_provider_request_shape() {
|
||||
fn augment_sync_report_context_attaches_provider_request_headers_only() {
|
||||
let report_context = augment_sync_report_context(
|
||||
Some(serde_json::json!({"trace_id": "abc"})),
|
||||
&BTreeMap::from([("content-type".to_string(), "application/json".to_string())]),
|
||||
@@ -246,16 +242,13 @@ mod tests {
|
||||
report_context.get("trace_id"),
|
||||
Some(&serde_json::json!("abc"))
|
||||
);
|
||||
assert_eq!(
|
||||
report_context.get("provider_request_body"),
|
||||
Some(&serde_json::json!({"model": "gpt-5"}))
|
||||
);
|
||||
assert_eq!(
|
||||
report_context
|
||||
.get("provider_request_headers")
|
||||
.and_then(|value| value.get("content-type")),
|
||||
Some(&serde_json::json!("application/json"))
|
||||
);
|
||||
assert!(report_context.get("provider_request_body").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1,7 +1,8 @@
|
||||
mod types;
|
||||
|
||||
pub use types::{
|
||||
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
|
||||
StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository,
|
||||
UsageRepository, UsageWriteRepository,
|
||||
parse_usage_body_ref, usage_body_ref, StoredProviderApiKeyUsageSummary,
|
||||
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
|
||||
UpsertUsageRecord, UsageAuditListQuery, UsageBodyField, UsageReadRepository, UsageRepository,
|
||||
UsageWriteRepository,
|
||||
};
|
||||
|
||||
@@ -48,18 +48,42 @@ pub struct StoredRequestUsageAudit {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub request_body: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub request_body_ref: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_request_headers: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_request_body: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_request_body_ref: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub response_headers: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub response_body: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub response_body_ref: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub client_response_headers: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub client_response_body: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub client_response_body_ref: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub candidate_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub candidate_index: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub key_name: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub planner_kind: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub route_family: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub route_kind: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub execution_path: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub local_execution_runtime_miss_reason: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub request_metadata: Option<Value>,
|
||||
pub created_at_unix_ms: u64,
|
||||
pub updated_at_unix_secs: u64,
|
||||
@@ -185,12 +209,24 @@ impl StoredRequestUsageAudit {
|
||||
billing_status,
|
||||
request_headers: None,
|
||||
request_body: None,
|
||||
request_body_ref: None,
|
||||
provider_request_headers: None,
|
||||
provider_request_body: None,
|
||||
provider_request_body_ref: None,
|
||||
response_headers: None,
|
||||
response_body: None,
|
||||
response_body_ref: None,
|
||||
client_response_headers: None,
|
||||
client_response_body: None,
|
||||
client_response_body_ref: None,
|
||||
candidate_id: None,
|
||||
candidate_index: None,
|
||||
key_name: None,
|
||||
planner_kind: None,
|
||||
route_family: None,
|
||||
route_kind: None,
|
||||
execution_path: None,
|
||||
local_execution_runtime_miss_reason: None,
|
||||
request_metadata: None,
|
||||
created_at_unix_ms: parse_timestamp(created_at_unix_ms, "usage.created_at_unix_ms")?,
|
||||
updated_at_unix_secs: parse_timestamp(
|
||||
@@ -212,6 +248,154 @@ impl StoredRequestUsageAudit {
|
||||
self.cache_read_input_tokens = cache_read_input_tokens;
|
||||
self
|
||||
}
|
||||
|
||||
fn request_metadata_object(&self) -> Option<&serde_json::Map<String, Value>> {
|
||||
self.request_metadata.as_ref().and_then(Value::as_object)
|
||||
}
|
||||
|
||||
fn request_metadata_number(&self, key: &str) -> Option<f64> {
|
||||
self.request_metadata_object()
|
||||
.and_then(|metadata| metadata.get(key))
|
||||
.and_then(Value::as_f64)
|
||||
.filter(|value| value.is_finite())
|
||||
}
|
||||
|
||||
fn request_metadata_u64(&self, key: &str) -> Option<u64> {
|
||||
self.request_metadata_object()
|
||||
.and_then(|metadata| metadata.get(key))
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_u64()
|
||||
.or_else(|| value.as_i64().and_then(|n| u64::try_from(n).ok()))
|
||||
})
|
||||
}
|
||||
|
||||
fn request_metadata_bool(&self, key: &str) -> Option<bool> {
|
||||
self.request_metadata_object()
|
||||
.and_then(|metadata| metadata.get(key))
|
||||
.and_then(Value::as_bool)
|
||||
}
|
||||
|
||||
fn request_metadata_string(&self, key: &str) -> Option<&str> {
|
||||
self.request_metadata_object()
|
||||
.and_then(|metadata| metadata.get(key))
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn billing_snapshot_resolved_number(&self, key: &str) -> Option<f64> {
|
||||
self.request_metadata_object()
|
||||
.and_then(|metadata| metadata.get("billing_snapshot"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|snapshot| snapshot.get("resolved_variables"))
|
||||
.and_then(Value::as_object)
|
||||
.and_then(|variables| variables.get(key))
|
||||
.and_then(Value::as_f64)
|
||||
.filter(|value| value.is_finite())
|
||||
}
|
||||
|
||||
pub fn settlement_billing_snapshot_schema_version(&self) -> Option<&str> {
|
||||
self.request_metadata_string("billing_snapshot_schema_version")
|
||||
}
|
||||
|
||||
pub fn settlement_billing_snapshot_status(&self) -> Option<&str> {
|
||||
self.request_metadata_string("billing_snapshot_status")
|
||||
}
|
||||
|
||||
pub fn settlement_rate_multiplier(&self) -> Option<f64> {
|
||||
self.request_metadata_number("rate_multiplier")
|
||||
}
|
||||
|
||||
pub fn settlement_is_free_tier(&self) -> Option<bool> {
|
||||
self.request_metadata_bool("is_free_tier")
|
||||
}
|
||||
|
||||
pub fn settlement_input_price_per_1m(&self) -> Option<f64> {
|
||||
self.request_metadata_number("input_price_per_1m")
|
||||
.or_else(|| self.billing_snapshot_resolved_number("input_price_per_1m"))
|
||||
}
|
||||
|
||||
pub fn settlement_output_price_per_1m(&self) -> Option<f64> {
|
||||
self.request_metadata_number("output_price_per_1m")
|
||||
.or_else(|| self.billing_snapshot_resolved_number("output_price_per_1m"))
|
||||
.or(self.output_price_per_1m)
|
||||
}
|
||||
|
||||
pub fn settlement_cache_creation_price_per_1m(&self) -> Option<f64> {
|
||||
self.request_metadata_number("cache_creation_price_per_1m")
|
||||
.or_else(|| self.billing_snapshot_resolved_number("cache_creation_price_per_1m"))
|
||||
}
|
||||
|
||||
pub fn settlement_cache_read_price_per_1m(&self) -> Option<f64> {
|
||||
self.request_metadata_number("cache_read_price_per_1m")
|
||||
.or_else(|| self.billing_snapshot_resolved_number("cache_read_price_per_1m"))
|
||||
}
|
||||
|
||||
pub fn settlement_price_per_request(&self) -> Option<f64> {
|
||||
self.request_metadata_number("price_per_request")
|
||||
.or_else(|| self.billing_snapshot_resolved_number("price_per_request"))
|
||||
}
|
||||
|
||||
pub fn trace_id(&self) -> Option<&str> {
|
||||
self.request_metadata_string("trace_id")
|
||||
}
|
||||
|
||||
pub fn body_ref(&self, field: UsageBodyField) -> Option<&str> {
|
||||
match field {
|
||||
UsageBodyField::RequestBody => self.request_body_ref.as_deref(),
|
||||
UsageBodyField::ProviderRequestBody => self.provider_request_body_ref.as_deref(),
|
||||
UsageBodyField::ResponseBody => self.response_body_ref.as_deref(),
|
||||
UsageBodyField::ClientResponseBody => self.client_response_body_ref.as_deref(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn routing_candidate_id(&self) -> Option<&str> {
|
||||
self.candidate_id
|
||||
.as_deref()
|
||||
.or_else(|| self.request_metadata_string("candidate_id"))
|
||||
}
|
||||
|
||||
pub fn routing_candidate_index(&self) -> Option<u64> {
|
||||
self.candidate_index
|
||||
.or_else(|| self.request_metadata_u64("candidate_index"))
|
||||
}
|
||||
|
||||
pub fn routing_key_name(&self) -> Option<&str> {
|
||||
self.key_name
|
||||
.as_deref()
|
||||
.or_else(|| self.request_metadata_string("key_name"))
|
||||
}
|
||||
|
||||
pub fn routing_planner_kind(&self) -> Option<&str> {
|
||||
self.planner_kind
|
||||
.as_deref()
|
||||
.or_else(|| self.request_metadata_string("planner_kind"))
|
||||
}
|
||||
|
||||
pub fn routing_route_family(&self) -> Option<&str> {
|
||||
self.route_family
|
||||
.as_deref()
|
||||
.or_else(|| self.request_metadata_string("route_family"))
|
||||
}
|
||||
|
||||
pub fn routing_route_kind(&self) -> Option<&str> {
|
||||
self.route_kind
|
||||
.as_deref()
|
||||
.or_else(|| self.request_metadata_string("route_kind"))
|
||||
}
|
||||
|
||||
pub fn routing_execution_path(&self) -> Option<&str> {
|
||||
self.execution_path
|
||||
.as_deref()
|
||||
.or_else(|| self.request_metadata_string("execution_path"))
|
||||
}
|
||||
|
||||
pub fn routing_local_execution_runtime_miss_reason(&self) -> Option<&str> {
|
||||
self.local_execution_runtime_miss_reason
|
||||
.as_deref()
|
||||
.or_else(|| self.request_metadata_string("local_execution_runtime_miss_reason"))
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
@@ -298,6 +482,64 @@ pub struct UsageAuditListQuery {
|
||||
pub model: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum UsageBodyField {
|
||||
RequestBody,
|
||||
ProviderRequestBody,
|
||||
ResponseBody,
|
||||
ClientResponseBody,
|
||||
}
|
||||
|
||||
impl UsageBodyField {
|
||||
pub fn as_ref_key(&self) -> &'static str {
|
||||
match self {
|
||||
Self::RequestBody => "request_body_ref",
|
||||
Self::ProviderRequestBody => "provider_request_body_ref",
|
||||
Self::ResponseBody => "response_body_ref",
|
||||
Self::ClientResponseBody => "client_response_body_ref",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_storage_field(&self) -> &'static str {
|
||||
match self {
|
||||
Self::RequestBody => "request_body",
|
||||
Self::ProviderRequestBody => "provider_request_body",
|
||||
Self::ResponseBody => "response_body",
|
||||
Self::ClientResponseBody => "client_response_body",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn from_storage_field(value: &str) -> Option<Self> {
|
||||
match value {
|
||||
"request_body" => Some(Self::RequestBody),
|
||||
"provider_request_body" => Some(Self::ProviderRequestBody),
|
||||
"response_body" => Some(Self::ResponseBody),
|
||||
"client_response_body" => Some(Self::ClientResponseBody),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn usage_body_ref(request_id: &str, field: UsageBodyField) -> String {
|
||||
format!("usage://request/{request_id}/{}", field.as_storage_field())
|
||||
}
|
||||
|
||||
pub fn parse_usage_body_ref(body_ref: &str) -> Option<(String, UsageBodyField)> {
|
||||
let body_ref = body_ref.trim();
|
||||
let prefix = "usage://request/";
|
||||
let suffix = body_ref.strip_prefix(prefix)?;
|
||||
let (request_id, field) = suffix.rsplit_once('/')?;
|
||||
let request_id = request_id.trim();
|
||||
if request_id.is_empty() {
|
||||
return None;
|
||||
}
|
||||
Some((
|
||||
request_id.to_string(),
|
||||
UsageBodyField::from_storage_field(field.trim())?,
|
||||
))
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait UsageReadRepository: Send + Sync {
|
||||
async fn find_by_id(
|
||||
@@ -310,6 +552,11 @@ pub trait UsageReadRepository: Send + Sync {
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, crate::DataLayerError>;
|
||||
|
||||
async fn resolve_body_ref(
|
||||
&self,
|
||||
body_ref: &str,
|
||||
) -> Result<Option<Value>, crate::DataLayerError>;
|
||||
|
||||
async fn list_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
@@ -384,12 +631,24 @@ pub struct UpsertUsageRecord {
|
||||
pub billing_status: String,
|
||||
pub request_headers: Option<Value>,
|
||||
pub request_body: Option<Value>,
|
||||
pub request_body_ref: Option<String>,
|
||||
pub provider_request_headers: Option<Value>,
|
||||
pub provider_request_body: Option<Value>,
|
||||
pub provider_request_body_ref: Option<String>,
|
||||
pub response_headers: Option<Value>,
|
||||
pub response_body: Option<Value>,
|
||||
pub response_body_ref: Option<String>,
|
||||
pub client_response_headers: Option<Value>,
|
||||
pub client_response_body: Option<Value>,
|
||||
pub client_response_body_ref: Option<String>,
|
||||
pub candidate_id: Option<String>,
|
||||
pub candidate_index: Option<u64>,
|
||||
pub key_name: Option<String>,
|
||||
pub planner_kind: Option<String>,
|
||||
pub route_family: Option<String>,
|
||||
pub route_kind: Option<String>,
|
||||
pub execution_path: Option<String>,
|
||||
pub local_execution_runtime_miss_reason: Option<String>,
|
||||
pub request_metadata: Option<Value>,
|
||||
pub finalized_at_unix_secs: Option<u64>,
|
||||
pub created_at_unix_ms: Option<u64>,
|
||||
@@ -511,9 +770,51 @@ fn parse_timestamp(value: i64, field_name: &str) -> Result<u64, crate::DataLayer
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{StoredRequestUsageAudit, UpsertUsageRecord};
|
||||
use super::{StoredRequestUsageAudit, UpsertUsageRecord, UsageBodyField};
|
||||
use serde_json::json;
|
||||
|
||||
fn sample_usage() -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
"usage-1".to_string(),
|
||||
"req-1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
"OpenAI".to_string(),
|
||||
"gpt-4.1".to_string(),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
Some("openai:chat".to_string()),
|
||||
Some("openai".to_string()),
|
||||
Some("chat".to_string()),
|
||||
false,
|
||||
false,
|
||||
10,
|
||||
20,
|
||||
30,
|
||||
0.1,
|
||||
0.1,
|
||||
Some(200),
|
||||
None,
|
||||
None,
|
||||
Some(120),
|
||||
Some(80),
|
||||
"completed".to_string(),
|
||||
"settled".to_string(),
|
||||
100,
|
||||
101,
|
||||
Some(102),
|
||||
)
|
||||
.expect("usage should build")
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_request_id() {
|
||||
assert!(StoredRequestUsageAudit::new(
|
||||
@@ -644,12 +945,24 @@ mod tests {
|
||||
billing_status: "pending".to_string(),
|
||||
request_headers: Some(json!({"authorization": "Bearer test"})),
|
||||
request_body: Some(json!({"model": "gpt-5"})),
|
||||
request_body_ref: None,
|
||||
provider_request_headers: None,
|
||||
provider_request_body: None,
|
||||
provider_request_body_ref: None,
|
||||
response_headers: None,
|
||||
response_body: None,
|
||||
response_body_ref: None,
|
||||
client_response_headers: None,
|
||||
client_response_body: None,
|
||||
client_response_body_ref: None,
|
||||
candidate_id: None,
|
||||
candidate_index: None,
|
||||
key_name: None,
|
||||
planner_kind: None,
|
||||
route_family: None,
|
||||
route_kind: None,
|
||||
execution_path: None,
|
||||
local_execution_runtime_miss_reason: None,
|
||||
request_metadata: None,
|
||||
finalized_at_unix_secs: None,
|
||||
created_at_unix_ms: Some(100),
|
||||
@@ -658,4 +971,133 @@ mod tests {
|
||||
|
||||
assert!(record.validate().is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settlement_accessors_prefer_typed_metadata() {
|
||||
let mut usage = sample_usage();
|
||||
usage.output_price_per_1m = Some(11.0);
|
||||
usage.request_metadata = Some(json!({
|
||||
"billing_snapshot_schema_version": "v2",
|
||||
"billing_snapshot_status": "resolved",
|
||||
"rate_multiplier": 0.5,
|
||||
"is_free_tier": false,
|
||||
"input_price_per_1m": 3.0,
|
||||
"output_price_per_1m": 9.0,
|
||||
"cache_creation_price_per_1m": 3.75,
|
||||
"cache_read_price_per_1m": 0.3,
|
||||
"price_per_request": 0.02,
|
||||
"billing_snapshot": {
|
||||
"resolved_variables": {
|
||||
"output_price_per_1m": 11.0
|
||||
}
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(
|
||||
usage.settlement_billing_snapshot_schema_version(),
|
||||
Some("v2")
|
||||
);
|
||||
assert_eq!(usage.settlement_billing_snapshot_status(), Some("resolved"));
|
||||
assert_eq!(usage.settlement_rate_multiplier(), Some(0.5));
|
||||
assert_eq!(usage.settlement_is_free_tier(), Some(false));
|
||||
assert_eq!(usage.settlement_input_price_per_1m(), Some(3.0));
|
||||
assert_eq!(usage.settlement_output_price_per_1m(), Some(9.0));
|
||||
assert_eq!(usage.settlement_cache_creation_price_per_1m(), Some(3.75));
|
||||
assert_eq!(usage.settlement_cache_read_price_per_1m(), Some(0.3));
|
||||
assert_eq!(usage.settlement_price_per_request(), Some(0.02));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settlement_accessors_fall_back_to_billing_snapshot_and_legacy_output_price() {
|
||||
let mut usage = sample_usage();
|
||||
usage.output_price_per_1m = Some(15.0);
|
||||
usage.request_metadata = Some(json!({
|
||||
"billing_snapshot": {
|
||||
"resolved_variables": {
|
||||
"input_price_per_1m": 3.0,
|
||||
"cache_creation_price_per_1m": 3.75,
|
||||
"cache_read_price_per_1m": 0.3,
|
||||
"price_per_request": 0.02
|
||||
}
|
||||
}
|
||||
}));
|
||||
|
||||
assert_eq!(usage.settlement_input_price_per_1m(), Some(3.0));
|
||||
assert_eq!(usage.settlement_output_price_per_1m(), Some(15.0));
|
||||
assert_eq!(usage.settlement_cache_creation_price_per_1m(), Some(3.75));
|
||||
assert_eq!(usage.settlement_cache_read_price_per_1m(), Some(0.3));
|
||||
assert_eq!(usage.settlement_price_per_request(), Some(0.02));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn body_ref_and_routing_accessors_prefer_typed_fields() {
|
||||
let mut usage = sample_usage();
|
||||
usage.request_body_ref = Some("usage://request/req-1/request_body".to_string());
|
||||
usage.provider_request_body_ref =
|
||||
Some("usage://request/req-1/provider_request_body".to_string());
|
||||
usage.response_body_ref = Some("usage://request/req-1/response_body".to_string());
|
||||
usage.client_response_body_ref =
|
||||
Some("usage://request/req-1/client_response_body".to_string());
|
||||
usage.candidate_id = Some("cand-typed".to_string());
|
||||
usage.key_name = Some("primary-typed".to_string());
|
||||
usage.planner_kind = Some("claude_cli_sync".to_string());
|
||||
usage.route_family = Some("claude".to_string());
|
||||
usage.route_kind = Some("cli".to_string());
|
||||
usage.execution_path = Some("local_execution_runtime_miss".to_string());
|
||||
usage.local_execution_runtime_miss_reason = Some("all_candidates_skipped".to_string());
|
||||
usage.request_metadata = Some(json!({
|
||||
"request_body_ref": "blob://legacy-request",
|
||||
"provider_request_body_ref": "blob://legacy-provider",
|
||||
"response_body_ref": "blob://legacy-response",
|
||||
"client_response_body_ref": "blob://legacy-client-response",
|
||||
"candidate_id": "cand-legacy",
|
||||
"key_name": "primary-legacy"
|
||||
}));
|
||||
|
||||
assert_eq!(
|
||||
usage.body_ref(UsageBodyField::RequestBody),
|
||||
Some("usage://request/req-1/request_body")
|
||||
);
|
||||
assert_eq!(
|
||||
usage.body_ref(UsageBodyField::ProviderRequestBody),
|
||||
Some("usage://request/req-1/provider_request_body")
|
||||
);
|
||||
assert_eq!(
|
||||
usage.body_ref(UsageBodyField::ResponseBody),
|
||||
Some("usage://request/req-1/response_body")
|
||||
);
|
||||
assert_eq!(
|
||||
usage.body_ref(UsageBodyField::ClientResponseBody),
|
||||
Some("usage://request/req-1/client_response_body")
|
||||
);
|
||||
assert_eq!(usage.routing_candidate_id(), Some("cand-typed"));
|
||||
assert_eq!(usage.routing_key_name(), Some("primary-typed"));
|
||||
assert_eq!(usage.routing_planner_kind(), Some("claude_cli_sync"));
|
||||
assert_eq!(usage.routing_route_family(), Some("claude"));
|
||||
assert_eq!(usage.routing_route_kind(), Some("cli"));
|
||||
assert_eq!(
|
||||
usage.routing_execution_path(),
|
||||
Some("local_execution_runtime_miss")
|
||||
);
|
||||
assert_eq!(
|
||||
usage.routing_local_execution_runtime_miss_reason(),
|
||||
Some("all_candidates_skipped")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn body_ref_accessor_ignores_legacy_metadata_compatibility_keys() {
|
||||
let mut usage = sample_usage();
|
||||
usage.request_metadata = Some(json!({
|
||||
"request_body_ref": "blob://legacy-request",
|
||||
"provider_request_body_ref": "blob://legacy-provider",
|
||||
"response_body_ref": "blob://legacy-response",
|
||||
"client_response_body_ref": "blob://legacy-client-response"
|
||||
}));
|
||||
|
||||
assert_eq!(usage.body_ref(UsageBodyField::RequestBody), None);
|
||||
assert_eq!(usage.body_ref(UsageBodyField::ProviderRequestBody), None);
|
||||
assert_eq!(usage.body_ref(UsageBodyField::ResponseBody), None);
|
||||
assert_eq!(usage.body_ref(UsageBodyField::ClientResponseBody), None);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -13,6 +13,7 @@ aether-wallet.workspace = true
|
||||
async-trait.workspace = true
|
||||
chrono.workspace = true
|
||||
futures-util.workspace = true
|
||||
flate2.workspace = true
|
||||
redis.workspace = true
|
||||
serde.workspace = true
|
||||
serde_json.workspace = true
|
||||
|
||||
4740
crates/aether-data/bootstrap/20260413020000_baseline_v2.sql
Normal file
4740
crates/aether-data/bootstrap/20260413020000_baseline_v2.sql
Normal file
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,147 @@
|
||||
-- Squashed unreleased usage schema split:
|
||||
-- 20260412000000_add_usage_body_blobs.sql
|
||||
-- 20260412010000_add_usage_http_audits.sql
|
||||
-- 20260412020000_add_usage_routing_snapshots.sql
|
||||
-- 20260412030000_add_usage_settlement_snapshots.sql
|
||||
-- 20260413000000_expand_usage_settlement_snapshots_for_pricing.sql
|
||||
-- 20260413010000_mark_usage_legacy_columns_deprecated.sql
|
||||
-- 20260413020000_add_candidate_index_to_usage_routing_snapshots.sql
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.usage_body_blobs (
|
||||
body_ref character varying(160) NOT NULL,
|
||||
request_id character varying(100) NOT NULL,
|
||||
body_field character varying(50) NOT NULL,
|
||||
payload_gzip bytea NOT NULL,
|
||||
created_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
updated_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
CONSTRAINT usage_body_blobs_pkey PRIMARY KEY (body_ref),
|
||||
CONSTRAINT usage_body_blobs_request_id_field_key UNIQUE (request_id, body_field),
|
||||
CONSTRAINT usage_body_blobs_request_id_fkey
|
||||
FOREIGN KEY (request_id)
|
||||
REFERENCES public.usage(request_id)
|
||||
ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS ix_usage_body_blobs_request_id
|
||||
ON public.usage_body_blobs USING btree (request_id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.usage_http_audits (
|
||||
request_id character varying(100) NOT NULL,
|
||||
request_headers json,
|
||||
provider_request_headers json,
|
||||
response_headers json,
|
||||
client_response_headers json,
|
||||
request_body_ref character varying(160),
|
||||
provider_request_body_ref character varying(160),
|
||||
response_body_ref character varying(160),
|
||||
client_response_body_ref character varying(160),
|
||||
body_capture_mode character varying(32) DEFAULT 'none' NOT NULL,
|
||||
created_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
updated_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
CONSTRAINT usage_http_audits_pkey PRIMARY KEY (request_id),
|
||||
CONSTRAINT usage_http_audits_request_id_fkey
|
||||
FOREIGN KEY (request_id)
|
||||
REFERENCES public.usage(request_id)
|
||||
ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS ix_usage_http_audits_updated_at
|
||||
ON public.usage_http_audits USING btree (updated_at);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.usage_routing_snapshots (
|
||||
request_id character varying(100) NOT NULL,
|
||||
candidate_id character varying(160),
|
||||
candidate_index integer,
|
||||
key_name character varying(255),
|
||||
planner_kind character varying(120),
|
||||
route_family character varying(80),
|
||||
route_kind character varying(80),
|
||||
execution_path character varying(80),
|
||||
local_execution_runtime_miss_reason character varying(120),
|
||||
selected_provider_id character varying(100),
|
||||
selected_endpoint_id character varying(100),
|
||||
selected_provider_api_key_id character varying(100),
|
||||
has_format_conversion boolean,
|
||||
created_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
updated_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
CONSTRAINT usage_routing_snapshots_pkey PRIMARY KEY (request_id),
|
||||
CONSTRAINT usage_routing_snapshots_request_id_fkey
|
||||
FOREIGN KEY (request_id)
|
||||
REFERENCES public.usage(request_id)
|
||||
ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS ix_usage_routing_snapshots_route_family_kind
|
||||
ON public.usage_routing_snapshots USING btree (route_family, route_kind);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS ix_usage_routing_snapshots_candidate_id
|
||||
ON public.usage_routing_snapshots USING btree (candidate_id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS public.usage_settlement_snapshots (
|
||||
request_id character varying(100) NOT NULL,
|
||||
billing_status character varying(20) NOT NULL,
|
||||
wallet_id character varying(36),
|
||||
wallet_balance_before numeric(20,8),
|
||||
wallet_balance_after numeric(20,8),
|
||||
wallet_recharge_balance_before numeric(20,8),
|
||||
wallet_recharge_balance_after numeric(20,8),
|
||||
wallet_gift_balance_before numeric(20,8),
|
||||
wallet_gift_balance_after numeric(20,8),
|
||||
provider_monthly_used_usd numeric(20,8),
|
||||
finalized_at timestamp with time zone,
|
||||
created_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
updated_at timestamp with time zone DEFAULT now() NOT NULL,
|
||||
CONSTRAINT usage_settlement_snapshots_pkey PRIMARY KEY (request_id),
|
||||
CONSTRAINT usage_settlement_snapshots_request_id_fkey
|
||||
FOREIGN KEY (request_id)
|
||||
REFERENCES public.usage(request_id)
|
||||
ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS ix_usage_settlement_snapshots_wallet_id
|
||||
ON public.usage_settlement_snapshots USING btree (wallet_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS ix_usage_settlement_snapshots_billing_status
|
||||
ON public.usage_settlement_snapshots USING btree (billing_status);
|
||||
|
||||
ALTER TABLE IF EXISTS public.usage_settlement_snapshots
|
||||
ADD COLUMN IF NOT EXISTS billing_snapshot_schema_version character varying(20),
|
||||
ADD COLUMN IF NOT EXISTS billing_snapshot_status character varying(20),
|
||||
ADD COLUMN IF NOT EXISTS rate_multiplier numeric(10,6),
|
||||
ADD COLUMN IF NOT EXISTS is_free_tier boolean,
|
||||
ADD COLUMN IF NOT EXISTS input_price_per_1m numeric(20,8),
|
||||
ADD COLUMN IF NOT EXISTS output_price_per_1m numeric(20,8),
|
||||
ADD COLUMN IF NOT EXISTS cache_creation_price_per_1m numeric(20,8),
|
||||
ADD COLUMN IF NOT EXISTS cache_read_price_per_1m numeric(20,8),
|
||||
ADD COLUMN IF NOT EXISTS price_per_request numeric(20,8);
|
||||
|
||||
COMMENT ON COLUMN public.usage.wallet_id IS
|
||||
'DEPRECATED: settlement owner moved to public.usage_settlement_snapshots.wallet_id. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.wallet_balance_before IS
|
||||
'DEPRECATED: settlement owner moved to public.usage_settlement_snapshots.wallet_balance_before. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.wallet_balance_after IS
|
||||
'DEPRECATED: settlement owner moved to public.usage_settlement_snapshots.wallet_balance_after. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.wallet_recharge_balance_before IS
|
||||
'DEPRECATED: settlement owner moved to public.usage_settlement_snapshots.wallet_recharge_balance_before. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.wallet_recharge_balance_after IS
|
||||
'DEPRECATED: settlement owner moved to public.usage_settlement_snapshots.wallet_recharge_balance_after. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.wallet_gift_balance_before IS
|
||||
'DEPRECATED: settlement owner moved to public.usage_settlement_snapshots.wallet_gift_balance_before. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.wallet_gift_balance_after IS
|
||||
'DEPRECATED: settlement owner moved to public.usage_settlement_snapshots.wallet_gift_balance_after. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.rate_multiplier IS
|
||||
'DEPRECATED: settlement pricing owner moved to public.usage_settlement_snapshots.rate_multiplier. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.input_price_per_1m IS
|
||||
'DEPRECATED: settlement pricing owner moved to public.usage_settlement_snapshots.input_price_per_1m. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.output_price_per_1m IS
|
||||
'DEPRECATED: settlement pricing owner moved to public.usage_settlement_snapshots.output_price_per_1m. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.cache_creation_price_per_1m IS
|
||||
'DEPRECATED: settlement pricing owner moved to public.usage_settlement_snapshots.cache_creation_price_per_1m. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.cache_read_price_per_1m IS
|
||||
'DEPRECATED: settlement pricing owner moved to public.usage_settlement_snapshots.cache_read_price_per_1m. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.price_per_request IS
|
||||
'DEPRECATED: settlement pricing owner moved to public.usage_settlement_snapshots.price_per_request. Legacy compatibility only; do not write new values.';
|
||||
COMMENT ON COLUMN public.usage.username IS
|
||||
'DEPRECATED: display cache only. Prefer join-time lookup from user/auth records. Legacy compatibility only.';
|
||||
COMMENT ON COLUMN public.usage.api_key_name IS
|
||||
'DEPRECATED: display cache only. Prefer join-time lookup from API key records. Legacy compatibility only.';
|
||||
@@ -2,11 +2,44 @@ use std::collections::{HashMap, HashSet};
|
||||
|
||||
use sqlx::{
|
||||
migrate::{Migrate, MigrateError, Migrator},
|
||||
PgPool,
|
||||
query, query_scalar, Connection, PgConnection, PgPool,
|
||||
};
|
||||
use tracing::{error, info, warn};
|
||||
|
||||
static MIGRATOR: Migrator = sqlx::migrate!("./migrations");
|
||||
static BASELINE_V2_SQL: &str = include_str!("../bootstrap/20260413020000_baseline_v2.sql");
|
||||
const BASELINE_V2_CUTOFF_VERSION: i64 = 20260413020000;
|
||||
const MIGRATIONS_TABLE_EXISTS_SQL: &str =
|
||||
"SELECT to_regclass('public._sqlx_migrations') IS NOT NULL";
|
||||
const EMPTY_DATABASE_USER_TABLE_COUNT_SQL: &str = r#"
|
||||
SELECT COUNT(*)::BIGINT
|
||||
FROM information_schema.tables
|
||||
WHERE table_schema = 'public'
|
||||
AND table_type = 'BASE TABLE'
|
||||
AND table_name <> '_sqlx_migrations'
|
||||
"#;
|
||||
const INSERT_APPLIED_MIGRATION_SQL: &str = r#"
|
||||
INSERT INTO _sqlx_migrations (
|
||||
version,
|
||||
description,
|
||||
success,
|
||||
checksum,
|
||||
execution_time
|
||||
) VALUES (
|
||||
$1,
|
||||
$2,
|
||||
TRUE,
|
||||
$3,
|
||||
0
|
||||
)
|
||||
ON CONFLICT (version) DO NOTHING
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct PendingMigrationInfo {
|
||||
pub version: i64,
|
||||
pub description: String,
|
||||
}
|
||||
|
||||
/// Run all pending migrations embedded at compile time from `migrations/`.
|
||||
pub async fn run_migrations(pool: &PgPool) -> Result<(), MigrateError> {
|
||||
@@ -16,7 +49,7 @@ pub async fn run_migrations(pool: &PgPool) -> Result<(), MigrateError> {
|
||||
conn.lock().await?;
|
||||
}
|
||||
|
||||
let result = run_migrations_locked(&mut *conn).await;
|
||||
let result = run_migrations_locked(&mut conn).await;
|
||||
|
||||
if MIGRATOR.locking {
|
||||
match conn.unlock().await {
|
||||
@@ -34,11 +67,41 @@ pub async fn run_migrations(pool: &PgPool) -> Result<(), MigrateError> {
|
||||
result
|
||||
}
|
||||
|
||||
async fn run_migrations_locked<C>(conn: &mut C) -> Result<(), MigrateError>
|
||||
where
|
||||
C: Migrate,
|
||||
{
|
||||
pub async fn pending_migrations(pool: &PgPool) -> Result<Vec<PendingMigrationInfo>, MigrateError> {
|
||||
let mut conn = pool.acquire().await?;
|
||||
pending_migrations_locked(&mut conn).await
|
||||
}
|
||||
|
||||
pub async fn prepare_database_for_startup(
|
||||
pool: &PgPool,
|
||||
) -> Result<Vec<PendingMigrationInfo>, MigrateError> {
|
||||
let mut conn = pool.acquire().await?;
|
||||
|
||||
if MIGRATOR.locking {
|
||||
conn.lock().await?;
|
||||
}
|
||||
|
||||
let result = prepare_database_for_startup_locked(&mut conn).await;
|
||||
|
||||
if MIGRATOR.locking {
|
||||
match conn.unlock().await {
|
||||
Ok(()) => {}
|
||||
Err(unlock_error) if result.is_ok() => return Err(unlock_error),
|
||||
Err(unlock_error) => {
|
||||
warn!(
|
||||
error = %unlock_error,
|
||||
"database migration lock release failed after startup preparation error"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
async fn run_migrations_locked(conn: &mut PgConnection) -> Result<(), MigrateError> {
|
||||
conn.ensure_migrations_table().await?;
|
||||
bootstrap_empty_database_to_baseline_v2(conn).await?;
|
||||
|
||||
if let Some(version) = conn.dirty_version().await? {
|
||||
error!(version, "database migration state is dirty");
|
||||
@@ -118,6 +181,125 @@ where
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn prepare_database_for_startup_locked(
|
||||
conn: &mut PgConnection,
|
||||
) -> Result<Vec<PendingMigrationInfo>, MigrateError> {
|
||||
conn.ensure_migrations_table().await?;
|
||||
bootstrap_empty_database_to_baseline_v2(conn).await?;
|
||||
pending_migrations_locked(conn).await
|
||||
}
|
||||
|
||||
async fn pending_migrations_locked(
|
||||
conn: &mut PgConnection,
|
||||
) -> Result<Vec<PendingMigrationInfo>, MigrateError> {
|
||||
if !migrations_table_exists(conn).await? {
|
||||
return Ok(all_up_migrations());
|
||||
}
|
||||
|
||||
if let Some(version) = conn.dirty_version().await? {
|
||||
error!(version, "database migration state is dirty");
|
||||
return Err(MigrateError::Dirty(version));
|
||||
}
|
||||
|
||||
let applied_migrations = conn.list_applied_migrations().await?;
|
||||
validate_applied_migrations(&applied_migrations)?;
|
||||
Ok(pending_migrations_from_applied(&applied_migrations))
|
||||
}
|
||||
|
||||
async fn bootstrap_empty_database_to_baseline_v2(
|
||||
conn: &mut PgConnection,
|
||||
) -> Result<(), MigrateError> {
|
||||
if !should_bootstrap_baseline_v2(conn).await? {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
let migrations = baseline_v2_migrations()?;
|
||||
info!(
|
||||
cutoff_version = BASELINE_V2_CUTOFF_VERSION,
|
||||
stamped_migrations = migrations.len(),
|
||||
"bootstrapping empty database from baseline_v2"
|
||||
);
|
||||
|
||||
let mut tx = conn.begin().await?;
|
||||
sqlx::raw_sql(BASELINE_V2_SQL).execute(&mut *tx).await?;
|
||||
for migration in migrations {
|
||||
query(INSERT_APPLIED_MIGRATION_SQL)
|
||||
.bind(migration.version)
|
||||
.bind(migration.description.as_ref())
|
||||
.bind(migration.checksum.as_ref())
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
}
|
||||
tx.commit().await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn migrations_table_exists(conn: &mut PgConnection) -> Result<bool, MigrateError> {
|
||||
let exists: bool = query_scalar(MIGRATIONS_TABLE_EXISTS_SQL)
|
||||
.fetch_one(&mut *conn)
|
||||
.await?;
|
||||
Ok(exists)
|
||||
}
|
||||
|
||||
async fn should_bootstrap_baseline_v2(conn: &mut PgConnection) -> Result<bool, MigrateError> {
|
||||
let applied_migrations = conn.list_applied_migrations().await?;
|
||||
if !applied_migrations.is_empty() {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let user_table_count: i64 = query_scalar(EMPTY_DATABASE_USER_TABLE_COUNT_SQL)
|
||||
.fetch_one(&mut *conn)
|
||||
.await?;
|
||||
Ok(user_table_count == 0)
|
||||
}
|
||||
|
||||
fn baseline_v2_migrations() -> Result<Vec<&'static sqlx::migrate::Migration>, MigrateError> {
|
||||
let migrations = MIGRATOR
|
||||
.iter()
|
||||
.filter(|migration| migration.migration_type.is_up_migration())
|
||||
.filter(|migration| migration.version <= BASELINE_V2_CUTOFF_VERSION)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
if migrations.is_empty() {
|
||||
return Err(MigrateError::Source(Box::new(std::io::Error::other(
|
||||
"baseline_v2 cutoff does not match any embedded migrations",
|
||||
))));
|
||||
}
|
||||
|
||||
Ok(migrations)
|
||||
}
|
||||
|
||||
fn all_up_migrations() -> Vec<PendingMigrationInfo> {
|
||||
MIGRATOR
|
||||
.iter()
|
||||
.filter(|migration| migration.migration_type.is_up_migration())
|
||||
.map(|migration| PendingMigrationInfo {
|
||||
version: migration.version,
|
||||
description: migration.description.to_string(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn pending_migrations_from_applied(
|
||||
applied_migrations: &[sqlx::migrate::AppliedMigration],
|
||||
) -> Vec<PendingMigrationInfo> {
|
||||
let applied_versions: HashSet<_> = applied_migrations
|
||||
.iter()
|
||||
.map(|migration| migration.version)
|
||||
.collect();
|
||||
|
||||
MIGRATOR
|
||||
.iter()
|
||||
.filter(|migration| migration.migration_type.is_up_migration())
|
||||
.filter(|migration| !applied_versions.contains(&migration.version))
|
||||
.map(|migration| PendingMigrationInfo {
|
||||
version: migration.version,
|
||||
description: migration.description.to_string(),
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn validate_applied_migrations(
|
||||
applied_migrations: &[sqlx::migrate::AppliedMigration],
|
||||
) -> Result<(), MigrateError> {
|
||||
@@ -165,7 +347,14 @@ fn validate_applied_migrations(
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::MIGRATOR;
|
||||
use std::borrow::Cow;
|
||||
|
||||
use sqlx::migrate::AppliedMigration;
|
||||
|
||||
use super::{
|
||||
all_up_migrations, baseline_v2_migrations, pending_migrations_from_applied,
|
||||
BASELINE_V2_SQL, MIGRATOR,
|
||||
};
|
||||
|
||||
#[test]
|
||||
fn baseline_migration_restores_search_path_for_sqlx_bookkeeping() {
|
||||
@@ -199,4 +388,109 @@ mod tests {
|
||||
"baseline migration must not persist a restored search_path at session scope",
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn baseline_v2_bootstrap_covers_current_cutoff_versions() {
|
||||
let versions = baseline_v2_migrations()
|
||||
.expect("baseline_v2 migrations should resolve")
|
||||
.into_iter()
|
||||
.map(|migration| migration.version)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(
|
||||
versions,
|
||||
vec![
|
||||
20260403000000,
|
||||
20260406000000,
|
||||
20260410000000,
|
||||
20260413020000,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn baseline_v2_sql_includes_usage_body_blobs() {
|
||||
assert!(BASELINE_V2_SQL.contains("CREATE TABLE IF NOT EXISTS public.usage_body_blobs"));
|
||||
assert!(BASELINE_V2_SQL.contains("ix_usage_body_blobs_request_id"));
|
||||
assert!(BASELINE_V2_SQL.contains("CREATE TABLE IF NOT EXISTS public.usage_http_audits"));
|
||||
assert!(
|
||||
BASELINE_V2_SQL.contains("CREATE TABLE IF NOT EXISTS public.usage_routing_snapshots")
|
||||
);
|
||||
assert!(BASELINE_V2_SQL
|
||||
.contains("CREATE TABLE IF NOT EXISTS public.usage_settlement_snapshots"));
|
||||
assert!(BASELINE_V2_SQL.contains("billing_snapshot_schema_version"));
|
||||
assert!(BASELINE_V2_SQL.contains("price_per_request"));
|
||||
assert!(BASELINE_V2_SQL.contains("candidate_index integer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deprecation_migration_and_baseline_mark_legacy_usage_columns() {
|
||||
let migration = MIGRATOR
|
||||
.iter()
|
||||
.find(|migration| migration.version == 20260413020000)
|
||||
.expect("deprecation migration should be embedded");
|
||||
|
||||
assert!(migration
|
||||
.sql
|
||||
.contains("COMMENT ON COLUMN public.usage.output_price_per_1m"));
|
||||
assert!(migration
|
||||
.sql
|
||||
.contains("COMMENT ON COLUMN public.usage.wallet_id"));
|
||||
assert!(migration
|
||||
.sql
|
||||
.contains("COMMENT ON COLUMN public.usage.username"));
|
||||
assert!(migration
|
||||
.sql
|
||||
.contains("COMMENT ON COLUMN public.usage.api_key_name"));
|
||||
assert!(BASELINE_V2_SQL.contains("COMMENT ON COLUMN public.usage.output_price_per_1m"));
|
||||
assert!(BASELINE_V2_SQL.contains("COMMENT ON COLUMN public.usage.wallet_id"));
|
||||
assert!(BASELINE_V2_SQL.contains("COMMENT ON COLUMN public.usage.username"));
|
||||
assert!(BASELINE_V2_SQL.contains("COMMENT ON COLUMN public.usage.api_key_name"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_migrations_from_applied_returns_all_versions_when_none_applied() {
|
||||
let pending = pending_migrations_from_applied(&[]);
|
||||
assert_eq!(pending, all_up_migrations());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_migrations_from_applied_skips_versions_already_applied() {
|
||||
let applied = vec![
|
||||
AppliedMigration {
|
||||
version: 20260403000000,
|
||||
checksum: Cow::Borrowed(&[]),
|
||||
},
|
||||
AppliedMigration {
|
||||
version: 20260406000000,
|
||||
checksum: Cow::Borrowed(&[]),
|
||||
},
|
||||
];
|
||||
|
||||
let pending_versions = pending_migrations_from_applied(&applied)
|
||||
.into_iter()
|
||||
.map(|migration| migration.version)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(pending_versions, vec![20260410000000, 20260413020000]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pending_migrations_from_applied_is_empty_after_baseline_v2_stamp() {
|
||||
let applied = baseline_v2_migrations()
|
||||
.expect("baseline_v2 migrations should resolve")
|
||||
.into_iter()
|
||||
.map(|migration| AppliedMigration {
|
||||
version: migration.version,
|
||||
checksum: migration.checksum.clone(),
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let pending = pending_migrations_from_applied(&applied);
|
||||
|
||||
assert!(
|
||||
pending.is_empty(),
|
||||
"baseline_v2-stamped empty databases should not require a manual migration before first startup"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -385,15 +385,15 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
StoredAuthApiKeySnapshot {
|
||||
api_key_id: record.api_key_id.clone(),
|
||||
api_key_name: record.name.clone(),
|
||||
api_key_is_active: true,
|
||||
api_key_is_active: record.is_active,
|
||||
api_key_is_locked: false,
|
||||
api_key_is_standalone: false,
|
||||
api_key_rate_limit: Some(record.rate_limit),
|
||||
api_key_concurrent_limit: Some(record.concurrent_limit),
|
||||
api_key_expires_at_unix_secs: None,
|
||||
api_key_allowed_providers: None,
|
||||
api_key_allowed_api_formats: None,
|
||||
api_key_allowed_models: None,
|
||||
api_key_expires_at_unix_secs: record.expires_at_unix_secs,
|
||||
api_key_allowed_providers: record.allowed_providers.clone(),
|
||||
api_key_allowed_api_formats: record.allowed_api_formats.clone(),
|
||||
api_key_allowed_models: record.allowed_models.clone(),
|
||||
..template
|
||||
}
|
||||
} else {
|
||||
@@ -413,15 +413,24 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
None,
|
||||
record.api_key_id.clone(),
|
||||
record.name.clone(),
|
||||
true,
|
||||
record.is_active,
|
||||
false,
|
||||
false,
|
||||
Some(record.rate_limit),
|
||||
Some(record.concurrent_limit),
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
record.expires_at_unix_secs.map(|value| value as i64),
|
||||
record
|
||||
.allowed_providers
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
record
|
||||
.allowed_api_formats
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
record
|
||||
.allowed_models
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
)?
|
||||
};
|
||||
|
||||
@@ -431,17 +440,26 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
record.key_hash.clone(),
|
||||
record.key_encrypted,
|
||||
record.name,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
record
|
||||
.allowed_providers
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
record
|
||||
.allowed_api_formats
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
record
|
||||
.allowed_models
|
||||
.as_ref()
|
||||
.map(|value| serde_json::json!(value)),
|
||||
Some(record.rate_limit),
|
||||
Some(record.concurrent_limit),
|
||||
None,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
0,
|
||||
0.0,
|
||||
record.force_capabilities,
|
||||
record.is_active,
|
||||
record.expires_at_unix_secs.map(|value| value as i64),
|
||||
record.auto_delete_on_expiry,
|
||||
record.total_requests as i64,
|
||||
record.total_cost_usd,
|
||||
false,
|
||||
)?;
|
||||
|
||||
@@ -487,12 +505,12 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
StoredAuthApiKeySnapshot {
|
||||
api_key_id: record.api_key_id.clone(),
|
||||
api_key_name: record.name.clone(),
|
||||
api_key_is_active: true,
|
||||
api_key_is_active: record.is_active,
|
||||
api_key_is_locked: false,
|
||||
api_key_is_standalone: true,
|
||||
api_key_rate_limit: Some(record.rate_limit),
|
||||
api_key_concurrent_limit: Some(record.concurrent_limit),
|
||||
api_key_expires_at_unix_secs: None,
|
||||
api_key_expires_at_unix_secs: record.expires_at_unix_secs,
|
||||
api_key_allowed_providers: record.allowed_providers.clone(),
|
||||
api_key_allowed_api_formats: record.allowed_api_formats.clone(),
|
||||
api_key_allowed_models: record.allowed_models.clone(),
|
||||
@@ -515,12 +533,12 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
None,
|
||||
record.api_key_id.clone(),
|
||||
record.name.clone(),
|
||||
true,
|
||||
record.is_active,
|
||||
false,
|
||||
true,
|
||||
Some(record.rate_limit),
|
||||
Some(record.concurrent_limit),
|
||||
None,
|
||||
record.expires_at_unix_secs.map(|value| value as i64),
|
||||
record
|
||||
.allowed_providers
|
||||
.as_ref()
|
||||
@@ -556,12 +574,12 @@ impl AuthApiKeyWriteRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
.map(|value| serde_json::json!(value)),
|
||||
Some(record.rate_limit),
|
||||
Some(record.concurrent_limit),
|
||||
None,
|
||||
true,
|
||||
None,
|
||||
false,
|
||||
0,
|
||||
0.0,
|
||||
record.force_capabilities,
|
||||
record.is_active,
|
||||
record.expires_at_unix_secs.map(|value| value as i64),
|
||||
record.auto_delete_on_expiry,
|
||||
record.total_requests as i64,
|
||||
record.total_cost_usd,
|
||||
true,
|
||||
)?;
|
||||
|
||||
|
||||
@@ -311,10 +311,14 @@ INSERT INTO api_keys (
|
||||
key_hash,
|
||||
key_encrypted,
|
||||
name,
|
||||
allowed_providers,
|
||||
allowed_api_formats,
|
||||
allowed_models,
|
||||
rate_limit,
|
||||
concurrent_limit,
|
||||
force_capabilities,
|
||||
is_active,
|
||||
expires_at,
|
||||
is_locked,
|
||||
is_standalone,
|
||||
auto_delete_on_expiry,
|
||||
@@ -331,13 +335,17 @@ VALUES (
|
||||
$5,
|
||||
$6,
|
||||
$7,
|
||||
NULL,
|
||||
TRUE,
|
||||
$8,
|
||||
$9,
|
||||
$10,
|
||||
$11,
|
||||
$12,
|
||||
$13,
|
||||
FALSE,
|
||||
FALSE,
|
||||
FALSE,
|
||||
0,
|
||||
0,
|
||||
$14,
|
||||
$15,
|
||||
$16,
|
||||
NOW(),
|
||||
NOW()
|
||||
)
|
||||
@@ -375,6 +383,7 @@ INSERT INTO api_keys (
|
||||
concurrent_limit,
|
||||
force_capabilities,
|
||||
is_active,
|
||||
expires_at,
|
||||
is_locked,
|
||||
is_standalone,
|
||||
auto_delete_on_expiry,
|
||||
@@ -394,13 +403,14 @@ VALUES (
|
||||
$8,
|
||||
$9,
|
||||
$10,
|
||||
NULL,
|
||||
TRUE,
|
||||
$11,
|
||||
$12,
|
||||
$13,
|
||||
FALSE,
|
||||
TRUE,
|
||||
FALSE,
|
||||
0,
|
||||
0,
|
||||
$14,
|
||||
$15,
|
||||
$16,
|
||||
NOW(),
|
||||
NOW()
|
||||
)
|
||||
@@ -920,14 +930,46 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
|
||||
&self,
|
||||
record: CreateUserApiKeyRecord,
|
||||
) -> Result<Option<StoredAuthApiKeyExportRecord>, DataLayerError> {
|
||||
let allowed_providers = record
|
||||
.allowed_providers
|
||||
.map(serde_json::to_value)
|
||||
.transpose()
|
||||
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
|
||||
let allowed_api_formats = record
|
||||
.allowed_api_formats
|
||||
.map(serde_json::to_value)
|
||||
.transpose()
|
||||
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
|
||||
let allowed_models = record
|
||||
.allowed_models
|
||||
.map(serde_json::to_value)
|
||||
.transpose()
|
||||
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
|
||||
let expires_at = record
|
||||
.expires_at_unix_secs
|
||||
.map(|value| {
|
||||
chrono::DateTime::<chrono::Utc>::from_timestamp(value as i64, 0).ok_or_else(|| {
|
||||
DataLayerError::UnexpectedValue(format!("invalid api_keys.expires_at: {value}"))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let row = sqlx::query(CREATE_USER_API_KEY_SQL)
|
||||
.bind(record.api_key_id)
|
||||
.bind(record.user_id)
|
||||
.bind(record.key_hash)
|
||||
.bind(record.key_encrypted)
|
||||
.bind(record.name)
|
||||
.bind(allowed_providers)
|
||||
.bind(allowed_api_formats)
|
||||
.bind(allowed_models)
|
||||
.bind(record.rate_limit)
|
||||
.bind(record.concurrent_limit)
|
||||
.bind(record.force_capabilities)
|
||||
.bind(record.is_active)
|
||||
.bind(expires_at)
|
||||
.bind(record.auto_delete_on_expiry)
|
||||
.bind(record.total_requests as i64)
|
||||
.bind(record.total_cost_usd)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
@@ -953,6 +995,14 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
|
||||
.map(serde_json::to_value)
|
||||
.transpose()
|
||||
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
|
||||
let expires_at = record
|
||||
.expires_at_unix_secs
|
||||
.map(|value| {
|
||||
chrono::DateTime::<chrono::Utc>::from_timestamp(value as i64, 0).ok_or_else(|| {
|
||||
DataLayerError::UnexpectedValue(format!("invalid api_keys.expires_at: {value}"))
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
let row = sqlx::query(CREATE_STANDALONE_API_KEY_SQL)
|
||||
.bind(record.api_key_id)
|
||||
.bind(record.user_id)
|
||||
@@ -964,6 +1014,12 @@ impl AuthApiKeyWriteRepository for SqlxAuthApiKeySnapshotReadRepository {
|
||||
.bind(allowed_models)
|
||||
.bind(record.rate_limit)
|
||||
.bind(record.concurrent_limit)
|
||||
.bind(record.force_capabilities)
|
||||
.bind(record.is_active)
|
||||
.bind(expires_at)
|
||||
.bind(record.auto_delete_on_expiry)
|
||||
.bind(record.total_requests as i64)
|
||||
.bind(record.total_cost_usd)
|
||||
.fetch_optional(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
|
||||
@@ -352,15 +352,24 @@ pub struct StandaloneApiKeyExportListQuery {
|
||||
pub is_active: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct CreateUserApiKeyRecord {
|
||||
pub user_id: String,
|
||||
pub api_key_id: String,
|
||||
pub key_hash: String,
|
||||
pub key_encrypted: Option<String>,
|
||||
pub name: Option<String>,
|
||||
pub allowed_providers: Option<Vec<String>>,
|
||||
pub allowed_api_formats: Option<Vec<String>>,
|
||||
pub allowed_models: Option<Vec<String>>,
|
||||
pub rate_limit: i32,
|
||||
pub concurrent_limit: i32,
|
||||
pub force_capabilities: Option<serde_json::Value>,
|
||||
pub is_active: bool,
|
||||
pub expires_at_unix_secs: Option<u64>,
|
||||
pub auto_delete_on_expiry: bool,
|
||||
pub total_requests: u64,
|
||||
pub total_cost_usd: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
@@ -371,7 +380,7 @@ pub struct UpdateUserApiKeyBasicRecord {
|
||||
pub rate_limit: Option<i32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
#[derive(Debug, Clone, PartialEq)]
|
||||
pub struct CreateStandaloneApiKeyRecord {
|
||||
pub user_id: String,
|
||||
pub api_key_id: String,
|
||||
@@ -383,6 +392,12 @@ pub struct CreateStandaloneApiKeyRecord {
|
||||
pub allowed_models: Option<Vec<String>>,
|
||||
pub rate_limit: i32,
|
||||
pub concurrent_limit: i32,
|
||||
pub force_capabilities: Option<serde_json::Value>,
|
||||
pub is_active: bool,
|
||||
pub expires_at_unix_secs: Option<u64>,
|
||||
pub auto_delete_on_expiry: bool,
|
||||
pub total_requests: u64,
|
||||
pub total_cost_usd: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
|
||||
@@ -46,6 +46,7 @@ impl InMemorySettlementWalletStore {
|
||||
pub struct InMemorySettlementRepository {
|
||||
wallets: InMemorySettlementWalletStore,
|
||||
provider_monthly_used: RwLock<BTreeMap<String, f64>>,
|
||||
settlements: RwLock<BTreeMap<String, StoredUsageSettlement>>,
|
||||
}
|
||||
|
||||
impl InMemorySettlementRepository {
|
||||
@@ -56,6 +57,7 @@ impl InMemorySettlementRepository {
|
||||
Self {
|
||||
wallets: InMemorySettlementWalletStore::seeded(items),
|
||||
provider_monthly_used: RwLock::new(BTreeMap::new()),
|
||||
settlements: RwLock::new(BTreeMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -63,6 +65,7 @@ impl InMemorySettlementRepository {
|
||||
Self {
|
||||
wallets: InMemorySettlementWalletStore::Shared(wallet_repository),
|
||||
provider_monthly_used: RwLock::new(BTreeMap::new()),
|
||||
settlements: RwLock::new(BTreeMap::new()),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -75,7 +78,13 @@ impl SettlementWriteRepository for InMemorySettlementRepository {
|
||||
) -> Result<Option<StoredUsageSettlement>, DataLayerError> {
|
||||
input.validate()?;
|
||||
if input.billing_status != "pending" {
|
||||
return Ok(Some(StoredUsageSettlement {
|
||||
let existing = self
|
||||
.settlements
|
||||
.read()
|
||||
.expect("settlement snapshot lock")
|
||||
.get(&input.request_id)
|
||||
.cloned();
|
||||
return Ok(Some(existing.unwrap_or(StoredUsageSettlement {
|
||||
request_id: input.request_id,
|
||||
wallet_id: None,
|
||||
billing_status: input.billing_status,
|
||||
@@ -87,7 +96,7 @@ impl SettlementWriteRepository for InMemorySettlementRepository {
|
||||
wallet_gift_balance_after: None,
|
||||
provider_monthly_used_usd: None,
|
||||
finalized_at_unix_secs: input.finalized_at_unix_secs,
|
||||
}));
|
||||
})));
|
||||
}
|
||||
|
||||
let final_billing_status = if input.status == "completed" {
|
||||
@@ -172,6 +181,11 @@ impl SettlementWriteRepository for InMemorySettlementRepository {
|
||||
}
|
||||
}
|
||||
|
||||
self.settlements
|
||||
.write()
|
||||
.expect("settlement snapshot lock")
|
||||
.insert(settlement.request_id.clone(), settlement.clone());
|
||||
|
||||
Ok(Some(settlement))
|
||||
}
|
||||
}
|
||||
@@ -225,4 +239,42 @@ mod tests {
|
||||
assert_eq!(settlement.wallet_balance_after, Some(9.0));
|
||||
assert_eq!(settlement.provider_monthly_used_usd, Some(1.5));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn returns_stored_snapshot_when_usage_is_already_finalized() {
|
||||
let repository = InMemorySettlementRepository::seed(vec![sample_wallet()]);
|
||||
let settled = repository
|
||||
.settle_usage(UsageSettlementInput {
|
||||
request_id: "req-2".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("key-1".to_string()),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
status: "completed".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
total_cost_usd: 2.0,
|
||||
actual_total_cost_usd: 1.0,
|
||||
finalized_at_unix_secs: Some(250),
|
||||
})
|
||||
.await
|
||||
.expect("settlement should succeed")
|
||||
.expect("settlement should exist");
|
||||
|
||||
let replay = repository
|
||||
.settle_usage(UsageSettlementInput {
|
||||
request_id: "req-2".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("key-1".to_string()),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
status: "completed".to_string(),
|
||||
billing_status: "settled".to_string(),
|
||||
total_cost_usd: 2.0,
|
||||
actual_total_cost_usd: 1.0,
|
||||
finalized_at_unix_secs: Some(250),
|
||||
})
|
||||
.await
|
||||
.expect("replay should succeed")
|
||||
.expect("snapshot should exist");
|
||||
|
||||
assert_eq!(replay, settled);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -6,6 +6,49 @@ use crate::error::SqlxResultExt;
|
||||
use crate::postgres::PostgresTransactionRunner;
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_USAGE_FOR_SETTLEMENT_SQL: &str = r#"
|
||||
SELECT
|
||||
usage_record.request_id,
|
||||
COALESCE(usage_settlement_snapshots.wallet_id, usage_record.wallet_id) AS wallet_id,
|
||||
usage_record.billing_status,
|
||||
COALESCE(
|
||||
CAST(usage_settlement_snapshots.wallet_balance_before AS DOUBLE PRECISION),
|
||||
CAST(usage_record.wallet_balance_before AS DOUBLE PRECISION)
|
||||
) AS wallet_balance_before,
|
||||
COALESCE(
|
||||
CAST(usage_settlement_snapshots.wallet_balance_after AS DOUBLE PRECISION),
|
||||
CAST(usage_record.wallet_balance_after AS DOUBLE PRECISION)
|
||||
) AS wallet_balance_after,
|
||||
COALESCE(
|
||||
CAST(usage_settlement_snapshots.wallet_recharge_balance_before AS DOUBLE PRECISION),
|
||||
CAST(usage_record.wallet_recharge_balance_before AS DOUBLE PRECISION)
|
||||
) AS wallet_recharge_balance_before,
|
||||
COALESCE(
|
||||
CAST(usage_settlement_snapshots.wallet_recharge_balance_after AS DOUBLE PRECISION),
|
||||
CAST(usage_record.wallet_recharge_balance_after AS DOUBLE PRECISION)
|
||||
) AS wallet_recharge_balance_after,
|
||||
COALESCE(
|
||||
CAST(usage_settlement_snapshots.wallet_gift_balance_before AS DOUBLE PRECISION),
|
||||
CAST(usage_record.wallet_gift_balance_before AS DOUBLE PRECISION)
|
||||
) AS wallet_gift_balance_before,
|
||||
COALESCE(
|
||||
CAST(usage_settlement_snapshots.wallet_gift_balance_after AS DOUBLE PRECISION),
|
||||
CAST(usage_record.wallet_gift_balance_after AS DOUBLE PRECISION)
|
||||
) AS wallet_gift_balance_after,
|
||||
CAST(usage_settlement_snapshots.provider_monthly_used_usd AS DOUBLE PRECISION) AS provider_monthly_used_usd,
|
||||
usage_record.provider_id,
|
||||
CAST(
|
||||
EXTRACT(
|
||||
EPOCH FROM COALESCE(usage_settlement_snapshots.finalized_at, usage_record.finalized_at)
|
||||
) AS BIGINT
|
||||
) AS finalized_at_unix_secs
|
||||
FROM "usage" AS usage_record
|
||||
LEFT JOIN usage_settlement_snapshots
|
||||
ON usage_settlement_snapshots.request_id = usage_record.request_id
|
||||
WHERE usage_record.request_id = $1
|
||||
FOR UPDATE OF usage_record
|
||||
"#;
|
||||
|
||||
const FINALIZE_USAGE_BILLING_SQL: &str = r#"
|
||||
UPDATE "usage"
|
||||
SET
|
||||
@@ -14,6 +57,71 @@ SET
|
||||
WHERE request_id = $1
|
||||
"#;
|
||||
|
||||
const UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL: &str = r#"
|
||||
INSERT INTO usage_settlement_snapshots (
|
||||
request_id,
|
||||
billing_status,
|
||||
wallet_id,
|
||||
wallet_balance_before,
|
||||
wallet_balance_after,
|
||||
wallet_recharge_balance_before,
|
||||
wallet_recharge_balance_after,
|
||||
wallet_gift_balance_before,
|
||||
wallet_gift_balance_after,
|
||||
provider_monthly_used_usd,
|
||||
finalized_at
|
||||
) VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5,
|
||||
$6,
|
||||
$7,
|
||||
$8,
|
||||
$9,
|
||||
$10,
|
||||
CASE
|
||||
WHEN $11 IS NULL THEN NULL
|
||||
ELSE TO_TIMESTAMP($11::double precision)
|
||||
END
|
||||
)
|
||||
ON CONFLICT (request_id)
|
||||
DO UPDATE SET
|
||||
billing_status = EXCLUDED.billing_status,
|
||||
wallet_id = COALESCE(EXCLUDED.wallet_id, usage_settlement_snapshots.wallet_id),
|
||||
wallet_balance_before = COALESCE(
|
||||
EXCLUDED.wallet_balance_before,
|
||||
usage_settlement_snapshots.wallet_balance_before
|
||||
),
|
||||
wallet_balance_after = COALESCE(
|
||||
EXCLUDED.wallet_balance_after,
|
||||
usage_settlement_snapshots.wallet_balance_after
|
||||
),
|
||||
wallet_recharge_balance_before = COALESCE(
|
||||
EXCLUDED.wallet_recharge_balance_before,
|
||||
usage_settlement_snapshots.wallet_recharge_balance_before
|
||||
),
|
||||
wallet_recharge_balance_after = COALESCE(
|
||||
EXCLUDED.wallet_recharge_balance_after,
|
||||
usage_settlement_snapshots.wallet_recharge_balance_after
|
||||
),
|
||||
wallet_gift_balance_before = COALESCE(
|
||||
EXCLUDED.wallet_gift_balance_before,
|
||||
usage_settlement_snapshots.wallet_gift_balance_before
|
||||
),
|
||||
wallet_gift_balance_after = COALESCE(
|
||||
EXCLUDED.wallet_gift_balance_after,
|
||||
usage_settlement_snapshots.wallet_gift_balance_after
|
||||
),
|
||||
provider_monthly_used_usd = COALESCE(
|
||||
EXCLUDED.provider_monthly_used_usd,
|
||||
usage_settlement_snapshots.provider_monthly_used_usd
|
||||
),
|
||||
finalized_at = COALESCE(EXCLUDED.finalized_at, usage_settlement_snapshots.finalized_at),
|
||||
updated_at = NOW()
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxSettlementRepository {
|
||||
tx_runner: PostgresTransactionRunner,
|
||||
@@ -26,6 +134,62 @@ impl SqlxSettlementRepository {
|
||||
}
|
||||
}
|
||||
|
||||
fn settlement_from_row(
|
||||
row: &sqlx::postgres::PgRow,
|
||||
) -> Result<StoredUsageSettlement, DataLayerError> {
|
||||
Ok(StoredUsageSettlement {
|
||||
request_id: row.try_get("request_id").map_postgres_err()?,
|
||||
wallet_id: row.try_get("wallet_id").map_postgres_err()?,
|
||||
billing_status: row.try_get("billing_status").map_postgres_err()?,
|
||||
wallet_balance_before: row.try_get("wallet_balance_before").map_postgres_err()?,
|
||||
wallet_balance_after: row.try_get("wallet_balance_after").map_postgres_err()?,
|
||||
wallet_recharge_balance_before: row
|
||||
.try_get("wallet_recharge_balance_before")
|
||||
.map_postgres_err()?,
|
||||
wallet_recharge_balance_after: row
|
||||
.try_get("wallet_recharge_balance_after")
|
||||
.map_postgres_err()?,
|
||||
wallet_gift_balance_before: row
|
||||
.try_get("wallet_gift_balance_before")
|
||||
.map_postgres_err()?,
|
||||
wallet_gift_balance_after: row
|
||||
.try_get("wallet_gift_balance_after")
|
||||
.map_postgres_err()?,
|
||||
provider_monthly_used_usd: row
|
||||
.try_get("provider_monthly_used_usd")
|
||||
.map_postgres_err()?,
|
||||
finalized_at_unix_secs: row
|
||||
.try_get::<Option<i64>, _>("finalized_at_unix_secs")
|
||||
.map_postgres_err()?
|
||||
.map(|value| value as u64),
|
||||
})
|
||||
}
|
||||
|
||||
async fn sync_usage_settlement_snapshot<'e, E>(
|
||||
executor: E,
|
||||
settlement: &StoredUsageSettlement,
|
||||
) -> Result<(), DataLayerError>
|
||||
where
|
||||
E: sqlx::Executor<'e, Database = sqlx::Postgres>,
|
||||
{
|
||||
sqlx::query(UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL)
|
||||
.bind(&settlement.request_id)
|
||||
.bind(&settlement.billing_status)
|
||||
.bind(settlement.wallet_id.as_deref())
|
||||
.bind(settlement.wallet_balance_before)
|
||||
.bind(settlement.wallet_balance_after)
|
||||
.bind(settlement.wallet_recharge_balance_before)
|
||||
.bind(settlement.wallet_recharge_balance_after)
|
||||
.bind(settlement.wallet_gift_balance_before)
|
||||
.bind(settlement.wallet_gift_balance_after)
|
||||
.bind(settlement.provider_monthly_used_usd)
|
||||
.bind(settlement.finalized_at_unix_secs.map(|value| value as f64))
|
||||
.execute(executor)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SettlementWriteRepository for SqlxSettlementRepository {
|
||||
async fn settle_usage(
|
||||
@@ -36,29 +200,11 @@ impl SettlementWriteRepository for SqlxSettlementRepository {
|
||||
self.tx_runner
|
||||
.run_read_write(|tx| {
|
||||
Box::pin(async move {
|
||||
let row = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
request_id,
|
||||
wallet_id,
|
||||
billing_status,
|
||||
CAST(wallet_balance_before AS DOUBLE PRECISION) AS wallet_balance_before,
|
||||
CAST(wallet_balance_after AS DOUBLE PRECISION) AS wallet_balance_after,
|
||||
CAST(wallet_recharge_balance_before AS DOUBLE PRECISION) AS wallet_recharge_balance_before,
|
||||
CAST(wallet_recharge_balance_after AS DOUBLE PRECISION) AS wallet_recharge_balance_after,
|
||||
CAST(wallet_gift_balance_before AS DOUBLE PRECISION) AS wallet_gift_balance_before,
|
||||
CAST(wallet_gift_balance_after AS DOUBLE PRECISION) AS wallet_gift_balance_after,
|
||||
provider_id,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
FROM "usage"
|
||||
WHERE request_id = $1
|
||||
FOR UPDATE
|
||||
"#,
|
||||
)
|
||||
.bind(&input.request_id)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
let row = sqlx::query(FIND_USAGE_FOR_SETTLEMENT_SQL)
|
||||
.bind(&input.request_id)
|
||||
.fetch_optional(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
|
||||
let Some(usage_row) = row else {
|
||||
return Ok(None);
|
||||
@@ -67,34 +213,7 @@ FOR UPDATE
|
||||
let current_billing_status: String =
|
||||
usage_row.try_get("billing_status").map_postgres_err()?;
|
||||
if current_billing_status == "settled" || current_billing_status == "void" {
|
||||
return Ok(Some(StoredUsageSettlement {
|
||||
request_id: usage_row.try_get("request_id").map_postgres_err()?,
|
||||
wallet_id: usage_row.try_get("wallet_id").map_postgres_err()?,
|
||||
billing_status: current_billing_status,
|
||||
wallet_balance_before: usage_row
|
||||
.try_get("wallet_balance_before")
|
||||
.map_postgres_err()?,
|
||||
wallet_balance_after: usage_row
|
||||
.try_get("wallet_balance_after")
|
||||
.map_postgres_err()?,
|
||||
wallet_recharge_balance_before: usage_row
|
||||
.try_get("wallet_recharge_balance_before")
|
||||
.map_postgres_err()?,
|
||||
wallet_recharge_balance_after: usage_row
|
||||
.try_get("wallet_recharge_balance_after")
|
||||
.map_postgres_err()?,
|
||||
wallet_gift_balance_before: usage_row
|
||||
.try_get("wallet_gift_balance_before")
|
||||
.map_postgres_err()?,
|
||||
wallet_gift_balance_after: usage_row
|
||||
.try_get("wallet_gift_balance_after")
|
||||
.map_postgres_err()?,
|
||||
provider_monthly_used_usd: None,
|
||||
finalized_at_unix_secs: usage_row
|
||||
.try_get::<Option<i64>, _>("finalized_at_unix_secs")
|
||||
.map_postgres_err()?
|
||||
.map(|value| value as u64),
|
||||
}));
|
||||
return settlement_from_row(&usage_row).map(Some);
|
||||
}
|
||||
|
||||
let final_billing_status = if input.status == "completed" {
|
||||
@@ -223,32 +342,6 @@ WHERE id = $1
|
||||
settlement.wallet_recharge_balance_after = Some(after_recharge);
|
||||
settlement.wallet_gift_balance_before = Some(before_gift);
|
||||
settlement.wallet_gift_balance_after = Some(after_gift);
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
UPDATE "usage"
|
||||
SET
|
||||
wallet_id = $2,
|
||||
wallet_balance_before = $3,
|
||||
wallet_balance_after = $4,
|
||||
wallet_recharge_balance_before = $5,
|
||||
wallet_recharge_balance_after = $6,
|
||||
wallet_gift_balance_before = $7,
|
||||
wallet_gift_balance_after = $8
|
||||
WHERE request_id = $1
|
||||
"#,
|
||||
)
|
||||
.bind(&input.request_id)
|
||||
.bind(&wallet_id)
|
||||
.bind(before_total)
|
||||
.bind(after_recharge + after_gift)
|
||||
.bind(before_recharge)
|
||||
.bind(after_recharge)
|
||||
.bind(before_gift)
|
||||
.bind(after_gift)
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
}
|
||||
|
||||
if let Some(provider_id) = input
|
||||
@@ -283,6 +376,7 @@ RETURNING CAST(monthly_used_usd AS DOUBLE PRECISION) AS monthly_used_usd
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
sync_usage_settlement_snapshot(&mut **tx, &settlement).await?;
|
||||
|
||||
Ok(Some(settlement))
|
||||
})
|
||||
@@ -297,4 +391,28 @@ mod tests {
|
||||
fn finalize_usage_billing_sql_does_not_require_usage_updated_at_column() {
|
||||
assert!(!super::FINALIZE_USAGE_BILLING_SQL.contains("updated_at"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settlement_sql_reads_settlement_snapshots_before_legacy_usage_columns() {
|
||||
assert!(
|
||||
super::FIND_USAGE_FOR_SETTLEMENT_SQL.contains("LEFT JOIN usage_settlement_snapshots")
|
||||
);
|
||||
assert!(super::FIND_USAGE_FOR_SETTLEMENT_SQL.contains("COALESCE("));
|
||||
assert!(super::FIND_USAGE_FOR_SETTLEMENT_SQL.contains("FOR UPDATE OF usage_record"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settlement_sql_dual_writes_usage_settlement_snapshots() {
|
||||
assert!(super::UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL
|
||||
.contains("INSERT INTO usage_settlement_snapshots"));
|
||||
assert!(super::UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL.contains("provider_monthly_used_usd"));
|
||||
assert!(super::UPSERT_USAGE_SETTLEMENT_SNAPSHOT_SQL
|
||||
.contains("TO_TIMESTAMP($11::double precision)"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn settlement_sql_no_longer_dual_writes_wallet_snapshots_to_usage_rows() {
|
||||
let source = include_str!("sql.rs");
|
||||
assert!(!source.contains("UPDATE \"usage\"\nSET\n wallet_id = $2"));
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -9,3 +9,92 @@ pub(crate) use aether_data_contracts::repository::usage::{
|
||||
};
|
||||
pub use memory::InMemoryUsageReadRepository;
|
||||
pub use sql::SqlxUsageReadRepository;
|
||||
|
||||
pub(crate) fn strip_deprecated_usage_display_fields(
|
||||
mut usage: UpsertUsageRecord,
|
||||
) -> UpsertUsageRecord {
|
||||
usage.username = None;
|
||||
usage.api_key_name = None;
|
||||
usage
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{strip_deprecated_usage_display_fields, UpsertUsageRecord};
|
||||
|
||||
#[test]
|
||||
fn strip_deprecated_usage_display_fields_clears_legacy_display_columns() {
|
||||
let usage = strip_deprecated_usage_display_fields(UpsertUsageRecord {
|
||||
request_id: "req-1".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("key-1".to_string()),
|
||||
username: Some("alice".to_string()),
|
||||
api_key_name: Some("default".to_string()),
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
target_model: None,
|
||||
provider_id: None,
|
||||
provider_endpoint_id: None,
|
||||
provider_api_key_id: None,
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
provider_api_family: Some("openai".to_string()),
|
||||
provider_endpoint_kind: Some("chat".to_string()),
|
||||
has_format_conversion: Some(false),
|
||||
is_stream: Some(false),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(20),
|
||||
total_tokens: Some(30),
|
||||
cache_creation_input_tokens: None,
|
||||
cache_creation_ephemeral_5m_input_tokens: None,
|
||||
cache_creation_ephemeral_1h_input_tokens: None,
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_cost_usd: None,
|
||||
cache_read_cost_usd: None,
|
||||
output_price_per_1m: None,
|
||||
total_cost_usd: Some(0.25),
|
||||
actual_total_cost_usd: Some(0.15),
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
error_category: None,
|
||||
response_time_ms: Some(120),
|
||||
first_byte_time_ms: Some(40),
|
||||
status: "completed".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
request_headers: None,
|
||||
request_body: None,
|
||||
request_body_ref: None,
|
||||
provider_request_headers: None,
|
||||
provider_request_body: None,
|
||||
provider_request_body_ref: None,
|
||||
response_headers: None,
|
||||
response_body: None,
|
||||
response_body_ref: None,
|
||||
client_response_headers: None,
|
||||
client_response_body: None,
|
||||
client_response_body_ref: None,
|
||||
candidate_id: None,
|
||||
candidate_index: None,
|
||||
key_name: None,
|
||||
planner_kind: None,
|
||||
route_family: None,
|
||||
route_kind: None,
|
||||
execution_path: None,
|
||||
local_execution_runtime_miss_reason: None,
|
||||
request_metadata: None,
|
||||
finalized_at_unix_secs: None,
|
||||
created_at_unix_ms: Some(100),
|
||||
updated_at_unix_secs: 101,
|
||||
});
|
||||
|
||||
assert_eq!(usage.user_id.as_deref(), Some("user-1"));
|
||||
assert_eq!(usage.api_key_id.as_deref(), Some("key-1"));
|
||||
assert_eq!(usage.username, None);
|
||||
assert_eq!(usage.api_key_name, None);
|
||||
assert_eq!(usage.provider_name, "OpenAI");
|
||||
assert_eq!(usage.model, "gpt-5");
|
||||
}
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -242,7 +242,7 @@ fn merge_query_string(
|
||||
mod tests {
|
||||
use super::{
|
||||
build_gemini_content_url, build_gemini_files_passthrough_url,
|
||||
build_gemini_video_predict_long_running_url, build_openai_chat_url,
|
||||
build_gemini_video_predict_long_running_url, build_openai_chat_url, build_openai_cli_url,
|
||||
build_passthrough_path_url,
|
||||
};
|
||||
|
||||
@@ -257,6 +257,18 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn openai_cli_url_preserves_codex_path_prefix() {
|
||||
assert_eq!(
|
||||
build_openai_cli_url("https://tiger.bookapi.cc/codex", None, false),
|
||||
"https://tiger.bookapi.cc/codex/responses"
|
||||
);
|
||||
assert_eq!(
|
||||
build_openai_cli_url("https://tiger.bookapi.cc/codex?tenant=demo", None, true),
|
||||
"https://tiger.bookapi.cc/codex/responses/compact?tenant=demo"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn merges_base_url_query_for_dynamic_gemini_content_urls() {
|
||||
assert_eq!(
|
||||
|
||||
@@ -94,18 +94,42 @@ pub struct UsageEventData {
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub request_body: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub request_body_ref: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_request_headers: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_request_body: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub provider_request_body_ref: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub response_headers: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub response_body: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub response_body_ref: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub client_response_headers: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub client_response_body: Option<Value>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub client_response_body_ref: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub candidate_id: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub candidate_index: Option<u64>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub key_name: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub planner_kind: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub route_family: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub route_kind: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub execution_path: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub local_execution_runtime_miss_reason: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub request_metadata: Option<Value>,
|
||||
}
|
||||
|
||||
|
||||
48
crates/aether-usage-runtime/src/executor.rs
Normal file
48
crates/aether-usage-runtime/src/executor.rs
Normal file
@@ -0,0 +1,48 @@
|
||||
use std::future::Future;
|
||||
use std::sync::OnceLock;
|
||||
|
||||
const USAGE_BACKGROUND_RUNTIME_THREADS: usize = 2;
|
||||
const USAGE_BACKGROUND_RUNTIME_STACK_BYTES: usize = 8 * 1024 * 1024;
|
||||
const USAGE_BACKGROUND_RUNTIME_THREAD_NAME: &str = "aether-usage-runtime";
|
||||
|
||||
pub(crate) fn spawn_on_usage_background_runtime<F>(task: F) -> tokio::task::JoinHandle<F::Output>
|
||||
where
|
||||
F: Future + Send + 'static,
|
||||
F::Output: Send + 'static,
|
||||
{
|
||||
usage_background_runtime().handle().spawn(task)
|
||||
}
|
||||
|
||||
fn usage_background_runtime() -> &'static tokio::runtime::Runtime {
|
||||
static RUNTIME: OnceLock<&'static tokio::runtime::Runtime> = OnceLock::new();
|
||||
|
||||
RUNTIME.get_or_init(|| {
|
||||
let runtime = tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.worker_threads(USAGE_BACKGROUND_RUNTIME_THREADS)
|
||||
.thread_name(USAGE_BACKGROUND_RUNTIME_THREAD_NAME)
|
||||
.thread_stack_size(USAGE_BACKGROUND_RUNTIME_STACK_BYTES)
|
||||
.build()
|
||||
.expect("usage background runtime should build");
|
||||
Box::leak(Box::new(runtime))
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::spawn_on_usage_background_runtime;
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_background_runtime_runs_on_dedicated_named_threads() {
|
||||
let thread_name = spawn_on_usage_background_runtime(async move {
|
||||
std::thread::current()
|
||||
.name()
|
||||
.unwrap_or_default()
|
||||
.to_string()
|
||||
})
|
||||
.await
|
||||
.expect("background task should complete");
|
||||
|
||||
assert_eq!(thread_name, "aether-usage-runtime");
|
||||
}
|
||||
}
|
||||
@@ -1,5 +1,6 @@
|
||||
pub mod config;
|
||||
pub mod event;
|
||||
mod executor;
|
||||
pub mod queue;
|
||||
pub mod record;
|
||||
pub mod report;
|
||||
@@ -38,9 +39,14 @@ pub use worker::{
|
||||
UsageQueueWorker, UsageRecordWriter,
|
||||
};
|
||||
pub use write::{
|
||||
build_pending_usage_record, build_stream_terminal_usage_event,
|
||||
build_stream_terminal_usage_outcome, build_streaming_usage_record,
|
||||
build_lifecycle_usage_seed, build_pending_usage_record, build_pending_usage_record_from_seed,
|
||||
build_stream_terminal_usage_event, build_stream_terminal_usage_outcome,
|
||||
build_stream_terminal_usage_payload_seed, build_stream_terminal_usage_seed,
|
||||
build_streaming_usage_record, build_streaming_usage_record_from_seed,
|
||||
build_sync_terminal_usage_event, build_sync_terminal_usage_outcome,
|
||||
build_terminal_usage_event_from_outcome, build_usage_event_data_seed, TerminalUsageOutcome,
|
||||
UsageTerminalState,
|
||||
build_sync_terminal_usage_payload_seed, build_sync_terminal_usage_seed,
|
||||
build_terminal_usage_context_seed, build_terminal_usage_event_from_outcome,
|
||||
build_terminal_usage_event_from_seed, build_usage_event_data_seed, LifecycleUsageSeed,
|
||||
StreamTerminalUsagePayloadSeed, SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed,
|
||||
TerminalUsageOutcome, TerminalUsageSeed, UsageTerminalState,
|
||||
};
|
||||
|
||||
@@ -4,6 +4,27 @@ use aether_data_contracts::DataLayerError;
|
||||
use crate::request_metadata::sanitize_usage_request_metadata;
|
||||
use crate::{UsageEvent, UsageEventType};
|
||||
|
||||
fn metadata_string(metadata: Option<&serde_json::Value>, key: &str) -> Option<String> {
|
||||
metadata
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|object| object.get(key))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned)
|
||||
}
|
||||
|
||||
fn metadata_u64(metadata: Option<&serde_json::Value>, key: &str) -> Option<u64> {
|
||||
metadata
|
||||
.and_then(serde_json::Value::as_object)
|
||||
.and_then(|object| object.get(key))
|
||||
.and_then(|value| {
|
||||
value
|
||||
.as_u64()
|
||||
.or_else(|| value.as_i64().and_then(|number| u64::try_from(number).ok()))
|
||||
})
|
||||
}
|
||||
|
||||
pub fn build_upsert_usage_record_from_event(
|
||||
event: &UsageEvent,
|
||||
) -> Result<UpsertUsageRecord, DataLayerError> {
|
||||
@@ -53,12 +74,51 @@ pub fn build_upsert_usage_record_from_event(
|
||||
billing_status: billing_status.to_string(),
|
||||
request_headers: data.request_headers,
|
||||
request_body: data.request_body,
|
||||
request_body_ref: empty_to_none(data.request_body_ref)
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "request_body_ref")),
|
||||
provider_request_headers: data.provider_request_headers,
|
||||
provider_request_body: data.provider_request_body,
|
||||
provider_request_body_ref: empty_to_none(data.provider_request_body_ref).or_else(|| {
|
||||
metadata_string(data.request_metadata.as_ref(), "provider_request_body_ref")
|
||||
}),
|
||||
response_headers: data.response_headers,
|
||||
response_body: data.response_body,
|
||||
response_body_ref: empty_to_none(data.response_body_ref)
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "response_body_ref")),
|
||||
client_response_headers: data.client_response_headers,
|
||||
client_response_body: data.client_response_body,
|
||||
client_response_body_ref: empty_to_none(data.client_response_body_ref).or_else(|| {
|
||||
metadata_string(data.request_metadata.as_ref(), "client_response_body_ref")
|
||||
}),
|
||||
candidate_id: data
|
||||
.candidate_id
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "candidate_id")),
|
||||
candidate_index: data
|
||||
.candidate_index
|
||||
.or_else(|| metadata_u64(data.request_metadata.as_ref(), "candidate_index")),
|
||||
key_name: data
|
||||
.key_name
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "key_name")),
|
||||
planner_kind: data
|
||||
.planner_kind
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "planner_kind")),
|
||||
route_family: data
|
||||
.route_family
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "route_family")),
|
||||
route_kind: data
|
||||
.route_kind
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "route_kind")),
|
||||
execution_path: data
|
||||
.execution_path
|
||||
.or_else(|| metadata_string(data.request_metadata.as_ref(), "execution_path")),
|
||||
local_execution_runtime_miss_reason: data.local_execution_runtime_miss_reason.or_else(
|
||||
|| {
|
||||
metadata_string(
|
||||
data.request_metadata.as_ref(),
|
||||
"local_execution_runtime_miss_reason",
|
||||
)
|
||||
},
|
||||
),
|
||||
request_metadata: sanitize_usage_request_metadata(data.request_metadata),
|
||||
finalized_at_unix_secs: Some(now_unix_secs),
|
||||
created_at_unix_ms: Some(now_unix_secs),
|
||||
@@ -138,11 +198,11 @@ mod tests {
|
||||
})
|
||||
.expect("record should build");
|
||||
|
||||
assert_eq!(record.candidate_id.as_deref(), Some("cand-2"));
|
||||
assert_eq!(record.key_name.as_deref(), Some("upstream-primary"));
|
||||
assert_eq!(
|
||||
record.request_metadata,
|
||||
Some(serde_json::json!({
|
||||
"candidate_id": "cand-2",
|
||||
"key_name": "upstream-primary",
|
||||
"billing_snapshot": { "status": "complete" }
|
||||
}))
|
||||
);
|
||||
|
||||
@@ -1,43 +1,34 @@
|
||||
use aether_contracts::ExecutionPlan;
|
||||
use serde_json::{Map, Value};
|
||||
use serde_json::{json, Map, Value};
|
||||
|
||||
const MAX_USAGE_REQUEST_METADATA_DEPTH: usize = 32;
|
||||
const MAX_USAGE_REQUEST_METADATA_NODES: usize = 4_000;
|
||||
const MAX_USAGE_REQUEST_METADATA_BYTES: usize = 16 * 1024;
|
||||
const MAX_USAGE_REQUEST_METADATA_STRING_BYTES: usize = 1_024;
|
||||
|
||||
pub(crate) fn build_usage_request_metadata_seed(
|
||||
plan: &ExecutionPlan,
|
||||
_plan: &ExecutionPlan,
|
||||
context: Option<&Map<String, Value>>,
|
||||
) -> Option<Value> {
|
||||
let mut metadata = context.cloned().unwrap_or_default();
|
||||
if !has_non_empty_string(&metadata, "candidate_id") {
|
||||
if let Some(candidate_id) = plan
|
||||
.candidate_id
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
metadata.insert(
|
||||
"candidate_id".to_string(),
|
||||
Value::String(candidate_id.to_string()),
|
||||
);
|
||||
}
|
||||
let mut metadata = Map::new();
|
||||
if let Some(context) = context {
|
||||
copy_allowed_metadata_fields(context, &mut metadata);
|
||||
}
|
||||
sanitize_usage_request_metadata(Some(Value::Object(metadata)))
|
||||
(!metadata.is_empty()).then_some(Value::Object(metadata))
|
||||
}
|
||||
|
||||
pub(crate) fn merge_usage_request_metadata(
|
||||
base: Option<Value>,
|
||||
override_value: Option<Value>,
|
||||
) -> Option<Value> {
|
||||
let merged = match (base, override_value) {
|
||||
(Some(Value::Object(mut base)), Some(Value::Object(override_object))) => {
|
||||
for (key, value) in override_object {
|
||||
base.insert(key, value);
|
||||
}
|
||||
Some(Value::Object(base))
|
||||
}
|
||||
(Some(base), None) => Some(base),
|
||||
(_, Some(override_value)) => Some(override_value),
|
||||
(None, None) => None,
|
||||
};
|
||||
sanitize_usage_request_metadata(merged)
|
||||
let mut metadata = Map::new();
|
||||
if let Some(Value::Object(base)) = base.as_ref() {
|
||||
copy_allowed_metadata_fields(base, &mut metadata);
|
||||
}
|
||||
if let Some(Value::Object(override_object)) = override_value.as_ref() {
|
||||
copy_allowed_metadata_fields(override_object, &mut metadata);
|
||||
}
|
||||
(!metadata.is_empty()).then_some(Value::Object(metadata))
|
||||
}
|
||||
|
||||
pub(crate) fn sanitize_usage_request_metadata(value: Option<Value>) -> Option<Value> {
|
||||
@@ -46,26 +37,29 @@ pub(crate) fn sanitize_usage_request_metadata(value: Option<Value>) -> Option<Va
|
||||
};
|
||||
|
||||
let mut filtered = Map::new();
|
||||
copy_non_empty_string(&object, &mut filtered, "candidate_id");
|
||||
copy_number(&object, &mut filtered, "candidate_index");
|
||||
copy_non_empty_string(&object, &mut filtered, "key_name");
|
||||
copy_non_empty_string(&object, &mut filtered, "trace_id");
|
||||
copy_non_null_value(&object, &mut filtered, "billing_snapshot");
|
||||
copy_non_null_value(&object, &mut filtered, "dimensions");
|
||||
copy_non_null_value(&object, &mut filtered, "billing_rule_snapshot");
|
||||
copy_non_null_value(&object, &mut filtered, "scheduling_audit");
|
||||
copy_number(&object, &mut filtered, "rate_multiplier");
|
||||
copy_bool(&object, &mut filtered, "is_free_tier");
|
||||
copy_allowed_metadata_fields(&object, &mut filtered);
|
||||
|
||||
(!filtered.is_empty()).then_some(Value::Object(filtered))
|
||||
}
|
||||
|
||||
fn has_non_empty_string(object: &Map<String, Value>, key: &str) -> bool {
|
||||
object
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.is_some_and(|value| !value.is_empty())
|
||||
fn copy_allowed_metadata_fields(source: &Map<String, Value>, target: &mut Map<String, Value>) {
|
||||
copy_non_empty_string(source, target, "trace_id");
|
||||
copy_number(source, target, "provider_request_body_base64_bytes");
|
||||
copy_number(source, target, "provider_response_body_base64_bytes");
|
||||
copy_number(source, target, "client_response_body_base64_bytes");
|
||||
copy_non_null_value(source, target, "billing_snapshot");
|
||||
copy_non_empty_string(source, target, "billing_snapshot_schema_version");
|
||||
copy_non_empty_string(source, target, "billing_snapshot_status");
|
||||
copy_non_null_value(source, target, "dimensions");
|
||||
copy_non_null_value(source, target, "billing_rule_snapshot");
|
||||
copy_non_null_value(source, target, "scheduling_audit");
|
||||
copy_number(source, target, "rate_multiplier");
|
||||
copy_bool(source, target, "is_free_tier");
|
||||
copy_number(source, target, "input_price_per_1m");
|
||||
copy_number(source, target, "output_price_per_1m");
|
||||
copy_number(source, target, "cache_creation_price_per_1m");
|
||||
copy_number(source, target, "cache_read_price_per_1m");
|
||||
copy_number(source, target, "price_per_request");
|
||||
}
|
||||
|
||||
fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
|
||||
@@ -77,7 +71,10 @@ fn copy_non_empty_string(source: &Map<String, Value>, target: &mut Map<String, V
|
||||
else {
|
||||
return;
|
||||
};
|
||||
target.insert(key.to_string(), Value::String(value.to_string()));
|
||||
target.insert(
|
||||
key.to_string(),
|
||||
Value::String(truncate_usage_request_metadata_string(value)),
|
||||
);
|
||||
}
|
||||
|
||||
fn copy_number(source: &Map<String, Value>, target: &mut Map<String, Value>, key: &str) {
|
||||
@@ -98,18 +95,130 @@ fn copy_non_null_value(source: &Map<String, Value>, target: &mut Map<String, Val
|
||||
let Some(value) = source.get(key).filter(|value| !value.is_null()) else {
|
||||
return;
|
||||
};
|
||||
target.insert(key.to_string(), value.clone());
|
||||
target.insert(
|
||||
key.to_string(),
|
||||
sanitize_usage_request_metadata_value(value),
|
||||
);
|
||||
}
|
||||
|
||||
fn sanitize_usage_request_metadata_value(value: &Value) -> Value {
|
||||
match value {
|
||||
Value::String(text) => Value::String(truncate_usage_request_metadata_string(text)),
|
||||
_ if usage_request_metadata_within_limits(value) => value.clone(),
|
||||
_ => truncated_usage_request_metadata_value(value),
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate_usage_request_metadata_string(value: &str) -> String {
|
||||
const TRUNCATED_SUFFIX: &str = "...[truncated]";
|
||||
|
||||
if value.len() <= MAX_USAGE_REQUEST_METADATA_STRING_BYTES {
|
||||
return value.to_string();
|
||||
}
|
||||
|
||||
let target_bytes =
|
||||
MAX_USAGE_REQUEST_METADATA_STRING_BYTES.saturating_sub(TRUNCATED_SUFFIX.len());
|
||||
let mut end = 0usize;
|
||||
for (idx, ch) in value.char_indices() {
|
||||
let next = idx + ch.len_utf8();
|
||||
if next > target_bytes {
|
||||
break;
|
||||
}
|
||||
end = next;
|
||||
}
|
||||
|
||||
if end == 0 {
|
||||
return TRUNCATED_SUFFIX.to_string();
|
||||
}
|
||||
|
||||
format!("{}{TRUNCATED_SUFFIX}", &value[..end])
|
||||
}
|
||||
|
||||
fn truncated_usage_request_metadata_value(value: &Value) -> Value {
|
||||
json!({
|
||||
"truncated": true,
|
||||
"reason": "usage_request_metadata_limits_exceeded",
|
||||
"max_depth": MAX_USAGE_REQUEST_METADATA_DEPTH,
|
||||
"max_nodes": MAX_USAGE_REQUEST_METADATA_NODES,
|
||||
"max_bytes": MAX_USAGE_REQUEST_METADATA_BYTES,
|
||||
"value_kind": usage_request_metadata_value_kind(value),
|
||||
})
|
||||
}
|
||||
|
||||
fn usage_request_metadata_within_limits(value: &Value) -> bool {
|
||||
let mut nodes = 0usize;
|
||||
let mut estimated_bytes = 0usize;
|
||||
let mut stack = vec![(value, 1usize)];
|
||||
|
||||
while let Some((current, depth)) = stack.pop() {
|
||||
nodes = nodes.saturating_add(1);
|
||||
estimated_bytes =
|
||||
estimated_bytes.saturating_add(usage_request_metadata_value_size_hint(current));
|
||||
if depth > MAX_USAGE_REQUEST_METADATA_DEPTH
|
||||
|| nodes > MAX_USAGE_REQUEST_METADATA_NODES
|
||||
|| estimated_bytes > MAX_USAGE_REQUEST_METADATA_BYTES
|
||||
{
|
||||
return false;
|
||||
}
|
||||
match current {
|
||||
Value::Array(items) => {
|
||||
estimated_bytes = estimated_bytes.saturating_add(items.len().saturating_mul(2));
|
||||
for item in items.iter().rev() {
|
||||
stack.push((item, depth + 1));
|
||||
}
|
||||
}
|
||||
Value::Object(object) => {
|
||||
estimated_bytes = estimated_bytes
|
||||
.saturating_add(object.len().saturating_mul(3))
|
||||
.saturating_add(
|
||||
object
|
||||
.keys()
|
||||
.map(|key| key.len().saturating_add(2))
|
||||
.sum::<usize>(),
|
||||
);
|
||||
for item in object.values() {
|
||||
stack.push((item, depth + 1));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
|
||||
fn usage_request_metadata_value_kind(value: &Value) -> &'static str {
|
||||
match value {
|
||||
Value::Null => "null",
|
||||
Value::Bool(_) => "bool",
|
||||
Value::Number(_) => "number",
|
||||
Value::String(_) => "string",
|
||||
Value::Array(_) => "array",
|
||||
Value::Object(_) => "object",
|
||||
}
|
||||
}
|
||||
|
||||
fn usage_request_metadata_value_size_hint(value: &Value) -> usize {
|
||||
match value {
|
||||
Value::Null => 4,
|
||||
Value::Bool(false) => 5,
|
||||
Value::Bool(true) => 4,
|
||||
Value::Number(number) => number.to_string().len(),
|
||||
Value::String(text) => text.len().saturating_add(2),
|
||||
Value::Array(_) | Value::Object(_) => 2,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use aether_contracts::{ExecutionPlan, RequestBody};
|
||||
use serde_json::json;
|
||||
use serde_json::{json, Value};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use super::{
|
||||
build_usage_request_metadata_seed, merge_usage_request_metadata,
|
||||
sanitize_usage_request_metadata,
|
||||
sanitize_usage_request_metadata, MAX_USAGE_REQUEST_METADATA_BYTES,
|
||||
MAX_USAGE_REQUEST_METADATA_DEPTH, MAX_USAGE_REQUEST_METADATA_NODES,
|
||||
};
|
||||
|
||||
fn sample_plan() -> ExecutionPlan {
|
||||
@@ -143,14 +252,22 @@ mod tests {
|
||||
"provider_id": "provider-1",
|
||||
"provider_name": "OpenAI",
|
||||
"model": "gpt-5",
|
||||
"candidate_id": "cand-1",
|
||||
"candidate_index": 2,
|
||||
"key_name": "upstream-primary",
|
||||
"trace_id": "trace-1",
|
||||
"provider_request_body_base64_bytes": 512,
|
||||
"provider_response_body_base64_bytes": 1024,
|
||||
"client_response_body_base64_bytes": 2048,
|
||||
"billing_snapshot": {"status": "complete"},
|
||||
"billing_snapshot_schema_version": "2.0",
|
||||
"billing_snapshot_status": "complete",
|
||||
"dimensions": {"total_input_context": 10},
|
||||
"rate_multiplier": 1.25,
|
||||
"is_free_tier": false,
|
||||
"input_price_per_1m": 3.0,
|
||||
"output_price_per_1m": 15.0,
|
||||
"cache_creation_price_per_1m": 3.75,
|
||||
"cache_read_price_per_1m": 0.3,
|
||||
"price_per_request": 0.02,
|
||||
"original_headers": {"authorization": "Bearer secret"},
|
||||
"original_request_body": {"messages": []},
|
||||
"provider_request_headers": {"authorization": "Bearer secret"},
|
||||
@@ -161,27 +278,60 @@ mod tests {
|
||||
assert_eq!(
|
||||
metadata,
|
||||
json!({
|
||||
"candidate_id": "cand-1",
|
||||
"candidate_index": 2,
|
||||
"key_name": "upstream-primary",
|
||||
"trace_id": "trace-1",
|
||||
"provider_request_body_base64_bytes": 512,
|
||||
"provider_response_body_base64_bytes": 1024,
|
||||
"client_response_body_base64_bytes": 2048,
|
||||
"billing_snapshot": {"status": "complete"},
|
||||
"billing_snapshot_schema_version": "2.0",
|
||||
"billing_snapshot_status": "complete",
|
||||
"dimensions": {"total_input_context": 10},
|
||||
"rate_multiplier": 1.25,
|
||||
"is_free_tier": false
|
||||
"is_free_tier": false,
|
||||
"input_price_per_1m": 3.0,
|
||||
"output_price_per_1m": 15.0,
|
||||
"cache_creation_price_per_1m": 3.75,
|
||||
"cache_read_price_per_1m": 0.3,
|
||||
"price_per_request": 0.02
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_seed_from_context_and_plan_candidate_id() {
|
||||
fn sanitizes_large_allowed_metadata_values_to_bounded_representations() {
|
||||
let metadata = sanitize_usage_request_metadata(Some(json!({
|
||||
"trace_id": "t".repeat(2_048),
|
||||
"billing_snapshot": {
|
||||
"payload": "x".repeat(32 * 1024)
|
||||
}
|
||||
})))
|
||||
.expect("metadata should remain");
|
||||
|
||||
assert!(metadata
|
||||
.get("trace_id")
|
||||
.and_then(Value::as_str)
|
||||
.is_some_and(|value| value.ends_with("...[truncated]")));
|
||||
assert_eq!(
|
||||
metadata.get("billing_snapshot"),
|
||||
Some(&json!({
|
||||
"truncated": true,
|
||||
"reason": "usage_request_metadata_limits_exceeded",
|
||||
"max_depth": MAX_USAGE_REQUEST_METADATA_DEPTH,
|
||||
"max_nodes": MAX_USAGE_REQUEST_METADATA_NODES,
|
||||
"max_bytes": MAX_USAGE_REQUEST_METADATA_BYTES,
|
||||
"value_kind": "object",
|
||||
}))
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn builds_seed_from_context_and_allowlisted_metadata_only() {
|
||||
let metadata = build_usage_request_metadata_seed(
|
||||
&sample_plan(),
|
||||
Some(
|
||||
json!({
|
||||
"request_id": "req-1",
|
||||
"candidate_index": 0,
|
||||
"key_name": "upstream-primary",
|
||||
"provider_id": "provider-1",
|
||||
"billing_snapshot": {"status": "complete"}
|
||||
})
|
||||
@@ -194,9 +344,6 @@ mod tests {
|
||||
assert_eq!(
|
||||
metadata,
|
||||
json!({
|
||||
"candidate_id": "cand-1",
|
||||
"candidate_index": 0,
|
||||
"key_name": "upstream-primary",
|
||||
"billing_snapshot": {"status": "complete"}
|
||||
})
|
||||
);
|
||||
@@ -206,24 +353,14 @@ mod tests {
|
||||
fn merges_and_filters_request_metadata() {
|
||||
let metadata = merge_usage_request_metadata(
|
||||
Some(json!({
|
||||
"candidate_id": "cand-1",
|
||||
"request_id": "req-1"
|
||||
})),
|
||||
Some(json!({
|
||||
"candidate_index": 0,
|
||||
"key_name": "upstream-primary",
|
||||
"provider_name": "OpenAI"
|
||||
})),
|
||||
)
|
||||
.expect("metadata should remain");
|
||||
|
||||
assert_eq!(
|
||||
metadata,
|
||||
json!({
|
||||
"candidate_id": "cand-1",
|
||||
"candidate_index": 0,
|
||||
"key_name": "upstream-primary"
|
||||
})
|
||||
);
|
||||
|
||||
assert_eq!(metadata, None);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,17 +1,22 @@
|
||||
use std::future::Future;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
|
||||
use aether_contracts::{ExecutionPlan, ExecutionTelemetry};
|
||||
use aether_contracts::ExecutionTelemetry;
|
||||
use aether_data::redis::RedisStreamRunner;
|
||||
use aether_data_contracts::repository::usage::UpsertUsageRecord;
|
||||
use aether_data_contracts::DataLayerError;
|
||||
use async_trait::async_trait;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::executor::spawn_on_usage_background_runtime;
|
||||
use crate::{
|
||||
build_pending_usage_record, build_stream_terminal_usage_outcome, build_streaming_usage_record,
|
||||
build_sync_terminal_usage_outcome, build_terminal_usage_event_from_outcome,
|
||||
build_upsert_usage_record_from_event, build_usage_queue_worker, settle_usage_if_needed,
|
||||
GatewayStreamReportRequest, GatewaySyncReportRequest, UsageEvent, UsageQueue,
|
||||
UsageRecordWriter, UsageRuntimeConfig, UsageSettlementWriter, UsageTerminalState,
|
||||
build_pending_usage_record_from_seed, build_stream_terminal_usage_seed,
|
||||
build_streaming_usage_record_from_seed, build_sync_terminal_usage_seed,
|
||||
build_terminal_usage_event_from_seed, build_upsert_usage_record_from_event,
|
||||
build_usage_queue_worker, settle_usage_if_needed, LifecycleUsageSeed,
|
||||
StreamTerminalUsagePayloadSeed, SyncTerminalUsagePayloadSeed, TerminalUsageContextSeed,
|
||||
UsageEvent, UsageQueue, UsageRecordWriter, UsageRuntimeConfig, UsageSettlementWriter,
|
||||
};
|
||||
|
||||
#[async_trait]
|
||||
@@ -32,6 +37,17 @@ pub struct UsageRuntime {
|
||||
config: UsageRuntimeConfig,
|
||||
}
|
||||
|
||||
struct SyncTerminalUsageTaskInput {
|
||||
context_seed: TerminalUsageContextSeed,
|
||||
payload_seed: SyncTerminalUsagePayloadSeed,
|
||||
}
|
||||
|
||||
struct StreamTerminalUsageTaskInput {
|
||||
context_seed: TerminalUsageContextSeed,
|
||||
payload_seed: StreamTerminalUsagePayloadSeed,
|
||||
cancelled: bool,
|
||||
}
|
||||
|
||||
impl Default for UsageRuntime {
|
||||
fn default() -> Self {
|
||||
Self::disabled()
|
||||
@@ -73,169 +89,197 @@ impl UsageRuntime {
|
||||
Some(worker.spawn())
|
||||
}
|
||||
|
||||
pub async fn record_pending<T>(
|
||||
&self,
|
||||
data: &T,
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
) where
|
||||
T: UsageRuntimeAccess,
|
||||
pub fn record_pending<T>(&self, data: &T, seed: &LifecycleUsageSeed)
|
||||
where
|
||||
T: UsageRuntimeAccess + Clone + 'static,
|
||||
{
|
||||
if !self.is_enabled() {
|
||||
return;
|
||||
}
|
||||
let now_unix_secs = now_unix_secs();
|
||||
match build_pending_usage_record(plan, report_context, now_unix_secs) {
|
||||
Ok(record) => {
|
||||
if let Err(err) = data.upsert_usage_record(record).await {
|
||||
let data = T::clone(data);
|
||||
let seed = seed.clone();
|
||||
let request_id = seed.request_id.clone();
|
||||
spawn_on_usage_background_runtime(boxed_usage_task(async move {
|
||||
let now_unix_secs = now_unix_secs();
|
||||
match build_pending_usage_record_offthread(&seed, now_unix_secs).await {
|
||||
Ok(record) => {
|
||||
if let Err(err) = data.upsert_usage_record(record).await {
|
||||
warn!(
|
||||
event_name = "usage_pending_record_failed",
|
||||
log_type = "event",
|
||||
request_id = %request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to record sync pending usage"
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "usage_pending_record_failed",
|
||||
event_name = "usage_pending_build_failed",
|
||||
log_type = "event",
|
||||
request_id = %plan.request_id,
|
||||
request_id = %request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to record sync pending usage"
|
||||
);
|
||||
"usage runtime failed to build sync pending usage"
|
||||
)
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "usage_pending_build_failed",
|
||||
log_type = "event",
|
||||
request_id = %plan.request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to build sync pending usage"
|
||||
)
|
||||
}
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
pub async fn record_stream_started<T>(
|
||||
pub fn record_stream_started<T>(
|
||||
&self,
|
||||
data: &T,
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
seed: &LifecycleUsageSeed,
|
||||
status_code: u16,
|
||||
headers: &std::collections::BTreeMap<String, String>,
|
||||
telemetry: Option<&ExecutionTelemetry>,
|
||||
) where
|
||||
T: UsageRuntimeAccess,
|
||||
T: UsageRuntimeAccess + Clone + 'static,
|
||||
{
|
||||
if !self.is_enabled() {
|
||||
return;
|
||||
}
|
||||
let now_unix_secs = now_unix_secs();
|
||||
match build_streaming_usage_record(
|
||||
plan,
|
||||
report_context,
|
||||
status_code,
|
||||
headers,
|
||||
telemetry,
|
||||
now_unix_secs,
|
||||
) {
|
||||
Ok(record) => {
|
||||
if let Err(err) = data.upsert_usage_record(record).await {
|
||||
let data = T::clone(data);
|
||||
let seed = seed.clone();
|
||||
let telemetry = telemetry.cloned();
|
||||
let request_id = seed.request_id.clone();
|
||||
spawn_on_usage_background_runtime(boxed_usage_task(async move {
|
||||
let now_unix_secs = now_unix_secs();
|
||||
match build_streaming_usage_record_offthread(
|
||||
&seed,
|
||||
status_code,
|
||||
telemetry.as_ref(),
|
||||
now_unix_secs,
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(record) => {
|
||||
if let Err(err) = data.upsert_usage_record(record).await {
|
||||
warn!(
|
||||
event_name = "usage_stream_record_failed",
|
||||
log_type = "event",
|
||||
request_id = %request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to record stream usage"
|
||||
);
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "usage_stream_record_failed",
|
||||
event_name = "usage_stream_build_failed",
|
||||
log_type = "event",
|
||||
request_id = %plan.request_id,
|
||||
request_id = %request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to record stream usage"
|
||||
);
|
||||
"usage runtime failed to build stream usage"
|
||||
)
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "usage_stream_build_failed",
|
||||
log_type = "event",
|
||||
request_id = %plan.request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to build stream usage"
|
||||
)
|
||||
}
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
pub async fn record_sync_terminal<T>(
|
||||
pub fn record_sync_terminal<T>(
|
||||
&self,
|
||||
data: &T,
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
payload: &GatewaySyncReportRequest,
|
||||
context_seed: &TerminalUsageContextSeed,
|
||||
payload_seed: &SyncTerminalUsagePayloadSeed,
|
||||
) where
|
||||
T: UsageRuntimeAccess,
|
||||
T: UsageRuntimeAccess + Clone + 'static,
|
||||
{
|
||||
if !self.is_enabled() {
|
||||
return;
|
||||
}
|
||||
match build_terminal_usage_event_from_outcome(build_sync_terminal_usage_outcome(
|
||||
plan,
|
||||
report_context,
|
||||
payload,
|
||||
)) {
|
||||
Ok(mut event) => {
|
||||
if let Err(err) = data.enrich_usage_event(&mut event).await {
|
||||
warn!(
|
||||
event_name = "usage_sync_terminal_billing_enrichment_failed",
|
||||
log_type = "event",
|
||||
request_id = %plan.request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to enrich sync usage event with billing"
|
||||
);
|
||||
let runtime = self.clone();
|
||||
let data = T::clone(data);
|
||||
let request_id = context_seed.request_id.clone();
|
||||
let input = Box::new(SyncTerminalUsageTaskInput {
|
||||
context_seed: context_seed.clone(),
|
||||
payload_seed: payload_seed.clone(),
|
||||
});
|
||||
spawn_on_usage_background_runtime(boxed_usage_task(async move {
|
||||
match build_sync_terminal_usage_event_offthread(input).await {
|
||||
Ok(mut event) => {
|
||||
if let Err(err) = data.enrich_usage_event(&mut event).await {
|
||||
warn!(
|
||||
event_name = "usage_sync_terminal_billing_enrichment_failed",
|
||||
log_type = "event",
|
||||
request_id = %request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to enrich sync usage event with billing"
|
||||
);
|
||||
}
|
||||
runtime.enqueue_or_write_terminal(&data, event).await
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "usage_sync_terminal_build_failed",
|
||||
log_type = "event",
|
||||
request_id = %request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to build sync terminal usage event"
|
||||
)
|
||||
}
|
||||
self.enqueue_or_write_terminal(data, event).await
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "usage_sync_terminal_build_failed",
|
||||
log_type = "event",
|
||||
request_id = %plan.request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to build sync terminal usage event"
|
||||
)
|
||||
}
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
pub async fn record_stream_terminal<T>(
|
||||
pub fn record_stream_terminal<T>(
|
||||
&self,
|
||||
data: &T,
|
||||
plan: &ExecutionPlan,
|
||||
report_context: Option<&serde_json::Value>,
|
||||
payload: &GatewayStreamReportRequest,
|
||||
context_seed: &TerminalUsageContextSeed,
|
||||
payload_seed: &StreamTerminalUsagePayloadSeed,
|
||||
cancelled: bool,
|
||||
) where
|
||||
T: UsageRuntimeAccess,
|
||||
T: UsageRuntimeAccess + Clone + 'static,
|
||||
{
|
||||
if !self.is_enabled() {
|
||||
return;
|
||||
}
|
||||
let mut outcome = build_stream_terminal_usage_outcome(plan, report_context, payload);
|
||||
if cancelled {
|
||||
outcome.terminal_state = UsageTerminalState::Cancelled;
|
||||
}
|
||||
match build_terminal_usage_event_from_outcome(outcome) {
|
||||
Ok(mut event) => {
|
||||
if let Err(err) = data.enrich_usage_event(&mut event).await {
|
||||
warn!(
|
||||
event_name = "usage_stream_terminal_billing_enrichment_failed",
|
||||
log_type = "event",
|
||||
request_id = %plan.request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to enrich stream usage event with billing"
|
||||
);
|
||||
let runtime = self.clone();
|
||||
let data = T::clone(data);
|
||||
let request_id = context_seed.request_id.clone();
|
||||
let input = Box::new(StreamTerminalUsageTaskInput {
|
||||
context_seed: context_seed.clone(),
|
||||
payload_seed: payload_seed.clone(),
|
||||
cancelled,
|
||||
});
|
||||
spawn_on_usage_background_runtime(boxed_usage_task(async move {
|
||||
match build_stream_terminal_usage_event_offthread(input).await {
|
||||
Ok(mut event) => {
|
||||
if let Err(err) = data.enrich_usage_event(&mut event).await {
|
||||
warn!(
|
||||
event_name = "usage_stream_terminal_billing_enrichment_failed",
|
||||
log_type = "event",
|
||||
request_id = %request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to enrich stream usage event with billing"
|
||||
);
|
||||
}
|
||||
runtime.enqueue_or_write_terminal(&data, event).await
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "usage_stream_terminal_build_failed",
|
||||
log_type = "event",
|
||||
request_id = %request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to build stream terminal usage event"
|
||||
)
|
||||
}
|
||||
self.enqueue_or_write_terminal(data, event).await
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
event_name = "usage_stream_terminal_build_failed",
|
||||
log_type = "event",
|
||||
request_id = %plan.request_id,
|
||||
error = %err,
|
||||
"usage runtime failed to build stream terminal usage event"
|
||||
)
|
||||
}
|
||||
}));
|
||||
}
|
||||
|
||||
pub fn submit_terminal_event<T>(&self, data: &T, event: UsageEvent)
|
||||
where
|
||||
T: UsageRuntimeAccess + Clone + 'static,
|
||||
{
|
||||
if !self.is_enabled() {
|
||||
return;
|
||||
}
|
||||
let runtime = self.clone();
|
||||
let data = T::clone(data);
|
||||
spawn_on_usage_background_runtime(boxed_usage_task(async move {
|
||||
runtime.record_terminal_event(&data, event).await;
|
||||
}));
|
||||
}
|
||||
|
||||
pub async fn record_terminal_event<T>(&self, data: &T, mut event: UsageEvent)
|
||||
@@ -326,6 +370,74 @@ impl UsageRuntime {
|
||||
}
|
||||
}
|
||||
|
||||
async fn build_pending_usage_record_offthread(
|
||||
seed: &LifecycleUsageSeed,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<UpsertUsageRecord, DataLayerError> {
|
||||
let seed = seed.clone();
|
||||
tokio::task::spawn_blocking(move || build_pending_usage_record_from_seed(&seed, now_unix_secs))
|
||||
.await
|
||||
.map_err(join_error_to_data_layer)?
|
||||
}
|
||||
|
||||
async fn build_streaming_usage_record_offthread(
|
||||
seed: &LifecycleUsageSeed,
|
||||
status_code: u16,
|
||||
telemetry: Option<&ExecutionTelemetry>,
|
||||
now_unix_secs: u64,
|
||||
) -> Result<UpsertUsageRecord, DataLayerError> {
|
||||
let seed = seed.clone();
|
||||
let telemetry = telemetry.cloned();
|
||||
tokio::task::spawn_blocking(move || {
|
||||
build_streaming_usage_record_from_seed(
|
||||
&seed,
|
||||
status_code,
|
||||
telemetry.as_ref(),
|
||||
now_unix_secs,
|
||||
)
|
||||
})
|
||||
.await
|
||||
.map_err(join_error_to_data_layer)?
|
||||
}
|
||||
|
||||
async fn build_sync_terminal_usage_event_offthread(
|
||||
input: Box<SyncTerminalUsageTaskInput>,
|
||||
) -> Result<UsageEvent, DataLayerError> {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
build_terminal_usage_event_from_seed(build_sync_terminal_usage_seed(
|
||||
input.context_seed,
|
||||
input.payload_seed,
|
||||
))
|
||||
})
|
||||
.await
|
||||
.map_err(join_error_to_data_layer)?
|
||||
}
|
||||
|
||||
async fn build_stream_terminal_usage_event_offthread(
|
||||
input: Box<StreamTerminalUsageTaskInput>,
|
||||
) -> Result<UsageEvent, DataLayerError> {
|
||||
tokio::task::spawn_blocking(move || {
|
||||
build_terminal_usage_event_from_seed(build_stream_terminal_usage_seed(
|
||||
input.context_seed,
|
||||
input.payload_seed,
|
||||
input.cancelled,
|
||||
))
|
||||
})
|
||||
.await
|
||||
.map_err(join_error_to_data_layer)?
|
||||
}
|
||||
|
||||
fn join_error_to_data_layer(err: tokio::task::JoinError) -> DataLayerError {
|
||||
DataLayerError::UnexpectedValue(format!("usage builder task join failed: {err}"))
|
||||
}
|
||||
|
||||
fn boxed_usage_task<F>(task: F) -> Pin<Box<dyn Future<Output = ()> + Send>>
|
||||
where
|
||||
F: Future<Output = ()> + Send + 'static,
|
||||
{
|
||||
Box::pin(task)
|
||||
}
|
||||
|
||||
fn now_unix_secs() -> u64 {
|
||||
std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
|
||||
@@ -7,6 +7,7 @@ use aether_data_contracts::DataLayerError;
|
||||
use async_trait::async_trait;
|
||||
use tracing::warn;
|
||||
|
||||
use crate::executor::spawn_on_usage_background_runtime;
|
||||
use crate::{
|
||||
build_upsert_usage_record_from_event, settle_usage_if_needed, UsageEvent, UsageQueue,
|
||||
UsageRuntimeConfig, UsageSettlementWriter,
|
||||
@@ -69,7 +70,7 @@ impl UsageQueueWorker {
|
||||
}
|
||||
|
||||
pub fn spawn(self) -> tokio::task::JoinHandle<()> {
|
||||
tokio::spawn(async move { self.run_forever().await })
|
||||
spawn_on_usage_background_runtime(async move { self.run_forever().await })
|
||||
}
|
||||
|
||||
async fn run_forever(self) {
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user