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

持久化 provider key 使用统计并移除 usage 汇总覆盖 (#313)

Browse Source 
* 持久化 provider key 使用统计并移除 usage 汇总覆盖

- 在 usage upsert 时按请求前后差值同步 provider_api_keys 统计
- 增加基于保留 usage 记录的 provider key 统计重建能力
- 号池管理列表直接读取 provider_api_keys 统计字段

* fix: harden provider key usage stats sync

---------

Co-authored-by: fawney19 <elky0401@gmail.com>
This commit is contained in:
AAEE86
2026-04-19 01:02:09 +08:00
committed by GitHub
parent 425227509a
commit bafb63a665
9 changed files with 1166 additions and 191 deletions
Download Patch File Download Diff File Expand all files Collapse all files

112
crates/aether-data/src/repository/provider_catalog/memory.rs
View File

@@ -8,6 +8,7 @@ use super::{
ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use crate::repository::usage::{ProviderApiKeyUsageContribution, ProviderApiKeyUsageDelta};
use crate::DataLayerError;
#[derive(Debug, Default)]
@@ -42,6 +43,117 @@ impl InMemoryProviderCatalogReadRepository {
}),
}
}
pub(crate) fn apply_usage_stats_delta(
&self,
key_id: &str,
delta: &ProviderApiKeyUsageDelta,
recomputed_last_used_at_unix_secs: Option<u64>,
) {
let mut index = self
.index
.write()
.expect("provider catalog repository lock");
let Some(key) = index.keys.get_mut(key_id) else {
return;
};
key.request_count = Some(apply_i64_delta_to_u32(
key.request_count.unwrap_or_default(),
delta.request_count,
));
key.success_count = Some(apply_i64_delta_to_u32(
key.success_count.unwrap_or_default(),
delta.success_count,
));
key.error_count = Some(apply_i64_delta_to_u32(
key.error_count.unwrap_or_default(),
delta.error_count,
));
key.total_tokens = apply_i64_delta_to_u64(key.total_tokens, delta.total_tokens);
key.total_cost_usd = apply_f64_delta(key.total_cost_usd, delta.total_cost_usd);
key.total_response_time_ms = Some(apply_i64_delta_to_u32(
key.total_response_time_ms.unwrap_or_default(),
delta.total_response_time_ms,
));
if let Some(candidate) = delta.candidate_last_used_at_unix_secs {
key.last_used_at_unix_secs = Some(
key.last_used_at_unix_secs
.map(|existing| existing.max(candidate))
.unwrap_or(candidate),
);
} else if delta.removed_last_used_at_unix_secs.is_some()
&& key.last_used_at_unix_secs == delta.removed_last_used_at_unix_secs
{
key.last_used_at_unix_secs = recomputed_last_used_at_unix_secs;
}
}
pub(crate) fn rebuild_usage_stats(
&self,
contributions: &BTreeMap<String, ProviderApiKeyUsageContribution>,
) {
let mut index = self
.index
.write()
.expect("provider catalog repository lock");
for key in index.keys.values_mut() {
key.request_count = Some(0);
key.success_count = Some(0);
key.error_count = Some(0);
key.total_tokens = 0;
key.total_cost_usd = 0.0;
key.total_response_time_ms = Some(0);
key.last_used_at_unix_secs = None;
}
for (key_id, contribution) in contributions {
let Some(key) = index.keys.get_mut(key_id) else {
continue;
};
key.request_count = Some(clamp_i64_to_u32(contribution.request_count));
key.success_count = Some(clamp_i64_to_u32(contribution.success_count));
key.error_count = Some(clamp_i64_to_u32(contribution.error_count));
key.total_tokens = clamp_i64_to_u64(contribution.total_tokens);
key.total_cost_usd = contribution.total_cost_usd.max(0.0);
key.total_response_time_ms =
Some(clamp_i64_to_u32(contribution.total_response_time_ms));
key.last_used_at_unix_secs = contribution.last_used_at_unix_secs;
}
}
}
fn apply_i64_delta_to_u32(current: u32, delta: i64) -> u32 {
clamp_i64_to_u32(i64::from(current).saturating_add(delta))
}
fn apply_i64_delta_to_u64(current: u64, delta: i64) -> u64 {
if delta >= 0 {
current.saturating_add(delta as u64)
} else {
current.saturating_sub(delta.unsigned_abs())
}
}
fn clamp_i64_to_u32(value: i64) -> u32 {
value.clamp(0, i64::from(u32::MAX)) as u32
}
fn clamp_i64_to_u64(value: i64) -> u64 {
value.max(0) as u64
}
fn apply_f64_delta(current: f64, delta: f64) -> f64 {
if !current.is_finite() && !delta.is_finite() {
return 0.0;
}
let next = current.max(0.0) + delta;
if next.is_finite() {
next.max(0.0)
} else {
0.0
}
}
#[async_trait]

556
crates/aether-data/src/repository/usage/memory.rs
View File

@@ -1,4 +1,5 @@
use std::collections::BTreeMap;
use std::sync::Arc;
use std::sync::RwLock;
use aether_data_contracts::repository::usage::{
@@ -23,11 +24,13 @@ use chrono::Utc;
use serde_json::Value;
use super::{
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
provider_api_key_usage_contribution, strip_deprecated_usage_display_fields,
usage_can_recover_terminal_failure, ProviderApiKeyUsageContribution, ProviderApiKeyUsageDelta,
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery,
UsageDailyHeatmapQuery, UsageReadRepository, UsageWriteRepository,
};
use crate::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use crate::DataLayerError;
#[derive(Debug, Default)]
@@ -35,6 +38,7 @@ pub struct InMemoryUsageReadRepository {
by_request_id: RwLock<BTreeMap<String, StoredRequestUsageAudit>>,
detached_bodies: RwLock<BTreeMap<String, Value>>,
provider_usage_windows: RwLock<Vec<StoredProviderUsageWindow>>,
provider_catalog: Option<Arc<InMemoryProviderCatalogReadRepository>>,
}
impl InMemoryUsageReadRepository {
@@ -51,6 +55,7 @@ impl InMemoryUsageReadRepository {
by_request_id: RwLock::new(by_request_id),
detached_bodies: RwLock::new(BTreeMap::new()),
provider_usage_windows: RwLock::new(Vec::new()),
provider_catalog: None,
}
}
@@ -101,6 +106,7 @@ impl InMemoryUsageReadRepository {
by_request_id: RwLock::new(by_request_id),
detached_bodies: RwLock::new(detached_bodies),
provider_usage_windows: RwLock::new(Vec::new()),
provider_catalog: None,
}
}
@@ -112,8 +118,17 @@ impl InMemoryUsageReadRepository {
by_request_id: self.by_request_id,
detached_bodies: self.detached_bodies,
provider_usage_windows: RwLock::new(items.into_iter().collect()),
provider_catalog: self.provider_catalog,
}
}
pub fn with_provider_catalog_repository(
mut self,
repository: Arc<InMemoryProviderCatalogReadRepository>,
) -> Self {
self.provider_catalog = Some(repository);
self
}
}
fn usage_status_is_finalized(status: &str) -> bool {
@@ -124,6 +139,38 @@ fn usage_status_is_lifecycle(status: &str) -> bool {
matches!(status, "pending" | "streaming")
}
fn accumulate_provider_api_key_usage_contribution(
aggregates: &mut BTreeMap<String, ProviderApiKeyUsageContribution>,
contribution: ProviderApiKeyUsageContribution,
) {
let entry = aggregates
.entry(contribution.key_id.clone())
.or_insert_with(|| ProviderApiKeyUsageContribution {
key_id: contribution.key_id.clone(),
..ProviderApiKeyUsageContribution::default()
});
entry.request_count = entry
.request_count
.saturating_add(contribution.request_count);
entry.success_count = entry
.success_count
.saturating_add(contribution.success_count);
entry.error_count = entry.error_count.saturating_add(contribution.error_count);
entry.total_tokens = entry.total_tokens.saturating_add(contribution.total_tokens);
entry.total_cost_usd += contribution.total_cost_usd;
entry.total_response_time_ms = entry
.total_response_time_ms
.saturating_add(contribution.total_response_time_ms);
entry.last_used_at_unix_secs = match (
entry.last_used_at_unix_secs,
contribution.last_used_at_unix_secs,
) {
(Some(existing), Some(candidate)) => Some(existing.max(candidate)),
(None, Some(candidate)) => Some(candidate),
(existing, None) => existing,
};
}
fn usage_matches_list_query(item: &StoredRequestUsageAudit, query: &UsageAuditListQuery) -> bool {
// The field is historically named `created_at_unix_ms`, but usage audit rows
// across gateway handlers, SQL repositories and tests are stored as epoch seconds.
@@ -2039,6 +2086,7 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
usage.validate()?;
let usage = strip_deprecated_usage_display_fields(usage);
let mut by_request_id = self.by_request_id.write().expect("usage repository lock");
let existing = by_request_id.get(&usage.request_id).cloned();
let created_at_unix_ms = by_request_id
.get(&usage.request_id)
@@ -2055,8 +2103,7 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
)
})
.unwrap_or_default();
let existing = by_request_id.get(&usage.request_id);
if existing.is_some_and(|existing| {
if existing.as_ref().is_some_and(|existing| {
usage_status_is_finalized(existing.status.as_str())
&& usage_status_is_lifecycle(usage.status.as_str())
&& !usage_can_recover_terminal_failure(
@@ -2068,7 +2115,7 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
}) {
return Ok(existing.expect("existing usage should be present").clone());
}
if existing.is_some_and(|existing| {
if existing.as_ref().is_some_and(|existing| {
existing.billing_status == "pending"
&& existing.status == "streaming"
&& usage.status == "pending"
@@ -2076,47 +2123,53 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
return Ok(existing.expect("existing usage should be present").clone());
}
let request_metadata = usage
.request_metadata
.clone()
.or_else(|| existing.and_then(|existing| existing.request_metadata.clone()));
let request_metadata = usage.request_metadata.clone().or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.request_metadata.clone())
});
let request_body_ref = persisted_usage_body_ref(
usage.request_body_ref.as_deref(),
usage.request_body.as_ref(),
request_metadata.as_ref(),
existing,
existing.as_ref(),
UsageBodyField::RequestBody,
);
let provider_request_body_ref = persisted_usage_body_ref(
usage.provider_request_body_ref.as_deref(),
usage.provider_request_body.as_ref(),
request_metadata.as_ref(),
existing,
existing.as_ref(),
UsageBodyField::ProviderRequestBody,
);
let response_body_ref = persisted_usage_body_ref(
usage.response_body_ref.as_deref(),
usage.response_body.as_ref(),
request_metadata.as_ref(),
existing,
existing.as_ref(),
UsageBodyField::ResponseBody,
);
let client_response_body_ref = persisted_usage_body_ref(
usage.client_response_body_ref.as_deref(),
usage.client_response_body.as_ref(),
request_metadata.as_ref(),
existing,
existing.as_ref(),
UsageBodyField::ClientResponseBody,
);
let stored = StoredRequestUsageAudit {
id: existing
.as_ref()
.map(|existing| existing.id.clone())
.unwrap_or_else(|| format!("usage-{}", usage.request_id)),
request_id: usage.request_id.clone(),
user_id: usage.user_id,
api_key_id: usage.api_key_id,
username: existing.and_then(|existing| existing.username.clone()),
api_key_name: existing.and_then(|existing| existing.api_key_name.clone()),
username: existing
.as_ref()
.and_then(|existing| existing.username.clone()),
api_key_name: existing
.as_ref()
.and_then(|existing| existing.api_key_name.clone()),
provider_name: usage.provider_name,
model: usage.model,
target_model: usage.target_model,
@@ -2137,6 +2190,7 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
total_tokens,
cache_creation_input_tokens: usage.cache_creation_input_tokens.unwrap_or_else(|| {
existing
.as_ref()
.map(|existing| existing.cache_creation_input_tokens)
.unwrap_or_default()
}),
@@ -2144,6 +2198,7 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
.cache_creation_ephemeral_5m_input_tokens
.unwrap_or_else(|| {
existing
.as_ref()
.map(|existing| existing.cache_creation_ephemeral_5m_input_tokens)
.unwrap_or_default()
}),
@@ -2151,32 +2206,40 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
.cache_creation_ephemeral_1h_input_tokens
.unwrap_or_else(|| {
existing
.as_ref()
.map(|existing| existing.cache_creation_ephemeral_1h_input_tokens)
.unwrap_or_default()
}),
cache_read_input_tokens: usage.cache_read_input_tokens.unwrap_or_else(|| {
existing
.as_ref()
.map(|existing| existing.cache_read_input_tokens)
.unwrap_or_default()
}),
cache_creation_cost_usd: usage.cache_creation_cost_usd.unwrap_or_else(|| {
existing
.as_ref()
.map(|existing| existing.cache_creation_cost_usd)
.unwrap_or_default()
}),
cache_read_cost_usd: usage.cache_read_cost_usd.unwrap_or_else(|| {
existing
.as_ref()
.map(|existing| existing.cache_read_cost_usd)
.unwrap_or_default()
}),
output_price_per_1m: existing.and_then(|existing| existing.output_price_per_1m),
output_price_per_1m: existing
.as_ref()
.and_then(|existing| existing.output_price_per_1m),
total_cost_usd: usage.total_cost_usd.unwrap_or_else(|| {
existing
.as_ref()
.map(|existing| existing.total_cost_usd)
.unwrap_or_default()
}),
actual_total_cost_usd: usage.actual_total_cost_usd.unwrap_or_else(|| {
existing
.as_ref()
.map(|existing| existing.actual_total_cost_usd)
.unwrap_or_default()
}),
@@ -2187,71 +2250,108 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
first_byte_time_ms: usage.first_byte_time_ms,
status: usage.status,
billing_status: usage.billing_status,
request_headers: usage
.request_headers
.or_else(|| existing.and_then(|existing| existing.request_headers.clone())),
request_body: usage
.request_body
.or_else(|| existing.and_then(|existing| existing.request_body.clone())),
request_headers: usage.request_headers.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.request_headers.clone())
}),
request_body: usage.request_body.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.request_body.clone())
}),
request_body_ref,
request_body_state: usage
.request_body_state
.or_else(|| existing.and_then(|existing| existing.request_body_state)),
request_body_state: usage.request_body_state.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.request_body_state)
}),
provider_request_headers: usage.provider_request_headers.or_else(|| {
existing.and_then(|existing| existing.provider_request_headers.clone())
existing
.as_ref()
.and_then(|existing| existing.provider_request_headers.clone())
}),
provider_request_body: usage.provider_request_body.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.provider_request_body.clone())
}),
provider_request_body: usage
.provider_request_body
.or_else(|| existing.and_then(|existing| existing.provider_request_body.clone())),
provider_request_body_ref,
provider_request_body_state: usage
.provider_request_body_state
.or_else(|| existing.and_then(|existing| existing.provider_request_body_state)),
response_headers: usage
.response_headers
.or_else(|| existing.and_then(|existing| existing.response_headers.clone())),
response_body: usage
.response_body
.or_else(|| existing.and_then(|existing| existing.response_body.clone())),
response_body_ref,
response_body_state: usage
.response_body_state
.or_else(|| existing.and_then(|existing| existing.response_body_state)),
client_response_headers: usage
.client_response_headers
.or_else(|| existing.and_then(|existing| existing.client_response_headers.clone())),
client_response_body: usage
.client_response_body
.or_else(|| existing.and_then(|existing| existing.client_response_body.clone())),
client_response_body_ref,
client_response_body_state: usage
.client_response_body_state
.or_else(|| existing.and_then(|existing| existing.client_response_body_state)),
candidate_id: usage.candidate_id.or_else(|| {
existing.and_then(|existing| existing.routing_candidate_id().map(ToOwned::to_owned))
provider_request_body_state: usage.provider_request_body_state.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.provider_request_body_state)
}),
response_headers: usage.response_headers.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.response_headers.clone())
}),
response_body: usage.response_body.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.response_body.clone())
}),
response_body_ref,
response_body_state: usage.response_body_state.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.response_body_state)
}),
client_response_headers: usage.client_response_headers.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.client_response_headers.clone())
}),
client_response_body: usage.client_response_body.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.client_response_body.clone())
}),
client_response_body_ref,
client_response_body_state: usage.client_response_body_state.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.client_response_body_state)
}),
candidate_id: usage.candidate_id.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.routing_candidate_id().map(ToOwned::to_owned))
}),
candidate_index: usage.candidate_index.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.routing_candidate_index())
}),
candidate_index: usage
.candidate_index
.or_else(|| existing.and_then(|existing| existing.routing_candidate_index())),
key_name: usage.key_name.or_else(|| {
existing.and_then(|existing| existing.routing_key_name().map(ToOwned::to_owned))
existing
.as_ref()
.and_then(|existing| existing.routing_key_name().map(ToOwned::to_owned))
}),
planner_kind: usage.planner_kind.or_else(|| {
existing.and_then(|existing| existing.routing_planner_kind().map(ToOwned::to_owned))
existing
.as_ref()
.and_then(|existing| existing.routing_planner_kind().map(ToOwned::to_owned))
}),
route_family: usage.route_family.or_else(|| {
existing.and_then(|existing| existing.routing_route_family().map(ToOwned::to_owned))
existing
.as_ref()
.and_then(|existing| existing.routing_route_family().map(ToOwned::to_owned))
}),
route_kind: usage.route_kind.or_else(|| {
existing.and_then(|existing| existing.routing_route_kind().map(ToOwned::to_owned))
existing
.as_ref()
.and_then(|existing| existing.routing_route_kind().map(ToOwned::to_owned))
}),
execution_path: usage.execution_path.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.routing_execution_path().map(ToOwned::to_owned))
}),
local_execution_runtime_miss_reason: usage.local_execution_runtime_miss_reason.or_else(
|| {
existing.and_then(|existing| {
existing.as_ref().and_then(|existing| {
existing
.routing_local_execution_runtime_miss_reason()
.map(ToOwned::to_owned)
@@ -2265,13 +2365,76 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
};
by_request_id.insert(stored.request_id.clone(), stored.clone());
if let Some(provider_catalog) = self.provider_catalog.as_ref() {
let before_contribution = existing
.as_ref()
.and_then(provider_api_key_usage_contribution);
let after_contribution = provider_api_key_usage_contribution(&stored);
match (before_contribution.as_ref(), after_contribution.as_ref()) {
(Some(before), Some(after)) if before.key_id == after.key_id => {
let delta = ProviderApiKeyUsageDelta::between(before, after);
provider_catalog.apply_usage_stats_delta(before.key_id.as_str(), &delta, None);
}
_ => {
if let Some(before) = before_contribution.as_ref() {
let delta = ProviderApiKeyUsageDelta::removal(before);
let recomputed_last_used_at_unix_secs = by_request_id
.values()
.filter_map(|item| {
item.provider_api_key_id
.as_deref()
.filter(|key_id| *key_id == before.key_id.as_str())
.map(|_| item.created_at_unix_ms)
})
.max();
provider_catalog.apply_usage_stats_delta(
before.key_id.as_str(),
&delta,
recomputed_last_used_at_unix_secs,
);
}
if let Some(after) = after_contribution.as_ref() {
let delta = ProviderApiKeyUsageDelta::addition(after);
provider_catalog.apply_usage_stats_delta(
after.key_id.as_str(),
&delta,
None,
);
}
}
}
}
Ok(stored)
}
async fn rebuild_provider_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
let Some(provider_catalog) = self.provider_catalog.as_ref() else {
return Ok(0);
};
let by_request_id = self.by_request_id.read().expect("usage repository lock");
let mut aggregates = BTreeMap::new();
for usage in by_request_id.values() {
let Some(contribution) = provider_api_key_usage_contribution(usage) else {
continue;
};
accumulate_provider_api_key_usage_contribution(&mut aggregates, contribution);
}
provider_catalog.rebuild_usage_stats(&aggregates);
Ok(aggregates.len() as u64)
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::InMemoryUsageReadRepository;
use crate::repository::provider_catalog::{
InMemoryProviderCatalogReadRepository, ProviderCatalogReadRepository,
ProviderCatalogWriteRepository, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use crate::repository::usage::{
StoredProviderUsageWindow, StoredRequestUsageAudit, UpsertUsageRecord, UsageReadRepository,
UsageWriteRepository,
@@ -2321,6 +2484,78 @@ mod tests {
.expect("usage should build")
}
fn sample_upsert_usage_record(request_id: &str) -> UpsertUsageRecord {
UpsertUsageRecord {
request_id: request_id.to_string(),
user_id: None,
api_key_id: None,
username: None,
api_key_name: None,
provider_name: "OpenAI".to_string(),
model: "gpt-5".to_string(),
target_model: None,
provider_id: Some("provider-1".to_string()),
provider_endpoint_id: None,
provider_api_key_id: None,
request_type: None,
api_format: None,
api_family: None,
endpoint_kind: None,
endpoint_api_format: None,
provider_api_family: None,
provider_endpoint_kind: None,
has_format_conversion: Some(false),
is_stream: Some(false),
input_tokens: None,
output_tokens: None,
total_tokens: None,
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: None,
actual_total_cost_usd: None,
status_code: None,
error_message: None,
error_category: None,
response_time_ms: None,
first_byte_time_ms: None,
status: "pending".to_string(),
billing_status: "pending".to_string(),
request_headers: None,
request_body: None,
request_body_ref: None,
request_body_state: None,
provider_request_headers: None,
provider_request_body: None,
provider_request_body_ref: None,
provider_request_body_state: None,
response_headers: None,
response_body: None,
response_body_ref: None,
response_body_state: None,
client_response_headers: None,
client_response_body: None,
client_response_body_ref: None,
client_response_body_state: 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(1_700_000_000),
updated_at_unix_secs: 1_700_000_000,
}
}
#[tokio::test]
async fn finds_usage_by_request_id() {
let repository = InMemoryUsageReadRepository::seed(vec![
@@ -3719,4 +3954,203 @@ mod tests {
assert_eq!(usage.request_count, 2);
assert_eq!(usage.last_used_at_unix_secs, Some(2_500));
}
fn sample_provider_catalog_key(key_id: &str) -> StoredProviderCatalogKey {
StoredProviderCatalogKey::new(
key_id.to_string(),
"provider-1".to_string(),
"provider key".to_string(),
"api_key".to_string(),
None,
true,
)
.expect("provider key should build")
}
fn sample_provider_catalog_repository(
key_ids: &[&str],
) -> Arc<InMemoryProviderCatalogReadRepository> {
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![StoredProviderCatalogProvider::new(
"provider-1".to_string(),
"OpenAI".to_string(),
None,
"openai".to_string(),
)
.expect("provider should build")],
Vec::new(),
key_ids
.iter()
.map(|key_id| sample_provider_catalog_key(key_id))
.collect(),
))
}
#[tokio::test]
async fn upsert_syncs_linked_provider_key_stats_without_double_counting_request_count() {
let provider_catalog = sample_provider_catalog_repository(&["provider-key-1"]);
let repository = InMemoryUsageReadRepository::default()
.with_provider_catalog_repository(Arc::clone(&provider_catalog));
repository
.upsert(UpsertUsageRecord {
provider_api_key_id: Some("provider-key-1".to_string()),
total_tokens: Some(100),
total_cost_usd: Some(0.5),
created_at_unix_ms: Some(1_711_100_000),
updated_at_unix_secs: 1_711_100_000,
..sample_upsert_usage_record("req-linked-1")
})
.await
.expect("pending upsert should succeed");
repository
.upsert(UpsertUsageRecord {
provider_api_key_id: Some("provider-key-1".to_string()),
status: "completed".to_string(),
billing_status: "settled".to_string(),
status_code: Some(200),
response_time_ms: Some(240),
total_tokens: Some(180),
total_cost_usd: Some(0.75),
created_at_unix_ms: Some(1_711_100_000),
updated_at_unix_secs: 1_711_100_010,
finalized_at_unix_secs: Some(1_711_100_011),
..sample_upsert_usage_record("req-linked-1")
})
.await
.expect("completed upsert should succeed");
let key = provider_catalog
.list_keys_by_ids(&["provider-key-1".to_string()])
.await
.expect("key list should succeed")
.into_iter()
.next()
.expect("provider key should exist");
assert_eq!(key.request_count, Some(1));
assert_eq!(key.success_count, Some(1));
assert_eq!(key.error_count, Some(0));
assert_eq!(key.total_tokens, 180);
assert_eq!(key.total_cost_usd, 0.75);
assert_eq!(key.total_response_time_ms, Some(240));
assert_eq!(key.last_used_at_unix_secs, Some(1_711_100_000));
}
#[tokio::test]
async fn upsert_moves_linked_provider_key_stats_when_key_assignment_changes() {
let provider_catalog =
sample_provider_catalog_repository(&["provider-key-a", "provider-key-b"]);
let repository = InMemoryUsageReadRepository::default()
.with_provider_catalog_repository(Arc::clone(&provider_catalog));
repository
.upsert(UpsertUsageRecord {
provider_api_key_id: Some("provider-key-a".to_string()),
status: "completed".to_string(),
billing_status: "settled".to_string(),
status_code: Some(200),
response_time_ms: Some(100),
total_tokens: Some(120),
total_cost_usd: Some(0.4),
created_at_unix_ms: Some(1_711_200_000),
updated_at_unix_secs: 1_711_200_000,
finalized_at_unix_secs: Some(1_711_200_001),
..sample_upsert_usage_record("req-move-1")
})
.await
.expect("first upsert should succeed");
repository
.upsert(UpsertUsageRecord {
provider_api_key_id: Some("provider-key-b".to_string()),
status: "completed".to_string(),
billing_status: "settled".to_string(),
status_code: Some(200),
response_time_ms: Some(150),
total_tokens: Some(140),
total_cost_usd: Some(0.6),
created_at_unix_ms: Some(1_711_200_000),
updated_at_unix_secs: 1_711_200_010,
finalized_at_unix_secs: Some(1_711_200_011),
..sample_upsert_usage_record("req-move-1")
})
.await
.expect("moved upsert should succeed");
let keys = provider_catalog
.list_keys_by_ids(&["provider-key-a".to_string(), "provider-key-b".to_string()])
.await
.expect("key list should succeed");
let key_a = keys
.iter()
.find(|key| key.id == "provider-key-a")
.expect("key a should exist");
let key_b = keys
.iter()
.find(|key| key.id == "provider-key-b")
.expect("key b should exist");
assert_eq!(key_a.request_count, Some(0));
assert_eq!(key_a.success_count, Some(0));
assert_eq!(key_a.total_tokens, 0);
assert_eq!(key_a.total_cost_usd, 0.0);
assert_eq!(key_a.total_response_time_ms, Some(0));
assert_eq!(key_a.last_used_at_unix_secs, None);
assert_eq!(key_b.request_count, Some(1));
assert_eq!(key_b.success_count, Some(1));
assert_eq!(key_b.total_tokens, 140);
assert_eq!(key_b.total_cost_usd, 0.6);
assert_eq!(key_b.total_response_time_ms, Some(150));
assert_eq!(key_b.last_used_at_unix_secs, Some(1_711_200_000));
}
#[tokio::test]
async fn rebuild_provider_key_usage_stats_resets_linked_catalog_to_current_usage() {
let provider_catalog = sample_provider_catalog_repository(&["provider-key-1"]);
let mut stale_key = provider_catalog
.list_keys_by_ids(&["provider-key-1".to_string()])
.await
.expect("key list should succeed")
.into_iter()
.next()
.expect("provider key should exist");
stale_key.request_count = Some(99);
stale_key.success_count = Some(88);
stale_key.error_count = Some(11);
stale_key.total_tokens = 9_999;
stale_key.total_cost_usd = 42.0;
stale_key.total_response_time_ms = Some(9_999);
stale_key.last_used_at_unix_secs = Some(9_999);
provider_catalog
.update_key(&stale_key)
.await
.expect("stale key should update");
let repository = InMemoryUsageReadRepository::seed(vec![
sample_usage("req-1", 1_711_300_000),
sample_usage("req-2", 1_711_300_250),
])
.with_provider_catalog_repository(Arc::clone(&provider_catalog));
let rebuilt = repository
.rebuild_provider_api_key_usage_stats()
.await
.expect("rebuild should succeed");
assert_eq!(rebuilt, 1);
let key = provider_catalog
.list_keys_by_ids(&["provider-key-1".to_string()])
.await
.expect("key list should succeed")
.into_iter()
.next()
.expect("provider key should exist");
assert_eq!(key.request_count, Some(2));
assert_eq!(key.success_count, Some(2));
assert_eq!(key.error_count, Some(0));
assert_eq!(key.total_tokens, 300);
assert_eq!(key.total_cost_usd, 0.24);
assert_eq!(key.total_response_time_ms, Some(840));
assert_eq!(key.last_used_at_unix_secs, Some(1_711_300_250));
}
}

243
crates/aether-data/src/repository/usage/mod.rs
View File

@@ -26,6 +26,85 @@ pub(crate) use aether_data_contracts::repository::usage::{
pub use memory::InMemoryUsageReadRepository;
pub use sql::SqlxUsageReadRepository;
#[derive(Debug, Clone, PartialEq, Default)]
pub(crate) struct ProviderApiKeyUsageContribution {
pub key_id: String,
pub request_count: i64,
pub success_count: i64,
pub error_count: i64,
pub total_tokens: i64,
pub total_cost_usd: f64,
pub total_response_time_ms: i64,
pub last_used_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Default)]
pub(crate) struct ProviderApiKeyUsageDelta {
pub request_count: i64,
pub success_count: i64,
pub error_count: i64,
pub total_tokens: i64,
pub total_cost_usd: f64,
pub total_response_time_ms: i64,
pub candidate_last_used_at_unix_secs: Option<u64>,
pub removed_last_used_at_unix_secs: Option<u64>,
}
impl ProviderApiKeyUsageDelta {
pub(crate) fn between(
before: &ProviderApiKeyUsageContribution,
after: &ProviderApiKeyUsageContribution,
) -> Self {
Self {
request_count: after.request_count - before.request_count,
success_count: after.success_count - before.success_count,
error_count: after.error_count - before.error_count,
total_tokens: after.total_tokens - before.total_tokens,
total_cost_usd: after.total_cost_usd - before.total_cost_usd,
total_response_time_ms: after.total_response_time_ms - before.total_response_time_ms,
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
removed_last_used_at_unix_secs: None,
}
}
pub(crate) fn addition(after: &ProviderApiKeyUsageContribution) -> Self {
Self {
request_count: after.request_count,
success_count: after.success_count,
error_count: after.error_count,
total_tokens: after.total_tokens,
total_cost_usd: after.total_cost_usd,
total_response_time_ms: after.total_response_time_ms,
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
removed_last_used_at_unix_secs: None,
}
}
pub(crate) fn removal(before: &ProviderApiKeyUsageContribution) -> Self {
Self {
request_count: -before.request_count,
success_count: -before.success_count,
error_count: -before.error_count,
total_tokens: -before.total_tokens,
total_cost_usd: -before.total_cost_usd,
total_response_time_ms: -before.total_response_time_ms,
candidate_last_used_at_unix_secs: None,
removed_last_used_at_unix_secs: before.last_used_at_unix_secs,
}
}
pub(crate) fn is_noop(&self) -> bool {
self.request_count == 0
&& self.success_count == 0
&& self.error_count == 0
&& self.total_tokens == 0
&& self.total_cost_usd == 0.0
&& self.total_response_time_ms == 0
&& self.candidate_last_used_at_unix_secs.is_none()
&& self.removed_last_used_at_unix_secs.is_none()
}
}
pub(crate) fn incoming_usage_can_recover_terminal_failure(
incoming_status: &str,
incoming_billing_status: &str,
@@ -57,11 +136,77 @@ pub(crate) fn strip_deprecated_usage_display_fields(
usage
}
pub(crate) fn provider_api_key_usage_is_success(
status: &str,
status_code: Option<u16>,
error_message: Option<&str>,
) -> bool {
matches!(
status,
"completed" | "success" | "ok" | "billed" | "settled"
) && status_code.is_none_or(|code| code < 400)
&& error_message.is_none_or(|value| value.trim().is_empty())
}
pub(crate) fn provider_api_key_usage_is_error(
status: &str,
status_code: Option<u16>,
error_message: Option<&str>,
) -> bool {
!matches!(status, "pending" | "streaming")
&& !provider_api_key_usage_is_success(status, status_code, error_message)
}
pub(crate) fn provider_api_key_usage_contribution(
usage: &StoredRequestUsageAudit,
) -> Option<ProviderApiKeyUsageContribution> {
let key_id = usage
.provider_api_key_id
.as_deref()
.map(str::trim)
.filter(|value| !value.is_empty())?
.to_string();
let is_success = provider_api_key_usage_is_success(
usage.status.as_str(),
usage.status_code,
usage.error_message.as_deref(),
);
let is_error = provider_api_key_usage_is_error(
usage.status.as_str(),
usage.status_code,
usage.error_message.as_deref(),
);
Some(ProviderApiKeyUsageContribution {
key_id,
request_count: 1,
success_count: i64::from(is_success),
error_count: i64::from(is_error),
total_tokens: i64::try_from(usage.total_tokens).unwrap_or(i64::MAX),
total_cost_usd: if usage.total_cost_usd.is_finite() {
usage.total_cost_usd.max(0.0)
} else {
0.0
},
total_response_time_ms: if is_success {
usage
.response_time_ms
.and_then(|value| i64::try_from(value).ok())
.unwrap_or_default()
} else {
0
},
last_used_at_unix_secs: Some(usage.created_at_unix_ms),
})
}
#[cfg(test)]
mod tests {
use super::{
incoming_usage_can_recover_terminal_failure, strip_deprecated_usage_display_fields,
usage_can_recover_terminal_failure, UpsertUsageRecord,
incoming_usage_can_recover_terminal_failure, provider_api_key_usage_contribution,
provider_api_key_usage_is_error, provider_api_key_usage_is_success,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
StoredRequestUsageAudit, UpsertUsageRecord,
};
#[test]
@@ -187,4 +332,98 @@ mod tests {
"failed", "void", "failed", "void"
));
}
#[test]
fn provider_key_usage_success_requires_clean_terminal_success() {
assert!(provider_api_key_usage_is_success(
"completed",
Some(200),
None
));
assert!(!provider_api_key_usage_is_success(
"completed",
Some(500),
None
));
assert!(!provider_api_key_usage_is_success(
"completed",
Some(200),
Some("boom")
));
assert!(!provider_api_key_usage_is_success(
"streaming",
Some(200),
None
));
}
#[test]
fn provider_key_usage_error_ignores_pending_states() {
assert!(provider_api_key_usage_is_error(
"failed",
Some(500),
Some("boom")
));
assert!(provider_api_key_usage_is_error(
"completed",
Some(200),
Some("boom")
));
assert!(!provider_api_key_usage_is_error("pending", None, None));
assert!(!provider_api_key_usage_is_error("streaming", None, None));
}
#[test]
fn provider_key_usage_contribution_tracks_success_response_time() {
let usage = StoredRequestUsageAudit::new(
"usage-1".to_string(),
"request-1".to_string(),
None,
None,
None,
None,
"OpenAI".to_string(),
"gpt-5".to_string(),
None,
Some("provider-1".to_string()),
None,
Some("provider-key-1".to_string()),
None,
None,
None,
None,
None,
None,
None,
false,
false,
12,
8,
20,
0.25,
0.25,
Some(200),
None,
None,
Some(120),
None,
"completed".to_string(),
"settled".to_string(),
123,
124,
Some(125),
)
.expect("usage should build");
let contribution =
provider_api_key_usage_contribution(&usage).expect("contribution should exist");
assert_eq!(contribution.key_id, "provider-key-1");
assert_eq!(contribution.request_count, 1);
assert_eq!(contribution.success_count, 1);
assert_eq!(contribution.error_count, 0);
assert_eq!(contribution.total_tokens, 20);
assert_eq!(contribution.total_cost_usd, 0.25);
assert_eq!(contribution.total_response_time_ms, 120);
assert_eq!(contribution.last_used_at_unix_secs, Some(123));
}
}

283
crates/aether-data/src/repository/usage/sql.rs
View File

@@ -26,7 +26,8 @@ use std::io::{Read, Write};
use uuid::Uuid;
use super::{
incoming_usage_can_recover_terminal_failure, strip_deprecated_usage_display_fields,
incoming_usage_can_recover_terminal_failure, provider_api_key_usage_contribution,
strip_deprecated_usage_display_fields, ProviderApiKeyUsageDelta,
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredRequestUsageAudit,
StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery, UsageDailyHeatmapQuery,
UsageReadRepository, UsageWriteRepository,
@@ -78,6 +79,9 @@ SET
WHERE request_id = $1
AND billing_status = 'void'
"#;
const LOCK_USAGE_REQUEST_ID_SQL: &str = r#"
SELECT pg_advisory_xact_lock(hashtext($1)::BIGINT)
"#;
const UPSERT_USAGE_HTTP_AUDIT_SQL: &str = r#"
INSERT INTO usage_http_audits (
request_id,
@@ -563,6 +567,112 @@ GROUP BY provider_api_key_id
ORDER BY provider_api_key_id ASC
"#;
const APPLY_PROVIDER_API_KEY_USAGE_DELTA_SQL: &str = r#"
UPDATE provider_api_keys
SET
request_count = GREATEST(COALESCE(request_count, 0) + $2, 0),
success_count = GREATEST(COALESCE(success_count, 0) + $3, 0),
error_count = GREATEST(COALESCE(error_count, 0) + $4, 0),
total_tokens = GREATEST(total_tokens + $5, 0),
total_cost_usd = CAST(
GREATEST(CAST(total_cost_usd AS DOUBLE PRECISION) + $6, 0) AS NUMERIC(20,8)
),
total_response_time_ms = GREATEST(COALESCE(total_response_time_ms, 0) + $7, 0),
last_used_at = CASE
WHEN $8::double precision IS NOT NULL THEN CASE
WHEN last_used_at IS NULL THEN TO_TIMESTAMP($8::double precision)
ELSE GREATEST(last_used_at, TO_TIMESTAMP($8::double precision))
END
WHEN $9::double precision IS NOT NULL
AND last_used_at IS NOT NULL
AND EXTRACT(EPOCH FROM last_used_at)::BIGINT = $9::BIGINT
THEN (
SELECT MAX(created_at)
FROM "usage"
WHERE provider_api_key_id = $1
)
ELSE last_used_at
END
WHERE id = $1
"#;
const RESET_PROVIDER_API_KEY_USAGE_STATS_SQL: &str = r#"
UPDATE provider_api_keys
SET
request_count = 0,
success_count = 0,
error_count = 0,
total_tokens = 0,
total_cost_usd = 0,
total_response_time_ms = 0,
last_used_at = NULL
"#;
const REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL: &str = r#"
WITH aggregated AS (
SELECT
provider_api_key_id,
COUNT(*)::INTEGER AS request_count,
COALESCE(SUM(
CASE
WHEN status IN ('completed', 'success', 'ok', 'billed', 'settled')
AND (status_code IS NULL OR status_code < 400)
AND NULLIF(BTRIM(error_message), '') IS NULL
THEN 1
ELSE 0
END
), 0)::INTEGER AS success_count,
COALESCE(SUM(
CASE
WHEN status NOT IN ('pending', 'streaming')
AND NOT (
status IN ('completed', 'success', 'ok', 'billed', 'settled')
AND (status_code IS NULL OR status_code < 400)
AND NULLIF(BTRIM(error_message), '') IS NULL
)
THEN 1
ELSE 0
END
), 0)::INTEGER AS error_count,
COALESCE(SUM(
GREATEST(
COALESCE(
total_tokens,
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
),
0
)::BIGINT
), 0)::BIGINT AS total_tokens,
COALESCE(SUM(COALESCE(total_cost_usd, 0)), 0)::NUMERIC(20,8) AS total_cost_usd,
COALESCE(SUM(
CASE
WHEN status IN ('completed', 'success', 'ok', 'billed', 'settled')
AND (status_code IS NULL OR status_code < 400)
AND NULLIF(BTRIM(error_message), '') IS NULL
AND response_time_ms IS NOT NULL
THEN GREATEST(response_time_ms, 0)
ELSE 0
END
), 0)::INTEGER AS total_response_time_ms,
MAX(created_at) AS last_used_at
FROM "usage"
WHERE provider_api_key_id IS NOT NULL
AND BTRIM(provider_api_key_id) <> ''
GROUP BY provider_api_key_id
)
UPDATE provider_api_keys
SET
request_count = aggregated.request_count,
success_count = aggregated.success_count,
error_count = aggregated.error_count,
total_tokens = aggregated.total_tokens,
total_cost_usd = aggregated.total_cost_usd,
total_response_time_ms = aggregated.total_response_time_ms,
last_used_at = aggregated.last_used_at
FROM aggregated
WHERE provider_api_keys.id = aggregated.provider_api_key_id
"#;
const LIST_USAGE_AUDITS_PREFIX: &str = r#"
SELECT
"usage".id,
@@ -3971,6 +4081,8 @@ WHERE "usage".created_at >= TO_TIMESTAMP($1::double precision)"#,
self.tx_runner
.run_read_write(|tx| {
Box::pin(async move {
lock_usage_request_id_in_tx(tx, &usage.request_id).await?;
if incoming_usage_can_recover_terminal_failure(
usage.status.as_str(),
usage.billing_status.as_str(),
@@ -3987,6 +4099,9 @@ WHERE "usage".created_at >= TO_TIMESTAMP($1::double precision)"#,
.map_postgres_err()?;
}
let previous_usage =
find_usage_by_request_id_in_tx(tx, &usage.request_id).await?;
let request_headers_json = json_bind_text(usage.request_headers.as_ref())?;
let request_body_storage =
prepare_usage_body_storage(usage.request_body.as_ref())?;
@@ -4292,11 +4407,66 @@ WHERE "usage".created_at >= TO_TIMESTAMP($1::double precision)"#,
routing_snapshot.local_execution_runtime_miss_reason.clone();
stored.output_price_per_1m = settlement_pricing_snapshot.output_price_per_1m;
stored.request_metadata = request_metadata_value;
let before_contribution = previous_usage
.as_ref()
.and_then(provider_api_key_usage_contribution);
let after_contribution = provider_api_key_usage_contribution(&stored);
match (before_contribution.as_ref(), after_contribution.as_ref()) {
(Some(before), Some(after)) if before.key_id == after.key_id => {
let delta = ProviderApiKeyUsageDelta::between(before, after);
apply_provider_api_key_usage_delta_in_tx(
tx,
before.key_id.as_str(),
&delta,
)
.await?;
}
_ => {
if let Some(before) = before_contribution.as_ref() {
let delta = ProviderApiKeyUsageDelta::removal(before);
apply_provider_api_key_usage_delta_in_tx(
tx,
before.key_id.as_str(),
&delta,
)
.await?;
}
if let Some(after) = after_contribution.as_ref() {
let delta = ProviderApiKeyUsageDelta::addition(after);
apply_provider_api_key_usage_delta_in_tx(
tx,
after.key_id.as_str(),
&delta,
)
.await?;
}
}
}
Ok(stored)
}) as BoxFuture<'_, Result<StoredRequestUsageAudit, DataLayerError>>
})
.await
}
pub async fn rebuild_provider_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
self.tx_runner
.run_read_write(|tx| {
Box::pin(async move {
sqlx::query(RESET_PROVIDER_API_KEY_USAGE_STATS_SQL)
.execute(&mut **tx)
.await
.map_postgres_err()?;
let rows_affected = sqlx::query(REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL)
.execute(&mut **tx)
.await
.map_postgres_err()?
.rows_affected();
Ok(rows_affected)
}) as BoxFuture<'_, Result<u64, DataLayerError>>
})
.await
}
}
#[async_trait]
@@ -4517,6 +4687,97 @@ impl UsageWriteRepository for SqlxUsageReadRepository {
) -> Result<StoredRequestUsageAudit, DataLayerError> {
Self::upsert(self, usage).await
}
async fn rebuild_provider_api_key_usage_stats(&self) -> Result<u64, DataLayerError> {
Self::rebuild_provider_api_key_usage_stats(self).await
}
}
async fn find_usage_by_request_id_in_tx(
tx: &mut sqlx::Transaction<'_, Postgres>,
request_id: &str,
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
sqlx::query(FIND_BY_REQUEST_ID_SQL)
.bind(request_id)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?
.map(|row| map_usage_row(&row, false))
.transpose()
}
async fn lock_usage_request_id_in_tx(
tx: &mut sqlx::Transaction<'_, Postgres>,
request_id: &str,
) -> Result<(), DataLayerError> {
sqlx::query(LOCK_USAGE_REQUEST_ID_SQL)
.bind(request_id)
.execute(&mut **tx)
.await
.map_postgres_err()?;
Ok(())
}
async fn apply_provider_api_key_usage_delta_in_tx(
tx: &mut sqlx::Transaction<'_, Postgres>,
key_id: &str,
delta: &ProviderApiKeyUsageDelta,
) -> Result<(), DataLayerError> {
if key_id.trim().is_empty() {
return Ok(());
}
if delta.is_noop() {
return Ok(());
}
let total_cost_usd_delta = if delta.total_cost_usd.is_finite() {
delta.total_cost_usd
} else {
0.0
};
sqlx::query(APPLY_PROVIDER_API_KEY_USAGE_DELTA_SQL)
.bind(key_id)
.bind(i32::try_from(delta.request_count).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"provider_api_keys.request_count delta exceeds i32: {}",
delta.request_count
))
})?)
.bind(i32::try_from(delta.success_count).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"provider_api_keys.success_count delta exceeds i32: {}",
delta.success_count
))
})?)
.bind(i32::try_from(delta.error_count).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"provider_api_keys.error_count delta exceeds i32: {}",
delta.error_count
))
})?)
.bind(delta.total_tokens)
.bind(total_cost_usd_delta)
.bind(i32::try_from(delta.total_response_time_ms).map_err(|_| {
DataLayerError::UnexpectedValue(format!(
"provider_api_keys.total_response_time_ms delta exceeds i32: {}",
delta.total_response_time_ms
))
})?)
.bind(
delta
.candidate_last_used_at_unix_secs
.map(|value| value as f64),
)
.bind(
delta
.removed_last_used_at_unix_secs
.map(|value| value as f64),
)
.execute(&mut **tx)
.await
.map_postgres_err()?;
Ok(())
}
// Build the usage read model from the split storage layout.
@@ -5848,6 +6109,26 @@ mod tests {
assert!(super::SUMMARIZE_USAGE_BY_PROVIDER_API_KEY_IDS_SQL.contains("ANY($1::TEXT[])"));
}
#[test]
fn usage_sql_serializes_request_id_upserts_before_reading_previous_usage() {
assert!(super::LOCK_USAGE_REQUEST_ID_SQL.contains("pg_advisory_xact_lock"));
assert!(super::LOCK_USAGE_REQUEST_ID_SQL.contains("hashtext($1)::BIGINT"));
assert!(include_str!("sql.rs")
.contains("lock_usage_request_id_in_tx(tx, &usage.request_id).await?;"));
}
#[test]
fn usage_sql_rebuild_matches_online_provider_key_usage_semantics() {
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("NULLIF(BTRIM(error_message), '') IS NULL"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL.contains("COALESCE("));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL.contains("total_tokens,"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("AND BTRIM(provider_api_key_id) <> ''"));
}
#[test]
fn usage_sql_supports_recent_usage_audits_query() {
assert!(super::LIST_RECENT_USAGE_AUDITS_PREFIX.contains("FROM \"usage\""));