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

Merge origin/main into main

Browse Source 
Integrate upstream updates while preserving the local analytics dashboards and schema-only migration changes.

Combine user account analysis with upstream user/group usage statistics in separate tabs, retain all migration versions, and keep the deleted audit document removed.

Validation: gateway all-target cargo check, frontend type check and 57 focused tests, 48 migration tests, schema composition checks, and diff whitespace checks.
This commit is contained in:
elky
2026-10-02 11:57:18 +08:00
parent 066ea87d72 54fbcc25a1
commit 625456ff24
343 changed files with 27929 additions and 2549 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/adapters/postgres/migrations/20260923000000_add_admin_wallet_batch_idempotency.sql
+12
View File
@@ -0,0 +1,12 @@
CREATE TABLE IF NOT EXISTS public.admin_user_wallet_balance_batches (
admin_user_id character varying(64) NOT NULL,
idempotency_key character varying(128) NOT NULL,
request_fingerprint character varying(64) NOT NULL,
target_user_ids jsonb NOT NULL,
missing_user_ids jsonb NOT NULL,
warnings jsonb NOT NULL,
user_outcomes jsonb NOT NULL,
created_at_unix_secs bigint NOT NULL,
updated_at_unix_secs bigint NOT NULL,
CONSTRAINT admin_user_wallet_balance_batches_pkey PRIMARY KEY (admin_user_id, idempotency_key)
);
crates/aether-data/adapters/postgres/src/candidate_selection.rs
+58 -10
View File
@@ -102,6 +102,11 @@ INNER JOIN LATERAL (
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'xai'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'bearer', 'api_key')
AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:image', 'openai:video')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
@@ -127,7 +132,8 @@ INNER JOIN LATERAL (
'vertex_ai',
'antigravity',
'kiro',
'windsurf'
'windsurf',
'xai'
)
AND LOWER(BTRIM(pak.auth_type)) <> 'oauth'
)
@@ -187,6 +193,11 @@ WHERE p.is_active = TRUE
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($3) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'xai'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'bearer', 'api_key')
AND LOWER($3) IN ('openai:responses', 'openai:responses:compact', 'openai:image', 'openai:video')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
@@ -212,7 +223,8 @@ WHERE p.is_active = TRUE
'vertex_ai',
'antigravity',
'kiro',
'windsurf'
'windsurf',
'xai'
)
AND LOWER(BTRIM(pak.auth_type)) <> 'oauth'
)
@@ -365,6 +377,11 @@ INNER JOIN LATERAL (
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'xai'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'bearer', 'api_key')
AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:image', 'openai:video')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
@@ -390,7 +407,8 @@ INNER JOIN LATERAL (
'vertex_ai',
'antigravity',
'kiro',
'windsurf'
'windsurf',
'xai'
)
AND LOWER(BTRIM(pak.auth_type)) <> 'oauth'
)
@@ -451,6 +469,11 @@ WHERE p.is_active = TRUE
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($4) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'xai'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'bearer', 'api_key')
AND LOWER($4) IN ('openai:responses', 'openai:responses:compact', 'openai:image', 'openai:video')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
@@ -476,7 +499,8 @@ WHERE p.is_active = TRUE
'vertex_ai',
'antigravity',
'kiro',
'windsurf'
'windsurf',
'xai'
)
AND LOWER(BTRIM(pak.auth_type)) <> 'oauth'
)
@@ -632,11 +656,16 @@ WHERE p.is_active = TRUE
)
)
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($6) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'grok'
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
AND LOWER($6) IN ('openai:chat', 'openai:responses', 'claude:messages', 'openai:image')
)
OR (
LOWER(BTRIM(p.provider_type)) = 'xai'
AND LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'bearer', 'api_key')
AND LOWER($6) IN ('openai:responses', 'openai:responses:compact', 'openai:image', 'openai:video')
)
OR (
LOWER(BTRIM(p.provider_type)) IN ('gemini_cli', 'antigravity')
AND LOWER(BTRIM(pak.auth_type)) = 'oauth'
@@ -662,7 +691,8 @@ WHERE p.is_active = TRUE
'vertex_ai',
'antigravity',
'kiro',
'windsurf'
'windsurf',
'xai'
)
AND LOWER(BTRIM(pak.auth_type)) <> 'oauth'
)
@@ -1717,6 +1747,24 @@ mod tests {
}
}
#[test]
fn candidate_selection_sql_allows_xai_oauth_responses_auth() {
let requested_model_sql = requested_model_selection_sql();
for sql in [
LIST_FOR_EXACT_API_FORMAT_SQL,
LIST_FOR_EXACT_API_FORMAT_AND_GLOBAL_MODEL_SQL,
LIST_POOL_KEYS_FOR_GROUP_SQL,
requested_model_sql.as_str(),
] {
assert!(sql.contains("LOWER(BTRIM(p.provider_type)) = 'xai'"));
assert!(sql.contains("LOWER(BTRIM(pak.auth_type)) IN ('oauth', 'bearer', 'api_key')"));
assert!(sql.contains(
"'openai:responses', 'openai:responses:compact', 'openai:image', 'openai:video'"
));
assert!(sql.contains("'xai'"));
}
}
#[test]
fn candidate_selection_sql_allows_windsurf_openai_chat_managed_keys() {
let requested_model_sql = requested_model_selection_sql();
crates/aether-data/adapters/postgres/src/usage/cleanup.rs
+20
View File
@@ -687,11 +687,31 @@ async fn cleanup_usage_raw_body_fields(
Ok(total_cleaned)
}
async fn truncate_usage_body_blobs_table(pool: &PostgresPool) -> Result<(), DataLayerError> {
let mut tx = pool.begin().await.map_err(postgres_error)?;
sqlx::query("SET LOCAL lock_timeout = '2s'")
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
sqlx::query("TRUNCATE TABLE usage_body_blobs")
.execute(&mut *tx)
.await
.map_err(postgres_error)?;
tx.commit().await.map_err(postgres_error)?;
Ok(())
}
async fn cleanup_usage_compressed_body_fields(
pool: &PostgresPool,
cutoff_time: DateTime<Utc>,
batch_size: usize,
) -> Result<usize, DataLayerError> {
if let Err(err) = truncate_usage_body_blobs_table(pool).await {
warn!(
error = %err,
"usage cleanup truncate usage_body_blobs table failed or timed out, falling back to batch deletion"
);
}
let mut total_cleaned = 0usize;
loop {
let rows = fetch_usage_body_cleanup_rows(
crates/aether-data/adapters/postgres/src/usage/mod.rs
+218 -170
View File
@@ -1925,6 +1925,33 @@ fn usage_leaderboard_sql_fragments(
}
}
fn push_usage_user_scope(
builder: &mut QueryBuilder<'_, Postgres>,
column: &str,
user_id: Option<&str>,
user_ids: Option<&[String]>,
) {
if let Some(user_id) = user_id {
builder
.push(" AND ")
.push(column)
.push(" = ")
.push_bind(user_id.to_string());
}
if let Some(user_ids) = user_ids {
if user_ids.is_empty() {
builder.push(" AND FALSE");
} else {
builder.push(" AND ").push(column).push(" IN (");
let mut separated = builder.separated(", ");
for user_id in user_ids {
separated.push_bind(user_id.clone());
}
separated.push_unseparated(")");
}
}
}
const LIST_RECENT_USAGE_AUDITS_PREFIX: &str =
include_str!("queries/list_recent_usage_audits_prefix.sql");
@@ -3762,14 +3789,15 @@ OR (\"usage\".error_message IS NOT NULL AND BTRIM(\"usage\".error_message) <> ''
start_day_utc: DateTime<Utc>,
end_day_utc: DateTime<Utc>,
user_id: Option<&str>,
user_ids: Option<&[String]>,
) -> Result<StoredUsageAuditSummary, DataLayerError> {
if start_day_utc >= end_day_utc {
return Ok(StoredUsageAuditSummary::default());
}
let row = if let Some(user_id) = user_id {
sqlx::query(
r#"
let scoped_to_users = user_id.is_some() || user_ids.is_some();
let mut builder = QueryBuilder::<Postgres>::new(
r#"
SELECT
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
COALESCE(SUM(input_tokens), 0)::BIGINT AS input_tokens,
@@ -3794,68 +3822,41 @@ SELECT
COALESCE(SUM(cache_read_cost), 0)::DOUBLE PRECISION AS cache_read_cost_usd,
COALESCE(SUM(response_time_sum_ms), 0)::DOUBLE PRECISION AS total_response_time_ms,
COALESCE(SUM(error_requests), 0)::BIGINT AS error_requests
FROM stats_user_daily
WHERE user_id = $1
AND date >= $2
AND date < $3
"#,
)
.bind(user_id)
.bind(start_day_utc)
.bind(end_day_utc)
.fetch_one(&self.pool)
.await
.map_postgres_err()?
FROM "#,
);
builder.push(if scoped_to_users {
"stats_user_daily"
} else {
sqlx::query(
r#"
SELECT
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
COALESCE(SUM(input_tokens), 0)::BIGINT AS input_tokens,
COALESCE(SUM(output_tokens), 0)::BIGINT AS output_tokens,
COALESCE(SUM(
CASE
WHEN effective_input_tokens = 0 AND total_input_context = 0 AND input_tokens > 0
THEN input_tokens
ELSE effective_input_tokens
END
+ output_tokens + cache_creation_tokens + cache_read_tokens
), 0)::BIGINT AS recorded_total_tokens,
COALESCE(SUM(cache_creation_tokens), 0)::BIGINT AS cache_creation_tokens,
COALESCE(SUM(cache_creation_ephemeral_5m_tokens), 0)::BIGINT
AS cache_creation_ephemeral_5m_tokens,
COALESCE(SUM(cache_creation_ephemeral_1h_tokens), 0)::BIGINT
AS cache_creation_ephemeral_1h_tokens,
COALESCE(SUM(cache_read_tokens), 0)::BIGINT AS cache_read_tokens,
COALESCE(SUM(total_cost), 0)::DOUBLE PRECISION AS total_cost_usd,
COALESCE(SUM(actual_total_cost), 0)::DOUBLE PRECISION AS actual_total_cost_usd,
COALESCE(SUM(cache_creation_cost), 0)::DOUBLE PRECISION AS cache_creation_cost_usd,
COALESCE(SUM(cache_read_cost), 0)::DOUBLE PRECISION AS cache_read_cost_usd,
COALESCE(SUM(response_time_sum_ms), 0)::DOUBLE PRECISION AS total_response_time_ms,
COALESCE(SUM(error_requests), 0)::BIGINT AS error_requests
FROM stats_daily
WHERE date >= $1
AND date < $2
"#,
)
.bind(start_day_utc)
.bind(end_day_utc)
"stats_daily"
});
builder
.push(" WHERE date >= ")
.push_bind(start_day_utc)
.push(" AND date < ")
.push_bind(end_day_utc);
if scoped_to_users {
push_usage_user_scope(&mut builder, "user_id", user_id, user_ids);
}
let row = builder
.build()
.fetch_one(&self.pool)
.await
.map_postgres_err()?
};
.map_postgres_err()?;
decode_usage_audit_summary_row(&row)
}
async fn summarize_usage_audits_raw(
&self,
created_from_unix_secs: u64,
created_until_unix_secs: u64,
user_id: Option<&str>,
provider_name: Option<&str>,
model: Option<&str>,
query: &UsageAuditSummaryQuery,
) -> Result<StoredUsageAuditSummary, DataLayerError> {
let created_from_unix_secs = query.created_from_unix_secs;
let created_until_unix_secs = query.created_until_unix_secs;
let user_id = query.user_id.as_deref();
let user_ids = query.user_ids.as_deref();
let provider_names = query.provider_names.as_deref();
let provider_name = query.provider_name.as_deref();
let model = query.model.as_deref();
if created_from_unix_secs >= created_until_unix_secs {
return Ok(StoredUsageAuditSummary::default());
}
@@ -3907,11 +3908,12 @@ FROM usage_billing_facts AS "usage"
.push("\"usage\".created_at < TO_TIMESTAMP(")
.push_bind(created_until_unix_secs as f64)
.push("::double precision)");
if let Some(user_id) = user_id {
builder.push(if has_where { " AND " } else { " WHERE " });
push_usage_user_scope(&mut builder, "\"usage\".user_id", user_id, user_ids);
if let Some(names) = provider_names {
builder
.push("\"usage\".user_id = ")
.push_bind(user_id.to_string());
.push(" AND \"usage\".provider_name = ANY(")
.push_bind(names.to_vec())
.push("::text[])");
}
if let Some(provider_name) = provider_name {
builder.push(if has_where { " AND " } else { " WHERE " });
@@ -3939,55 +3941,32 @@ FROM usage_billing_facts AS "usage"
&self,
query: &UsageAuditSummaryQuery,
) -> Result<StoredUsageAuditSummary, DataLayerError> {
if query.provider_name.is_some() || query.model.is_some() {
return self
.summarize_usage_audits_raw(
query.created_from_unix_secs,
query.created_until_unix_secs,
query.user_id.as_deref(),
query.provider_name.as_deref(),
query.model.as_deref(),
)
.await;
// Provider rollups lack cache-cost/error detail required by this response.
// Keep these scopes on canonical facts rather than inventing missing daily totals.
if query.provider_names.is_some() || query.provider_name.is_some() || query.model.is_some()
{
return self.summarize_usage_audits_raw(query).await;
}
let Some(cutoff_utc) = self.read_stats_daily_cutoff_date().await? else {
return self
.summarize_usage_audits_raw(
query.created_from_unix_secs,
query.created_until_unix_secs,
query.user_id.as_deref(),
None,
None,
)
.await;
return self.summarize_usage_audits_raw(query).await;
};
let start_utc = dashboard_unix_secs_to_utc(query.created_from_unix_secs);
let end_utc = dashboard_unix_secs_to_utc(query.created_until_unix_secs);
let split = split_dashboard_daily_aggregate_range(start_utc, end_utc, cutoff_utc);
let Some(_) = split.aggregate else {
return self
.summarize_usage_audits_raw(
query.created_from_unix_secs,
query.created_until_unix_secs,
query.user_id.as_deref(),
None,
None,
)
.await;
return self.summarize_usage_audits_raw(query).await;
};
let mut summary = StoredUsageAuditSummary::default();
if let Some((raw_start, raw_end)) = split.raw_leading {
absorb_usage_audit_summary(
&mut summary,
self.summarize_usage_audits_raw(
dashboard_utc_to_unix_secs(raw_start),
dashboard_utc_to_unix_secs(raw_end),
query.user_id.as_deref(),
None,
None,
)
self.summarize_usage_audits_raw(&UsageAuditSummaryQuery {
created_from_unix_secs: dashboard_utc_to_unix_secs(raw_start),
created_until_unix_secs: dashboard_utc_to_unix_secs(raw_end),
..query.clone()
})
.await?,
);
}
@@ -3998,6 +3977,7 @@ FROM usage_billing_facts AS "usage"
aggregate_start,
aggregate_end,
query.user_id.as_deref(),
query.user_ids.as_deref(),
)
.await?,
);
@@ -4005,13 +3985,11 @@ FROM usage_billing_facts AS "usage"
if let Some((raw_start, raw_end)) = split.raw_trailing {
absorb_usage_audit_summary(
&mut summary,
self.summarize_usage_audits_raw(
dashboard_utc_to_unix_secs(raw_start),
dashboard_utc_to_unix_secs(raw_end),
query.user_id.as_deref(),
None,
None,
)
self.summarize_usage_audits_raw(&UsageAuditSummaryQuery {
created_from_unix_secs: dashboard_utc_to_unix_secs(raw_start),
created_until_unix_secs: dashboard_utc_to_unix_secs(raw_end),
..query.clone()
})
.await?,
);
}
@@ -6952,12 +6930,17 @@ FROM usage_billing_facts AS "usage"
.push("\"usage\".created_at < TO_TIMESTAMP(")
.push_bind(query.created_until_unix_secs as f64)
.push("::double precision)");
if let Some(user_id) = query.user_id.as_deref() {
builder.push(if has_where { " AND " } else { " WHERE " });
has_where = true;
push_usage_user_scope(
&mut builder,
"\"usage\".user_id",
query.user_id.as_deref(),
query.user_ids.as_deref(),
);
if let Some(names) = query.provider_names.as_ref() {
builder
.push("\"usage\".user_id = ")
.push_bind(user_id.to_string());
.push(" AND \"usage\".provider_name = ANY(")
.push_bind(names.clone())
.push("::text[])");
}
if let Some(provider_name) = query.provider_name.as_deref() {
builder.push(if has_where { " AND " } else { " WHERE " });
@@ -6998,62 +6981,54 @@ FROM usage_billing_facts AS "usage"
start_day_utc: DateTime<Utc>,
end_day_utc: DateTime<Utc>,
user_id: Option<&str>,
user_ids: Option<&[String]>,
provider_names: Option<&[String]>,
) -> Result<Vec<StoredUsageTimeSeriesBucket>, DataLayerError> {
if start_day_utc >= end_day_utc {
return Ok(Vec::new());
}
let rows = if let Some(user_id) = user_id {
sqlx::query(
r#"
let scoped_to_users = user_id.is_some() || user_ids.is_some();
let mut builder = QueryBuilder::<Postgres>::new(
r#"
SELECT
TO_CHAR(date, 'YYYY-MM-DD') AS bucket_key,
total_requests::BIGINT AS total_requests,
input_tokens::BIGINT AS input_tokens,
output_tokens::BIGINT AS output_tokens,
cache_creation_tokens::BIGINT AS cache_creation_tokens,
cache_read_tokens::BIGINT AS cache_read_tokens,
CAST(total_cost AS DOUBLE PRECISION) AS total_cost_usd,
CAST(response_time_sum_ms AS DOUBLE PRECISION) AS total_response_time_ms
FROM stats_user_daily
WHERE user_id = $1
AND date >= $2
AND date < $3
ORDER BY date ASC
"#,
)
.bind(user_id)
.bind(start_day_utc)
.bind(end_day_utc)
.fetch_all(&self.pool)
.await
.map_postgres_err()?
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
COALESCE(SUM(input_tokens), 0)::BIGINT AS input_tokens,
COALESCE(SUM(output_tokens), 0)::BIGINT AS output_tokens,
COALESCE(SUM(cache_creation_tokens), 0)::BIGINT AS cache_creation_tokens,
COALESCE(SUM(cache_read_tokens), 0)::BIGINT AS cache_read_tokens,
COALESCE(SUM(CAST(total_cost AS DOUBLE PRECISION)), 0) AS total_cost_usd,
COALESCE(SUM(CAST(response_time_sum_ms AS DOUBLE PRECISION)), 0)
AS total_response_time_ms
FROM "#,
);
builder.push(if provider_names.is_some() {
"stats_user_daily_provider"
} else if scoped_to_users {
"stats_user_daily"
} else {
sqlx::query(
r#"
SELECT
TO_CHAR(date, 'YYYY-MM-DD') AS bucket_key,
total_requests::BIGINT AS total_requests,
input_tokens::BIGINT AS input_tokens,
output_tokens::BIGINT AS output_tokens,
cache_creation_tokens::BIGINT AS cache_creation_tokens,
cache_read_tokens::BIGINT AS cache_read_tokens,
CAST(total_cost AS DOUBLE PRECISION) AS total_cost_usd,
CAST(response_time_sum_ms AS DOUBLE PRECISION) AS total_response_time_ms
FROM stats_daily
WHERE date >= $1
AND date < $2
ORDER BY date ASC
"#,
)
.bind(start_day_utc)
.bind(end_day_utc)
.fetch_all(&self.pool)
.await
.map_postgres_err()?
};
"stats_daily"
});
builder
.push(" WHERE date >= ")
.push_bind(start_day_utc)
.push(" AND date < ")
.push_bind(end_day_utc);
if scoped_to_users {
push_usage_user_scope(&mut builder, "user_id", user_id, user_ids);
}
if let Some(names) = provider_names {
builder
.push(" AND provider_name = ANY(")
.push_bind(names.to_vec())
.push("::text[])");
}
builder.push(" GROUP BY date ORDER BY date ASC");
let mut rows = builder.build().fetch(&self.pool);
let mut items = Vec::new();
for row in rows {
while let Some(row) = rows.try_next().await.map_postgres_err()? {
items.push(decode_usage_time_series_bucket_row(&row)?);
}
Ok(items)
@@ -7134,11 +7109,13 @@ WHERE is_complete IS TRUE
&mut grouped,
self.summarize_usage_time_series_raw(
&UsageTimeSeriesQuery {
provider_names: query.provider_names.clone(),
created_from_unix_secs: dashboard_utc_to_unix_secs(raw_start),
created_until_unix_secs: dashboard_utc_to_unix_secs(raw_end),
granularity: UsageTimeSeriesGranularity::Day,
tz_offset_minutes: 0,
user_id: query.user_id.clone(),
user_ids: query.user_ids.clone(),
provider_name: None,
model: None,
},
@@ -7153,6 +7130,8 @@ WHERE is_complete IS TRUE
aggregate_start,
aggregate_end,
query.user_id.as_deref(),
query.user_ids.as_deref(),
query.provider_names.as_deref(),
)
.await?,
);
@@ -7160,11 +7139,13 @@ WHERE is_complete IS TRUE
&mut grouped,
self.summarize_usage_time_series_raw(
&UsageTimeSeriesQuery {
provider_names: query.provider_names.clone(),
created_from_unix_secs: dashboard_utc_to_unix_secs(aggregate_start),
created_until_unix_secs: dashboard_utc_to_unix_secs(aggregate_end),
granularity: UsageTimeSeriesGranularity::Day,
tz_offset_minutes: 0,
user_id: query.user_id.clone(),
user_ids: query.user_ids.clone(),
provider_name: None,
model: None,
},
@@ -7177,11 +7158,13 @@ WHERE is_complete IS TRUE
&mut grouped,
self.summarize_usage_time_series_raw(
&UsageTimeSeriesQuery {
provider_names: query.provider_names.clone(),
created_from_unix_secs: dashboard_utc_to_unix_secs(raw_start),
created_until_unix_secs: dashboard_utc_to_unix_secs(raw_end),
granularity: UsageTimeSeriesGranularity::Day,
tz_offset_minutes: 0,
user_id: query.user_id.clone(),
user_ids: query.user_ids.clone(),
provider_name: None,
model: None,
},
@@ -7195,7 +7178,11 @@ WHERE is_complete IS TRUE
}
}
if query.user_id.is_none() && query.tz_offset_minutes % 60 == 0 {
if query.provider_names.is_none()
&& query.user_id.is_none()
&& query.user_ids.is_none()
&& query.tz_offset_minutes % 60 == 0
{
if let Some(cutoff_utc) = self.read_stats_hourly_cutoff().await? {
let start_utc = dashboard_unix_secs_to_utc(query.created_from_unix_secs);
let end_utc = dashboard_unix_secs_to_utc(query.created_until_unix_secs);
@@ -7207,11 +7194,13 @@ WHERE is_complete IS TRUE
&mut grouped,
self.summarize_usage_time_series_raw(
&UsageTimeSeriesQuery {
provider_names: query.provider_names.clone(),
created_from_unix_secs: dashboard_utc_to_unix_secs(raw_start),
created_until_unix_secs: dashboard_utc_to_unix_secs(raw_end),
granularity: query.granularity,
tz_offset_minutes: query.tz_offset_minutes,
user_id: None,
user_ids: None,
provider_name: None,
model: None,
},
@@ -7234,11 +7223,13 @@ WHERE is_complete IS TRUE
&mut grouped,
self.summarize_usage_time_series_raw(
&UsageTimeSeriesQuery {
provider_names: query.provider_names.clone(),
created_from_unix_secs: dashboard_utc_to_unix_secs(aggregate_start),
created_until_unix_secs: dashboard_utc_to_unix_secs(aggregate_end),
granularity: query.granularity,
tz_offset_minutes: query.tz_offset_minutes,
user_id: None,
user_ids: None,
provider_name: None,
model: None,
},
@@ -7251,11 +7242,13 @@ WHERE is_complete IS TRUE
&mut grouped,
self.summarize_usage_time_series_raw(
&UsageTimeSeriesQuery {
provider_names: query.provider_names.clone(),
created_from_unix_secs: dashboard_utc_to_unix_secs(raw_start),
created_until_unix_secs: dashboard_utc_to_unix_secs(raw_end),
granularity: query.granularity,
tz_offset_minutes: query.tz_offset_minutes,
user_id: None,
user_ids: None,
provider_name: None,
model: None,
},
@@ -7295,6 +7288,8 @@ WHERE "usage".created_at >= TO_TIMESTAMP($1::double precision)
AND ($3::varchar IS NULL OR "usage".user_id = $3)
AND ($4::varchar IS NULL OR "usage".provider_name = $4)
AND ($5::varchar IS NULL OR "usage".model = $5)
AND ($6::text[] IS NULL OR "usage".user_id::text = ANY($6))
AND ($7::text[] IS NULL OR "usage".provider_name = ANY($7))
GROUP BY group_key
ORDER BY group_key ASC
"#,
@@ -7308,6 +7303,8 @@ ORDER BY group_key ASC
.bind(query.user_id.as_deref())
.bind(query.provider_name.as_deref())
.bind(query.model.as_deref())
.bind(query.user_ids.clone())
.bind(query.provider_names.clone())
.fetch(&self.pool);
let mut items = Vec::new();
while let Some(row) = rows.try_next().await.map_postgres_err()? {
@@ -7322,6 +7319,51 @@ ORDER BY group_key ASC
end_day_utc: DateTime<Utc>,
query: &UsageLeaderboardQuery,
) -> Result<Option<Vec<StoredUsageLeaderboardSummary>>, DataLayerError> {
if let Some(names) = query.provider_names.as_ref() {
let group_key = match query.group_by {
UsageLeaderboardGroupBy::User => "user_id",
UsageLeaderboardGroupBy::Model => "model",
UsageLeaderboardGroupBy::ApiKey => return Ok(None),
};
let mut builder = QueryBuilder::<Postgres>::new(format!(
r#"
SELECT
{group_key} AS group_key,
MAX(NULLIF(BTRIM(username), '')) AS legacy_name,
COALESCE(SUM(total_requests), 0)::BIGINT AS request_count,
COALESCE(SUM(total_tokens), 0)::BIGINT AS total_tokens,
COALESCE(SUM(total_cost), 0)::DOUBLE PRECISION AS total_cost_usd
FROM stats_user_daily_model_provider
WHERE date >=
"#
));
builder
.push_bind(start_day_utc)
.push(" AND date < ")
.push_bind(end_day_utc);
builder
.push(" AND provider_name = ANY(")
.push_bind(names.clone())
.push("::text[])");
push_usage_user_scope(
&mut builder,
"user_id",
query.user_id.as_deref(),
query.user_ids.as_deref(),
);
if let Some(provider) = query.provider_name.as_ref() {
builder
.push(" AND provider_name = ")
.push_bind(provider.clone());
}
if let Some(model) = query.model.as_ref() {
builder.push(" AND model = ").push_bind(model.clone());
}
builder.push(format!(" GROUP BY {group_key} ORDER BY {group_key}"));
return fetch_usage_leaderboard_query(builder.build(), &self.pool)
.await
.map(Some);
}
let items = match query.group_by {
UsageLeaderboardGroupBy::Model => {
let mut builder = if let Some(user_id) = query.user_id.as_deref() {
@@ -7451,11 +7493,12 @@ WHERE date >=
.push_bind(end_day_utc)
.push(" AND provider_name = ")
.push_bind(provider_name.to_string());
if let Some(user_id) = query.user_id.as_deref() {
builder
.push(" AND user_id = ")
.push_bind(user_id.to_string());
}
push_usage_user_scope(
&mut builder,
"user_id",
query.user_id.as_deref(),
query.user_ids.as_deref(),
);
builder.push(" GROUP BY user_id ORDER BY user_id ASC");
builder
} else if let Some(model) = query.model.as_deref() {
@@ -7477,11 +7520,12 @@ WHERE date >=
.push_bind(end_day_utc)
.push(" AND model = ")
.push_bind(model.to_string());
if let Some(user_id) = query.user_id.as_deref() {
builder
.push(" AND user_id = ")
.push_bind(user_id.to_string());
}
push_usage_user_scope(
&mut builder,
"user_id",
query.user_id.as_deref(),
query.user_ids.as_deref(),
);
builder.push(" GROUP BY user_id ORDER BY user_id ASC");
builder
} else {
@@ -7505,11 +7549,12 @@ WHERE date >=
.push(" AND date < ")
.push_bind(end_day_utc)
.push(" AND user_id IS NOT NULL");
if let Some(user_id) = query.user_id.as_deref() {
builder
.push(" AND user_id = ")
.push_bind(user_id.to_string());
}
push_usage_user_scope(
&mut builder,
"user_id",
query.user_id.as_deref(),
query.user_ids.as_deref(),
);
builder.push(" GROUP BY user_id ORDER BY user_id ASC");
builder
};
@@ -7594,6 +7639,7 @@ WHERE stats_daily_api_key.date >=
absorb_usage_leaderboard_rows(
&mut grouped,
self.summarize_usage_leaderboard_raw(&UsageLeaderboardQuery {
provider_names: query.provider_names.clone(),
created_from_unix_secs: dashboard_utc_to_unix_secs(raw_start),
created_until_unix_secs: dashboard_utc_to_unix_secs(raw_end),
..query.clone()
@@ -7616,6 +7662,7 @@ WHERE stats_daily_api_key.date >=
absorb_usage_leaderboard_rows(
&mut grouped,
self.summarize_usage_leaderboard_raw(&UsageLeaderboardQuery {
provider_names: query.provider_names.clone(),
created_from_unix_secs: dashboard_utc_to_unix_secs(aggregate_start),
created_until_unix_secs: dashboard_utc_to_unix_secs(aggregate_end),
..query.clone()
@@ -7629,6 +7676,7 @@ WHERE stats_daily_api_key.date >=
absorb_usage_leaderboard_rows(
&mut grouped,
self.summarize_usage_leaderboard_raw(&UsageLeaderboardQuery {
provider_names: query.provider_names.clone(),
created_from_unix_secs: dashboard_utc_to_unix_secs(raw_start),
created_until_unix_secs: dashboard_utc_to_unix_secs(raw_end),
..query.clone()
crates/aether-data/adapters/postgres/src/usage/queries/find_by_id_sql.sql
+4
View File
@@ -185,6 +185,10 @@ SELECT
usage_http_audits.provider_request_body_ref AS http_provider_request_body_ref,
usage_http_audits.response_body_ref AS http_response_body_ref,
usage_http_audits.client_response_body_ref AS http_client_response_body_ref,
usage_http_audits.request_body_state AS http_request_body_state,
usage_http_audits.provider_request_body_state AS http_provider_request_body_state,
usage_http_audits.response_body_state AS http_response_body_state,
usage_http_audits.client_response_body_state AS http_client_response_body_state,
usage_routing_snapshots.candidate_id AS routing_candidate_id,
usage_routing_snapshots.candidate_index AS routing_candidate_index,
usage_routing_snapshots.key_name AS routing_key_name,
crates/aether-data/adapters/postgres/src/usage/queries/list_recent_usage_audits_prefix.sql
+3
View File
@@ -183,6 +183,7 @@ SELECT
OR NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_actual_service_tier'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_response_model'), '') IS NOT NULL
OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false')
OR ("usage".request_metadata->>'websocket_mode') IN ('true', 'false')
@@ -208,6 +209,8 @@ SELECT
NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''),
'provider_actual_service_tier',
NULLIF(BTRIM("usage".request_metadata->>'provider_actual_service_tier'), ''),
'provider_response_model',
NULLIF(BTRIM("usage".request_metadata->>'provider_response_model'), ''),
'client_requested_stream',
CASE
WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
crates/aether-data/adapters/postgres/src/usage/queries/list_usage_audits_prefix.sql
+3
View File
@@ -183,6 +183,7 @@ SELECT
OR NULLIF(BTRIM("usage".request_metadata->>'provider_reasoning_effort'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_actual_service_tier'), '') IS NOT NULL
OR NULLIF(BTRIM("usage".request_metadata->>'provider_response_model'), '') IS NOT NULL
OR ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
OR ("usage".request_metadata->>'upstream_is_stream') IN ('true', 'false')
OR ("usage".request_metadata->>'websocket_mode') IN ('true', 'false')
@@ -208,6 +209,8 @@ SELECT
NULLIF(BTRIM("usage".request_metadata->>'provider_service_tier'), ''),
'provider_actual_service_tier',
NULLIF(BTRIM("usage".request_metadata->>'provider_actual_service_tier'), ''),
'provider_response_model',
NULLIF(BTRIM("usage".request_metadata->>'provider_response_model'), ''),
'client_requested_stream',
CASE
WHEN ("usage".request_metadata->>'client_requested_stream') IN ('true', 'false')
crates/aether-data/adapters/postgres/src/usage/tests.rs
+18 -2
View File
@@ -3251,13 +3251,15 @@ fn usage_sql_canonical_openai_cache_case_preserves_effective_and_total_tokens()
aggregate_audit_summary
.matches("WHEN effective_input_tokens = 0 AND total_input_context = 0")
.count(),
2
1,
"the shared daily aggregate query should define the legacy token fallback once"
);
assert_eq!(
aggregate_audit_summary
.matches("+ output_tokens + cache_creation_tokens + cache_read_tokens")
.count(),
2
1,
"the shared daily aggregate query should define canonical total tokens once"
);
assert!(!aggregate_audit_summary.contains("SUM(input_tokens + output_tokens)"));
@@ -3467,6 +3469,18 @@ fn usage_sql_reads_http_audits_for_single_record_fetches() {
assert!(super::FIND_BY_ID_SQL.contains("LEFT JOIN usage_http_audits"));
assert!(super::FIND_BY_REQUEST_ID_SQL.contains("http_request_body_ref"));
assert!(super::FIND_BY_ID_SQL.contains("http_client_response_body_ref"));
for sql in [super::FIND_BY_REQUEST_ID_SQL, super::FIND_BY_ID_SQL] {
for field in [
"request_body",
"provider_request_body",
"response_body",
"client_response_body",
] {
assert!(sql.contains(&format!(
"usage_http_audits.{field}_state AS http_{field}_state"
)));
}
}
}
#[test]
@@ -3597,6 +3611,8 @@ fn usage_sql_uses_json_null_placeholders_for_usage_payload_columns() {
assert!(sql.contains("request_metadata->>'provider_reasoning_effort'"));
assert!(sql.contains("request_metadata->>'provider_service_tier'"));
assert!(sql.contains("request_metadata->>'provider_actual_service_tier'"));
assert!(sql.contains("request_metadata->>'provider_response_model'"));
assert!(sql.contains("'provider_response_model'"));
assert!(sql.contains("request_metadata->>'websocket_mode'"));
assert!(sql.contains("'websocket_mode'"));
assert!(sql.contains("AS client_family"));
crates/aether-data/adapters/postgres/src/wallet.rs
+653 -26
View File
@@ -24,8 +24,9 @@ use aether_data_contracts::repository::wallet::{
wallet_recharge_checkout_failed_response, wallet_recharge_checkout_uncertain_response,
wallet_recharge_order_is_reclaimable_placeholder, wallet_recharge_replay_matches,
wallet_recharge_response_is_checkout_placeholder, wallet_refund_proof_is_success,
AdjustWalletBalanceInput, AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery,
AdminRedeemCodeListQuery, AdminWalletLedgerQuery, AdminWalletListQuery,
AdjustWalletBalanceInBatchInput, AdjustWalletBalanceInput, AdminPaymentOrderListQuery,
AdminRedeemCodeBatchListQuery, AdminRedeemCodeListQuery, AdminUserWalletBalanceBatchContext,
AdminUserWalletBalanceBatchUserOutcome, AdminWalletLedgerQuery, AdminWalletListQuery,
AdminWalletRefundRequestListQuery, CompareAndSwapPaymentOrderStripeClientSecretInput,
CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput,
CreateAdminRedeemCodeBatchResult, CreateManualWalletRechargeInput,
@@ -34,21 +35,23 @@ use aether_data_contracts::repository::wallet::{
CreateWalletRefundRequestOutcome, CreatedAdminRedeemCodePlaintext,
CreditAdminPaymentOrderInput, DeleteAdminRedeemCodeBatchInput,
DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput, FailAdminWalletRefundInput,
FailWalletRechargeCheckoutInput, InitializeAuthWalletOutcome, ProcessAdminWalletRefundInput,
ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome, ReclaimWalletRechargeCheckoutInput,
RedeemWalletCodeInput, RedeemWalletCodeOutcome, StoredAdminPaymentCallback,
StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder, StoredAdminPaymentOrderPage,
StoredAdminRedeemCode, StoredAdminRedeemCodeBatch, StoredAdminRedeemCodeBatchPage,
StoredAdminRedeemCodePage, StoredAdminWalletLedgerItem, StoredAdminWalletLedgerPage,
StoredAdminWalletListItem, StoredAdminWalletListPage, StoredAdminWalletRefund,
StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestItem,
StoredAdminWalletRefundRequestPage, StoredAdminWalletTransaction,
StoredAdminWalletTransactionPage, StoredWalletDailyUsageLedger,
FailWalletRechargeCheckoutInput, InitializeAuthWalletOutcome,
PrepareAdminUserWalletBalanceBatchInput, PrepareAdminUserWalletBalanceBatchOutcome,
ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome,
ReclaimWalletRechargeCheckoutInput, RedeemWalletCodeInput, RedeemWalletCodeOutcome,
StoredAdminPaymentCallback, StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder,
StoredAdminPaymentOrderPage, StoredAdminRedeemCode, StoredAdminRedeemCodeBatch,
StoredAdminRedeemCodeBatchPage, StoredAdminRedeemCodePage, StoredAdminUserWalletBalanceBatch,
StoredAdminWalletLedgerItem, StoredAdminWalletLedgerPage, StoredAdminWalletListItem,
StoredAdminWalletListPage, StoredAdminWalletRefund, StoredAdminWalletRefundPage,
StoredAdminWalletRefundRequestItem, StoredAdminWalletRefundRequestPage,
StoredAdminWalletTransaction, StoredAdminWalletTransactionPage, StoredWalletDailyUsageLedger,
StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, UpdateAdminWalletRefundGatewayInput,
UpdateWalletRechargeCheckoutInput, WalletLookupKey, WalletMutationOutcome,
WalletReadRepository, WalletWriteRepository,
};
use aether_data_contracts::DataLayerError;
use std::collections::BTreeMap;
use crate::{
error::{postgres_error, SqlxResultExt},
@@ -809,6 +812,44 @@ impl SqlxWalletRepository {
}
}
fn effective_wallet_adjustment_amount(input: &AdjustWalletBalanceInput, before_total: f64) -> f64 {
if input.clamp_deduction_to_available_balance && input.amount_usd < 0.0 {
if before_total < 0.0 {
// Clear legacy negative totals to zero and record the actual ledger delta.
-before_total
} else {
-(-input.amount_usd).min(before_total)
}
} else {
input.amount_usd
}
}
async fn persist_admin_wallet_batch_outcomes(
connection: &mut sqlx::PgConnection,
context: &AdminUserWalletBalanceBatchContext,
outcomes: &BTreeMap<String, AdminUserWalletBalanceBatchUserOutcome>,
) -> Result<(), DataLayerError> {
let outcomes = serde_json::to_value(outcomes).map_err(|error| {
DataLayerError::UnexpectedValue(format!("admin wallet batch outcomes are invalid: {error}"))
})?;
sqlx::query(
r#"
UPDATE admin_user_wallet_balance_batches
SET user_outcomes = $3, updated_at_unix_secs = $4
WHERE admin_user_id = $1 AND idempotency_key = $2
"#,
)
.bind(&context.admin_user_id)
.bind(&context.idempotency_key)
.bind(outcomes)
.bind(Utc::now().timestamp().max(0))
.execute(connection)
.await
.map_postgres_err()?;
Ok(())
}
#[async_trait]
impl WalletReadRepository for SqlxWalletRepository {
async fn find(
@@ -4266,7 +4307,8 @@ RETURNING
async fn adjust_wallet_balance(
&self,
input: AdjustWalletBalanceInput,
) -> Result<Option<(StoredWalletSnapshot, StoredAdminWalletTransaction)>, DataLayerError> {
) -> Result<Option<(StoredWalletSnapshot, Option<StoredAdminWalletTransaction>)>, DataLayerError>
{
if !input.amount_usd.is_finite() || input.amount_usd == 0.0 {
return Err(DataLayerError::InvalidInput(
"adjustment amount must be finite and non-zero".to_string(),
@@ -4275,6 +4317,58 @@ RETURNING
self.tx_runner
.run_read_write(|tx| {
Box::pin(async move {
let batch_context = input.batch_context.clone();
let mut batch_user_outcomes = if let Some(context) = &batch_context {
let batch_row = sqlx::query(
r#"
SELECT target_user_ids, user_outcomes
FROM admin_user_wallet_balance_batches
WHERE admin_user_id = $1 AND idempotency_key = $2
FOR UPDATE
"#,
)
.bind(&context.admin_user_id)
.bind(&context.idempotency_key)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?
.ok_or_else(|| {
DataLayerError::InvalidInput(
"admin wallet batch was not prepared".to_string(),
)
})?;
let target_user_ids: Vec<String> =
serde_json::from_value(row_get(&batch_row, "target_user_ids")?)
.map_err(|error| {
DataLayerError::UnexpectedValue(format!(
"admin wallet batch target list is invalid: {error}"
))
})?;
if !target_user_ids.iter().any(|id| id == &context.user_id) {
return Err(DataLayerError::InvalidInput(
"user is outside the prepared admin wallet batch".to_string(),
));
}
let outcomes: BTreeMap<String, AdminUserWalletBalanceBatchUserOutcome> =
serde_json::from_value(row_get(&batch_row, "user_outcomes")?).map_err(
|error| {
DataLayerError::UnexpectedValue(format!(
"admin wallet batch outcomes are invalid: {error}"
))
},
)?;
if let Some(outcome) = outcomes.get(&context.user_id) {
match outcome {
AdminUserWalletBalanceBatchUserOutcome::Succeeded => {}
AdminUserWalletBalanceBatchUserOutcome::Failed(_) => {
return Ok(None);
}
}
}
Some(outcomes)
} else {
None
};
let Some(row) = sqlx::query(
r#"
SELECT
@@ -4289,13 +4383,20 @@ SELECT
CAST(total_recharged AS DOUBLE PRECISION) AS total_recharged,
CAST(total_consumed AS DOUBLE PRECISION) AS total_consumed,
CAST(total_refunded AS DOUBLE PRECISION) AS total_refunded,
CAST(total_adjusted AS DOUBLE PRECISION) AS total_adjusted
CAST(total_adjusted AS DOUBLE PRECISION) AS total_adjusted,
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs
FROM wallets
WHERE id = $1
AND ($2::character varying IS NULL OR user_id = $2::character varying)
FOR UPDATE
"#,
)
.bind(&input.wallet_id)
.bind(
batch_context
.as_ref()
.map(|context| context.user_id.as_str()),
)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?
@@ -4316,17 +4417,40 @@ FOR UPDATE
"wallet balance is invalid".to_string(),
));
}
let wallet = map_wallet_row(&row)?;
let already_applied = batch_context.as_ref().is_some_and(|context| {
batch_user_outcomes.as_ref().is_some_and(|outcomes| {
outcomes.get(&context.user_id)
== Some(&AdminUserWalletBalanceBatchUserOutcome::Succeeded)
})
});
if already_applied {
return Ok(Some((wallet, None)));
}
let amount_usd = effective_wallet_adjustment_amount(&input, before_total);
if amount_usd == 0.0 {
if let (Some(context), Some(outcomes)) =
(batch_context.as_ref(), batch_user_outcomes.as_mut())
{
outcomes.insert(
context.user_id.clone(),
AdminUserWalletBalanceBatchUserOutcome::Succeeded,
);
persist_admin_wallet_batch_outcomes(tx, context, outcomes).await?;
}
return Ok(Some((wallet, None)));
}
let mut after_recharge = before_recharge;
let mut after_gift = before_gift;
if input.amount_usd > 0.0 {
if amount_usd > 0.0 {
if input.balance_type.eq_ignore_ascii_case("gift") {
after_gift += input.amount_usd;
after_gift += amount_usd;
} else {
after_recharge += input.amount_usd;
after_recharge += amount_usd;
}
} else {
let mut remaining = -input.amount_usd;
let mut remaining = -amount_usd;
let consume_positive_bucket = |balance: &mut f64, to_consume: &mut f64| {
if *to_consume <= 0.0 {
return;
@@ -4348,7 +4472,7 @@ FOR UPDATE
}
}
let after_total = after_recharge + after_gift;
let after_total_adjusted = before_total_adjusted + input.amount_usd;
let after_total_adjusted = before_total_adjusted + amount_usd;
if !after_recharge.is_finite()
|| !after_gift.is_finite()
|| !after_total.is_finite()
@@ -4387,7 +4511,7 @@ RETURNING
.bind(&input.wallet_id)
.bind(after_recharge)
.bind(after_gift)
.bind(input.amount_usd)
.bind(amount_usd)
.fetch_one(&mut **tx)
.await
.map_postgres_err()?;
@@ -4443,7 +4567,7 @@ VALUES (
)
.bind(&transaction_id)
.bind(&input.wallet_id)
.bind(input.amount_usd)
.bind(amount_usd)
.bind(before_total)
.bind(after_total)
.bind(before_recharge)
@@ -4457,14 +4581,24 @@ VALUES (
.await
.map_postgres_err()?;
if let (Some(context), Some(outcomes)) =
(batch_context.as_ref(), batch_user_outcomes.as_mut())
{
outcomes.insert(
context.user_id.clone(),
AdminUserWalletBalanceBatchUserOutcome::Succeeded,
);
persist_admin_wallet_batch_outcomes(tx, context, outcomes).await?;
}
Ok(Some((
wallet,
StoredAdminWalletTransaction {
Some(StoredAdminWalletTransaction {
id: transaction_id,
wallet_id: input.wallet_id,
category: "adjust".to_string(),
reason_code: "adjust_admin".to_string(),
amount: input.amount_usd,
amount: amount_usd,
balance_before: before_total,
balance_after: after_total,
recharge_balance_before: before_recharge,
@@ -4478,13 +4612,244 @@ VALUES (
operator_email: None,
description: Some(description),
created_at_unix_ms: Some(created_at),
},
}),
)))
})
})
.await
}
async fn prepare_admin_user_wallet_balance_batch(
&self,
input: PrepareAdminUserWalletBalanceBatchInput,
) -> Result<PrepareAdminUserWalletBalanceBatchOutcome, DataLayerError> {
let target_user_ids = serde_json::to_value(&input.target_user_ids).map_err(|error| {
DataLayerError::InvalidInput(format!("invalid batch target users: {error}"))
})?;
let missing_user_ids = serde_json::to_value(&input.missing_user_ids).map_err(|error| {
DataLayerError::InvalidInput(format!("invalid missing batch users: {error}"))
})?;
let warnings = serde_json::to_value(&input.warnings).map_err(|error| {
DataLayerError::InvalidInput(format!("invalid batch warnings: {error}"))
})?;
let now = Utc::now().timestamp().max(0);
self.tx_runner
.run_read_write(|tx| {
Box::pin(async move {
sqlx::query(
r#"
INSERT INTO admin_user_wallet_balance_batches (
admin_user_id, idempotency_key, request_fingerprint, target_user_ids,
missing_user_ids, warnings, user_outcomes, created_at_unix_secs, updated_at_unix_secs
)
VALUES ($1, $2, $3, $4, $5, $6, '{}'::jsonb, $7, $7)
ON CONFLICT (admin_user_id, idempotency_key) DO NOTHING
"#,
)
.bind(&input.admin_user_id)
.bind(&input.idempotency_key)
.bind(&input.request_fingerprint)
.bind(target_user_ids)
.bind(missing_user_ids)
.bind(warnings)
.bind(now)
.execute(&mut **tx)
.await
.map_postgres_err()?;
let row = sqlx::query(
r#"
SELECT admin_user_id, idempotency_key, request_fingerprint, target_user_ids,
missing_user_ids, warnings, user_outcomes
FROM admin_user_wallet_balance_batches
WHERE admin_user_id = $1 AND idempotency_key = $2
FOR UPDATE
"#,
)
.bind(&input.admin_user_id)
.bind(&input.idempotency_key)
.fetch_one(&mut **tx)
.await
.map_postgres_err()?;
let request_fingerprint: String = row_get(&row, "request_fingerprint")?;
if request_fingerprint != input.request_fingerprint {
return Ok(PrepareAdminUserWalletBalanceBatchOutcome::Conflict);
}
let read_json = |column: &str| -> Result<serde_json::Value, DataLayerError> {
row_get(&row, column)
};
let batch = StoredAdminUserWalletBalanceBatch {
admin_user_id: row_get(&row, "admin_user_id")?,
idempotency_key: row_get(&row, "idempotency_key")?,
request_fingerprint,
target_user_ids: serde_json::from_value(read_json("target_user_ids")?)
.map_err(|error| DataLayerError::UnexpectedValue(error.to_string()))?,
missing_user_ids: serde_json::from_value(read_json("missing_user_ids")?)
.map_err(|error| DataLayerError::UnexpectedValue(error.to_string()))?,
warnings: serde_json::from_value(read_json("warnings")?)
.map_err(|error| DataLayerError::UnexpectedValue(error.to_string()))?,
user_outcomes: serde_json::from_value(read_json("user_outcomes")?)
.map_err(|error| DataLayerError::UnexpectedValue(error.to_string()))?,
};
Ok(PrepareAdminUserWalletBalanceBatchOutcome::Ready(batch))
})
})
.await
}
async fn get_admin_user_wallet_balance_batch(
&self,
admin_user_id: &str,
idempotency_key: &str,
expected_fingerprint: &str,
) -> Result<Option<PrepareAdminUserWalletBalanceBatchOutcome>, DataLayerError> {
let row = sqlx::query(
r#"
SELECT admin_user_id, idempotency_key, request_fingerprint, target_user_ids,
missing_user_ids, warnings, user_outcomes
FROM admin_user_wallet_balance_batches
WHERE admin_user_id = $1 AND idempotency_key = $2
"#,
)
.bind(admin_user_id)
.bind(idempotency_key)
.fetch_optional(&self.pool)
.await
.map_postgres_err()?;
let Some(row) = row else {
return Ok(None);
};
let request_fingerprint: String = row_get(&row, "request_fingerprint")?;
if request_fingerprint != expected_fingerprint {
return Ok(Some(PrepareAdminUserWalletBalanceBatchOutcome::Conflict));
}
let read_json =
|column: &str| -> Result<serde_json::Value, DataLayerError> { row_get(&row, column) };
let batch = StoredAdminUserWalletBalanceBatch {
admin_user_id: row_get(&row, "admin_user_id")?,
idempotency_key: row_get(&row, "idempotency_key")?,
request_fingerprint,
target_user_ids: serde_json::from_value(read_json("target_user_ids")?)
.map_err(|error| DataLayerError::UnexpectedValue(error.to_string()))?,
missing_user_ids: serde_json::from_value(read_json("missing_user_ids")?)
.map_err(|error| DataLayerError::UnexpectedValue(error.to_string()))?,
warnings: serde_json::from_value(read_json("warnings")?)
.map_err(|error| DataLayerError::UnexpectedValue(error.to_string()))?,
user_outcomes: serde_json::from_value(read_json("user_outcomes")?)
.map_err(|error| DataLayerError::UnexpectedValue(error.to_string()))?,
};
Ok(Some(PrepareAdminUserWalletBalanceBatchOutcome::Ready(
batch,
)))
}
async fn adjust_admin_user_wallet_balance_batch_user(
&self,
input: AdjustWalletBalanceInBatchInput,
) -> Result<AdminUserWalletBalanceBatchUserOutcome, DataLayerError> {
let context = AdminUserWalletBalanceBatchContext {
admin_user_id: input.admin_user_id.clone(),
idempotency_key: input.idempotency_key.clone(),
user_id: input.user_id.clone(),
};
let mut adjustment = input.adjustment;
adjustment.batch_context = Some(context);
if self.adjust_wallet_balance(adjustment).await?.is_some() {
return Ok(AdminUserWalletBalanceBatchUserOutcome::Succeeded);
}
self.record_admin_user_wallet_balance_batch_failure(
&input.admin_user_id,
&input.idempotency_key,
&input.user_id,
"用户钱包不可用",
)
.await
}
async fn record_admin_user_wallet_balance_batch_failure(
&self,
admin_user_id: &str,
idempotency_key: &str,
user_id: &str,
reason: &str,
) -> Result<AdminUserWalletBalanceBatchUserOutcome, DataLayerError> {
let admin_user_id = admin_user_id.to_string();
let idempotency_key = idempotency_key.to_string();
let user_id = user_id.to_string();
let reason = reason.to_string();
self.tx_runner
.run_read_write(|tx| {
Box::pin(async move {
let row = sqlx::query(
r#"
SELECT target_user_ids, user_outcomes
FROM admin_user_wallet_balance_batches
WHERE admin_user_id = $1 AND idempotency_key = $2
FOR UPDATE
"#,
)
.bind(&admin_user_id)
.bind(&idempotency_key)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?
.ok_or_else(|| {
DataLayerError::InvalidInput(
"admin wallet batch was not prepared".to_string(),
)
})?;
let target_user_ids: Vec<String> =
serde_json::from_value(row_get(&row, "target_user_ids")?).map_err(
|error| {
DataLayerError::UnexpectedValue(format!(
"admin wallet batch target list is invalid: {error}"
))
},
)?;
if !target_user_ids.iter().any(|id| id == &user_id) {
return Err(DataLayerError::InvalidInput(
"user is outside the prepared admin wallet batch".to_string(),
));
}
let mut outcomes: BTreeMap<String, AdminUserWalletBalanceBatchUserOutcome> =
serde_json::from_value(row_get(&row, "user_outcomes")?).map_err(
|error| {
DataLayerError::UnexpectedValue(format!(
"admin wallet batch outcomes are invalid: {error}"
))
},
)?;
if let Some(outcome) = outcomes.get(&user_id) {
return Ok(outcome.clone());
}
let outcome = AdminUserWalletBalanceBatchUserOutcome::Failed(reason);
outcomes.insert(user_id, outcome.clone());
let value = serde_json::to_value(&outcomes).map_err(|error| {
DataLayerError::UnexpectedValue(format!(
"admin wallet batch outcomes are invalid: {error}"
))
})?;
sqlx::query(
r#"
UPDATE admin_user_wallet_balance_batches
SET user_outcomes = $3, updated_at_unix_secs = $4
WHERE admin_user_id = $1 AND idempotency_key = $2
"#,
)
.bind(&admin_user_id)
.bind(&idempotency_key)
.bind(value)
.bind(Utc::now().timestamp().max(0))
.execute(&mut **tx)
.await
.map_postgres_err()?;
Ok(outcome)
})
})
.await
}
async fn create_manual_wallet_recharge(
&self,
mut input: CreateManualWalletRechargeInput,
@@ -8493,13 +8858,16 @@ VALUES ($1, $2, 'gift', 'gift_initial', $3, 0, $3, 0, 0, 0, $3, 'system_task', $
#[cfg(test)]
mod tests {
use aether_data_contracts::repository::wallet::{
CreateManualWalletRechargeInput, CreditAdminPaymentOrderInput, ProcessPaymentCallbackInput,
AdjustWalletBalanceInBatchInput, AdjustWalletBalanceInput,
AdminUserWalletBalanceBatchUserOutcome, CreateManualWalletRechargeInput,
CreditAdminPaymentOrderInput, PrepareAdminUserWalletBalanceBatchInput,
PrepareAdminUserWalletBalanceBatchOutcome, ProcessPaymentCallbackInput,
ProcessPaymentCallbackOutcome, RedeemWalletCodeInput, RedeemWalletCodeOutcome,
WalletLookupKey, WalletMutationOutcome, WalletReadRepository, WalletWriteRepository,
};
use sqlx::Row;
use super::SqlxWalletRepository;
use super::{effective_wallet_adjustment_amount, SqlxWalletRepository};
use crate::{PostgresPoolConfig, PostgresPoolFactory};
#[test]
@@ -8579,6 +8947,7 @@ mod tests {
"payment_orders",
"payment_callbacks",
"wallet_transactions",
"admin_user_wallet_balance_batches",
"user_plan_entitlements",
"redeem_code_batches",
"redeem_codes",
@@ -9225,6 +9594,264 @@ mod tests {
pool.close().await;
}
#[test]
fn bulk_adjustment_clamp_is_opt_in_and_uses_available_total() {
let input = AdjustWalletBalanceInput {
wallet_id: "wallet-1".to_string(),
amount_usd: -100.0,
balance_type: "recharge".to_string(),
operator_id: None,
description: None,
clamp_deduction_to_available_balance: true,
batch_context: None,
};
assert_eq!(effective_wallet_adjustment_amount(&input, 13.0), -13.0);
assert_eq!(effective_wallet_adjustment_amount(&input, 0.0), -0.0);
assert_eq!(effective_wallet_adjustment_amount(&input, -1.0), 1.0);
let legacy_input = AdjustWalletBalanceInput {
clamp_deduction_to_available_balance: false,
..input
};
assert_eq!(
effective_wallet_adjustment_amount(&legacy_input, 13.0),
-100.0
);
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL bootstrap schema"]
async fn live_bulk_wallet_adjustment_persists_actual_delta_and_skips_zero_ledger() {
let pool = isolated_wallet_test_pool().await;
let (wallet_id, _) = seed_wallet(&pool).await;
let repository = SqlxWalletRepository::new(pool.clone());
let (wallet, transaction) = repository
.adjust_wallet_balance(AdjustWalletBalanceInput {
wallet_id: wallet_id.clone(),
amount_usd: -100.0,
balance_type: "recharge".to_string(),
operator_id: Some("admin-user".to_string()),
description: Some("bulk deduction".to_string()),
clamp_deduction_to_available_balance: true,
batch_context: None,
})
.await
.expect("bulk adjustment should succeed")
.expect("wallet should exist");
let transaction =
transaction.expect("positive available balance should create a ledger row");
assert_eq!(transaction.amount, -13.0);
assert_eq!(transaction.balance_before, 13.0);
assert_eq!(transaction.balance_after, 0.0);
assert_eq!(wallet.balance + wallet.gift_balance, 0.0);
let persisted_amount: f64 = sqlx::query_scalar(
"SELECT amount::double precision FROM wallet_transactions WHERE id = $1",
)
.bind(&transaction.id)
.fetch_one(&pool)
.await
.expect("ledger should store the effective deduction");
assert_eq!(persisted_amount, -13.0);
let (wallet, transaction) = repository
.adjust_wallet_balance(AdjustWalletBalanceInput {
wallet_id: wallet_id.clone(),
amount_usd: -1.0,
balance_type: "recharge".to_string(),
operator_id: Some("admin-user".to_string()),
description: Some("bulk deduction at zero".to_string()),
clamp_deduction_to_available_balance: true,
batch_context: None,
})
.await
.expect("zero-balance adjustment should succeed")
.expect("wallet should still exist");
assert_eq!(wallet.balance + wallet.gift_balance, 0.0);
assert!(
transaction.is_none(),
"zero effective delta must not create a ledger row"
);
let transaction_count: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM wallet_transactions WHERE wallet_id = $1")
.bind(&wallet_id)
.fetch_one(&pool)
.await
.expect("ledger row count should be readable");
assert_eq!(transaction_count, 1);
sqlx::query("UPDATE wallets SET balance = -2, gift_balance = 1 WHERE id = $1")
.bind(&wallet_id)
.execute(&pool)
.await
.expect("legacy negative wallet balance should be seeded");
let (wallet, transaction) = repository
.adjust_wallet_balance(AdjustWalletBalanceInput {
wallet_id: wallet_id.clone(),
amount_usd: -1.0,
balance_type: "recharge".to_string(),
operator_id: Some("admin-user".to_string()),
description: Some("bulk deduction floors a negative balance".to_string()),
clamp_deduction_to_available_balance: true,
batch_context: None,
})
.await
.expect("legacy negative wallet should be floored at zero")
.expect("wallet should still exist");
assert_eq!(wallet.balance + wallet.gift_balance, 0.0);
let transaction = transaction.expect("negative balance correction should be ledgered");
assert_eq!(transaction.amount, 1.0);
assert_eq!(transaction.balance_before, -1.0);
assert_eq!(transaction.balance_after, 0.0);
let persisted_correction_amount: f64 = sqlx::query_scalar(
"SELECT amount::double precision FROM wallet_transactions WHERE id = $1",
)
.bind(&transaction.id)
.fetch_one(&pool)
.await
.expect("ledger should store the negative-balance correction");
assert_eq!(persisted_correction_amount, 1.0);
let transaction_count: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM wallet_transactions WHERE wallet_id = $1")
.bind(&wallet_id)
.fetch_one(&pool)
.await
.expect("ledger row count should be readable");
assert_eq!(transaction_count, 2);
pool.close().await;
}
#[tokio::test]
#[ignore = "requires AETHER_TEST_DATABASE_URL and PostgreSQL bootstrap schema"]
async fn live_admin_wallet_balance_batch_replays_committed_and_zero_delta_results_once() {
let pool = isolated_wallet_test_pool().await;
let (wallet_id, user_id) = seed_wallet(&pool).await;
let repository = SqlxWalletRepository::new(pool.clone());
let admin_user_id = "admin-user".to_string();
let make_adjustment =
|idempotency_key: &str, amount_usd: f64| AdjustWalletBalanceInBatchInput {
admin_user_id: admin_user_id.clone(),
idempotency_key: idempotency_key.to_string(),
user_id: user_id.clone(),
adjustment: AdjustWalletBalanceInput {
wallet_id: wallet_id.clone(),
amount_usd,
balance_type: "recharge".to_string(),
operator_id: Some(admin_user_id.clone()),
description: Some("idempotency integration test".to_string()),
clamp_deduction_to_available_balance: true,
batch_context: None,
},
};
let prepare_batch = |idempotency_key: &str, request_fingerprint: &str| {
PrepareAdminUserWalletBalanceBatchInput {
admin_user_id: admin_user_id.clone(),
idempotency_key: idempotency_key.to_string(),
request_fingerprint: request_fingerprint.to_string(),
target_user_ids: vec![user_id.clone()],
missing_user_ids: Vec::new(),
warnings: Vec::new(),
}
};
assert!(matches!(
repository
.prepare_admin_user_wallet_balance_batch(prepare_batch(
"deduct-key",
"deduct-fingerprint"
))
.await
.unwrap(),
PrepareAdminUserWalletBalanceBatchOutcome::Ready(_)
));
let deduct = make_adjustment("deduct-key", -50.0);
assert_eq!(
repository
.adjust_admin_user_wallet_balance_batch_user(deduct.clone())
.await
.unwrap(),
AdminUserWalletBalanceBatchUserOutcome::Succeeded
);
assert_eq!(
repository
.adjust_admin_user_wallet_balance_batch_user(deduct.clone())
.await
.unwrap(),
AdminUserWalletBalanceBatchUserOutcome::Succeeded
);
let wallet = repository
.find(WalletLookupKey::WalletId(&wallet_id))
.await
.unwrap()
.unwrap();
assert_eq!(wallet.balance + wallet.gift_balance, 0.0);
assert!(matches!(
repository
.prepare_admin_user_wallet_balance_batch(prepare_batch(
"zero-key",
"zero-fingerprint"
))
.await
.unwrap(),
PrepareAdminUserWalletBalanceBatchOutcome::Ready(_)
));
let zero_delta = make_adjustment("zero-key", -5.0);
assert_eq!(
repository
.adjust_admin_user_wallet_balance_batch_user(zero_delta.clone())
.await
.unwrap(),
AdminUserWalletBalanceBatchUserOutcome::Succeeded
);
repository
.adjust_wallet_balance(AdjustWalletBalanceInput {
wallet_id: wallet_id.clone(),
amount_usd: 8.0,
balance_type: "recharge".to_string(),
operator_id: Some(admin_user_id.clone()),
description: Some("recharge after zero-delta batch".to_string()),
clamp_deduction_to_available_balance: true,
batch_context: None,
})
.await
.unwrap()
.unwrap();
for replay in [zero_delta, deduct] {
assert_eq!(
repository
.adjust_admin_user_wallet_balance_batch_user(replay)
.await
.unwrap(),
AdminUserWalletBalanceBatchUserOutcome::Succeeded
);
}
assert_eq!(
repository
.prepare_admin_user_wallet_balance_batch(prepare_batch(
"zero-key",
"different-fingerprint"
))
.await
.unwrap(),
PrepareAdminUserWalletBalanceBatchOutcome::Conflict
);
let wallet = repository
.find(WalletLookupKey::WalletId(&wallet_id))
.await
.unwrap()
.unwrap();
assert_eq!(wallet.balance + wallet.gift_balance, 8.0);
let transaction_count: i64 =
sqlx::query_scalar("SELECT COUNT(*) FROM wallet_transactions WHERE wallet_id = $1")
.bind(&wallet_id)
.fetch_one(&pool)
.await
.unwrap();
assert_eq!(transaction_count, 2);
pool.close().await;
}
#[tokio::test]
async fn repository_constructs_from_lazy_pool() {
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
crates/aether-data/contracts/src/repository/usage/metadata_policy.rs
+27 -3
View File
@@ -8,9 +8,10 @@ use serde_json::{Map, Value};
use crate::repository::candidates::sanitize_request_candidate_skip_reason;
use super::{
LIVE_SESSION_METADATA_KEY, PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY,
PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY, PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY,
PROVIDER_REASONING_EFFORT_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY,
normalize_provider_response_model, LIVE_SESSION_METADATA_KEY,
PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY, PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY,
PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY,
PROVIDER_RESPONSE_MODEL_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY,
REALTIME_SESSION_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY,
ROUTING_CANDIDATE_SKIP_REASON_METADATA_KEY, ROUTING_FAILURE_DIAGNOSTIC_METADATA_KEY,
USAGE_AVAILABLE_METADATA_KEY, USAGE_PRICING_AVAILABLE_METADATA_KEY,
@@ -128,6 +129,12 @@ pub fn sanitize_usage_request_metadata_object(source: &Map<String, Value>) -> Op
] {
insert_known_string(source, &mut target, key, sanitize_service_tier);
}
insert_known_string(
source,
&mut target,
PROVIDER_RESPONSE_MODEL_METADATA_KEY,
normalize_provider_response_model,
);
insert_bounded_u64(
source,
&mut target,
@@ -1334,6 +1341,23 @@ mod tests {
}
}
#[test]
fn persistence_projection_keeps_bounded_response_model_only_as_a_string() {
let metadata = sanitize_usage_request_metadata(Some(json!({
"provider_response_model": " GPT-5.1 "
})))
.expect("response model should remain");
assert_eq!(metadata["provider_response_model"], "GPT-5.1");
assert!(sanitize_usage_request_metadata(Some(json!({
"provider_response_model": 42
})))
.is_none());
assert!(sanitize_usage_request_metadata(Some(json!({
"provider_response_model": "x".repeat(257)
})))
.is_none());
}
#[test]
fn persistence_projection_keeps_only_bounded_settlement_facts() {
let metadata = sanitize_usage_request_metadata(Some(json!({
crates/aether-data/contracts/src/repository/usage/mod.rs
+10 -9
View File
@@ -25,11 +25,12 @@ pub use policy::*;
pub use types::{
canonical_usage_body_ref_for, extract_provider_actual_service_tier_from_response,
extract_provider_cache_ttl_minutes_from_metadata, extract_provider_reasoning_effort_from_body,
extract_provider_service_tier_from_body, normalize_provider_service_tier, parse_usage_body_ref,
extract_provider_response_model_from_bodies, extract_provider_service_tier_from_body,
normalize_provider_response_model, normalize_provider_service_tier, parse_usage_body_ref,
resolve_provider_cache_ttl_minutes, resolve_provider_service_tier_from_request_capture,
usage_body_ref, usage_request_metadata_client_family, ApiKeyLastUsedDelta,
ManagementTokenCounterDelta, PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest,
ProxyNodeCounterDelta, StoredProviderApiKeyUsageSummary,
usage_body_capture_is_authoritative, usage_body_ref, usage_request_metadata_client_family,
ApiKeyLastUsedDelta, ManagementTokenCounterDelta, PendingUsageCleanupSummary,
ProviderApiKeyWindowUsageRequest, ProxyNodeCounterDelta, StoredProviderApiKeyUsageSummary,
StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, StoredUsageAuditAggregation, StoredUsageAuditSummary,
StoredUsageBodyPayload, StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,
@@ -56,9 +57,9 @@ pub use types::{
UsageTimeSeriesQuery, UsageWriteRepository, LIVE_SESSION_METADATA_KEY,
PLAN_USAGE_RESERVATION_DEFERRED_METADATA_KEY, PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY,
PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY, PROVIDER_REASONING_EFFORT_METADATA_KEY,
PROVIDER_SERVICE_TIER_METADATA_KEY, REALTIME_SESSION_METADATA_KEY,
REQUESTED_REASONING_EFFORT_METADATA_KEY, ROUTING_CANDIDATE_SKIP_REASON_METADATA_KEY,
ROUTING_FAILURE_DIAGNOSTIC_METADATA_KEY, USAGE_AVAILABLE_METADATA_KEY,
USAGE_PRICING_AVAILABLE_METADATA_KEY, WEBSOCKET_MODE_METADATA_KEY,
WEBSOCKET_TRANSPORT_METADATA_KEY,
PROVIDER_RESPONSE_MODEL_METADATA_KEY, PROVIDER_SERVICE_TIER_METADATA_KEY,
REALTIME_SESSION_METADATA_KEY, REQUESTED_REASONING_EFFORT_METADATA_KEY,
ROUTING_CANDIDATE_SKIP_REASON_METADATA_KEY, ROUTING_FAILURE_DIAGNOSTIC_METADATA_KEY,
USAGE_AVAILABLE_METADATA_KEY, USAGE_PRICING_AVAILABLE_METADATA_KEY,
WEBSOCKET_MODE_METADATA_KEY, WEBSOCKET_TRANSPORT_METADATA_KEY,
};
crates/aether-data/contracts/src/repository/usage/types.rs
+281 -4
View File
@@ -1,3 +1,4 @@
use aether_ai_formats::normalize_api_format_alias;
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde_json::Value;
@@ -6,6 +7,7 @@ pub const PROVIDER_REASONING_EFFORT_METADATA_KEY: &str = "provider_reasoning_eff
pub const REQUESTED_REASONING_EFFORT_METADATA_KEY: &str = "requested_reasoning_effort";
pub const PROVIDER_SERVICE_TIER_METADATA_KEY: &str = "provider_service_tier";
pub const PROVIDER_ACTUAL_SERVICE_TIER_METADATA_KEY: &str = "provider_actual_service_tier";
pub const PROVIDER_RESPONSE_MODEL_METADATA_KEY: &str = "provider_response_model";
pub const PROVIDER_CACHE_TTL_MINUTES_METADATA_KEY: &str = "provider_cache_ttl_minutes";
pub const ROUTING_CANDIDATE_SKIP_REASON_METADATA_KEY: &str = "routing_candidate_skip_reason";
pub const ROUTING_FAILURE_DIAGNOSTIC_METADATA_KEY: &str = "routing_failure_diagnostic";
@@ -119,6 +121,141 @@ pub fn normalize_provider_service_tier(value: &str) -> Option<String> {
Some(value.to_ascii_lowercase())
}
/// 清洗模型名称，保留大小写，只去除首尾空白。
pub fn normalize_provider_response_model(value: &str) -> Option<String> {
let value = value.trim();
if value.is_empty() || value.len() > 256 {
return None;
}
Some(value.to_string())
}
fn extract_model_at_paths(value: &Value, paths: &[&[&str]]) -> Option<String> {
paths.iter().find_map(|path| {
let value = path
.iter()
.try_fold(value, |current, key| current.as_object()?.get(*key))?;
value.as_str().and_then(normalize_provider_response_model)
})
}
fn response_model_paths(provider_api_format: Option<&str>) -> &'static [&'static [&'static str]] {
match normalize_api_format_alias(provider_api_format.unwrap_or_default()).as_str() {
"gemini:generate_content" => {
// Gemini 原生响应使用 modelVersion；部分网关会改写为 model。
&[&["modelVersion"], &["model_version"], &["model"]]
}
"gemini:embedding" => {
// Gemini Embedding 可能返回 model、modelVersion 或 Vertex 的 deployedModelId。
&[
&["model"],
&["modelVersion"],
&["model_version"],
&["deployedModelId"],
&["deployed_model_id"],
]
}
"gemini:interactions" => {
// Interactions 请求既可能叫 model，也可能叫 agent；响应优先读取 model。
&[
&["model"],
&["modelVersion"],
&["model_version"],
&["agent"],
]
}
_ => &[&["model"]],
}
}
fn extract_model_from_known_response_wrappers(
response_body: &Value,
paths: &[&[&str]],
) -> Option<String> {
// 只展开协议中已知的 response/chunks 包装，避免在候选内容、工具参数等任意嵌套
// 对象中搜索同名字段，误把 role="model" 一类内容当成响应模型。
extract_model_at_paths(response_body, paths)
.or_else(|| {
response_body
.get("response")
.and_then(|response| extract_model_at_paths(response, paths))
})
.or_else(|| {
response_body
.get("chunks")
.and_then(Value::as_array)
.and_then(|chunks| {
chunks.iter().rev().find_map(|chunk| {
extract_model_at_paths(chunk, paths).or_else(|| {
chunk
.get("response")
.and_then(|response| extract_model_at_paths(response, paths))
})
})
})
})
.or_else(|| {
response_body
.get("response")
.and_then(|response| response.get("chunks"))
.and_then(Value::as_array)
.and_then(|chunks| {
chunks.iter().rev().find_map(|chunk| {
extract_model_at_paths(chunk, paths).or_else(|| {
chunk
.get("response")
.and_then(|response| extract_model_at_paths(response, paths))
})
})
})
})
}
fn extract_provider_model_from_response_body(
response_body: &Value,
provider_api_format: Option<&str>,
) -> Option<String> {
extract_model_from_known_response_wrappers(
response_body,
response_model_paths(provider_api_format),
)
}
fn extract_provider_model_from_request_body(
request_body: &Value,
request_api_format: Option<&str>,
) -> Option<String> {
let paths: &[&[&str]] =
match normalize_api_format_alias(request_api_format.unwrap_or_default()).as_str() {
"gemini:interactions" => &[&["model"], &["agent"]],
_ => &[&["model"]],
};
extract_model_at_paths(request_body, paths)
}
/// 只有请求体和响应体都可作为完整事实时，才计算响应模型，避免用截断内容猜测。
pub fn extract_provider_response_model_from_bodies(
request_body: Option<&Value>,
request_body_state: Option<UsageBodyCaptureState>,
request_api_format: Option<&str>,
response_body: Option<&Value>,
response_body_state: Option<UsageBodyCaptureState>,
provider_api_format: Option<&str>,
) -> Option<String> {
if !usage_body_capture_is_authoritative(request_body, request_body_state)
|| !usage_body_capture_is_authoritative(response_body, response_body_state)
{
return None;
}
let request_model =
extract_provider_model_from_request_body(request_body?, request_api_format)?;
let response_model =
extract_provider_model_from_response_body(response_body?, provider_api_format)?;
(request_model != response_model).then_some(response_model)
}
/// Resolves a provider processing tier exclusively from the final upstream request.
///
/// A complete captured body is authoritative, including when it contains no tier. The metadata
@@ -129,7 +266,7 @@ pub fn resolve_provider_service_tier_from_request_capture(
provider_request_body_state: Option<UsageBodyCaptureState>,
request_metadata: Option<&Value>,
) -> Option<String> {
if request_body_capture_is_authoritative(provider_request_body, provider_request_body_state) {
if usage_body_capture_is_authoritative(provider_request_body, provider_request_body_state) {
return extract_provider_service_tier_from_body(provider_request_body);
}
@@ -153,7 +290,7 @@ pub fn resolve_provider_service_tier_from_request_capture(
.and_then(normalize_provider_service_tier)
}
fn request_body_capture_is_authoritative(
pub fn usage_body_capture_is_authoritative(
request_body: Option<&Value>,
request_body_state: Option<UsageBodyCaptureState>,
) -> bool {
@@ -184,7 +321,7 @@ fn resolve_reasoning_effort_from_request_capture(
request_metadata: Option<&Value>,
metadata_key: &str,
) -> Option<String> {
if request_body_capture_is_authoritative(request_body, request_body_state) {
if usage_body_capture_is_authoritative(request_body, request_body_state) {
return extract_provider_reasoning_effort_from_body(request_body);
}
@@ -701,6 +838,11 @@ impl StoredRequestUsageAudit {
})
}
pub fn provider_response_model(&self) -> Option<String> {
self.request_metadata_string(PROVIDER_RESPONSE_MODEL_METADATA_KEY)
.and_then(normalize_provider_response_model)
}
pub fn provider_cache_ttl_minutes(&self) -> Option<i64> {
resolve_provider_cache_ttl_minutes(
self.endpoint_api_format
@@ -1096,9 +1238,15 @@ pub struct UsageAuditAggregationQuery {
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub struct UsageAuditSummaryQuery {
/// Optional provider-name allowlist, intersected with provider_name; empty matches nothing.
#[serde(default)]
pub provider_names: Option<Vec<String>>,
pub created_from_unix_secs: u64,
pub created_until_unix_secs: u64,
pub user_id: Option<String>,
/// Optional bulk user scope used by current user-group reporting.
/// An empty list intentionally matches no usage rows.
pub user_ids: Option<Vec<String>>,
pub provider_name: Option<String>,
pub model: Option<String>,
}
@@ -1486,11 +1634,17 @@ pub enum UsageTimeSeriesGranularity {
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct UsageTimeSeriesQuery {
/// Optional provider-name allowlist, intersected with provider_name; empty matches nothing.
#[serde(default)]
pub provider_names: Option<Vec<String>>,
pub created_from_unix_secs: u64,
pub created_until_unix_secs: u64,
pub granularity: UsageTimeSeriesGranularity,
pub tz_offset_minutes: i32,
pub user_id: Option<String>,
/// Optional bulk user scope used by current user-group reporting.
/// An empty list intentionally matches no usage rows.
pub user_ids: Option<Vec<String>>,
pub provider_name: Option<String>,
pub model: Option<String>,
}
@@ -1517,10 +1671,16 @@ pub enum UsageLeaderboardGroupBy {
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct UsageLeaderboardQuery {
/// Optional provider-name allowlist, intersected with provider_name; empty matches nothing.
#[serde(default)]
pub provider_names: Option<Vec<String>>,
pub created_from_unix_secs: u64,
pub created_until_unix_secs: u64,
pub group_by: UsageLeaderboardGroupBy,
pub user_id: Option<String>,
/// Optional bulk user scope used by current user-group reporting.
/// An empty list intentionally matches no usage rows.
pub user_ids: Option<Vec<String>>,
pub provider_name: Option<String>,
pub model: Option<String>,
}
@@ -2709,7 +2869,8 @@ fn parse_timestamp(value: i64, field_name: &str) -> Result<u64, crate::DataLayer
mod tests {
use super::{
canonical_usage_body_ref_for, extract_provider_actual_service_tier_from_response,
extract_provider_service_tier_from_body, normalize_provider_reasoning_effort,
extract_provider_response_model_from_bodies, extract_provider_service_tier_from_body,
normalize_provider_reasoning_effort, normalize_provider_response_model,
normalize_provider_service_tier, resolve_provider_cache_ttl_minutes, usage_body_ref,
StoredRequestUsageAudit, UpsertUsageRecord, UsageBodyCaptureState, UsageBodyCaptureStorage,
UsageBodyField, UsageProviderPerformanceQuery, REALTIME_SESSION_METADATA_KEY,
@@ -2735,6 +2896,122 @@ mod tests {
}
}
#[test]
fn response_model_requires_authoritative_different_top_level_models() {
let request = json!({"model": " gpt-5 "});
let response = json!({"model": " gpt-5.1 "});
assert_eq!(
extract_provider_response_model_from_bodies(
Some(&request),
Some(UsageBodyCaptureState::Inline),
Some("openai:chat"),
Some(&response),
Some(UsageBodyCaptureState::Inline),
Some("openai:chat"),
),
Some("gpt-5.1".to_string())
);
assert_eq!(
extract_provider_response_model_from_bodies(
Some(&request),
Some(UsageBodyCaptureState::Inline),
Some("openai:responses"),
Some(&json!({"model": "gpt-5"})),
Some(UsageBodyCaptureState::Inline),
Some("openai:responses"),
),
None
);
assert_eq!(
extract_provider_response_model_from_bodies(
Some(&request),
Some(UsageBodyCaptureState::Truncated),
Some("openai:chat"),
Some(&response),
Some(UsageBodyCaptureState::Inline),
Some("openai:chat"),
),
None
);
}
#[test]
fn response_model_keeps_case_and_rejects_invalid_values() {
assert_eq!(
normalize_provider_response_model(" GPT-5.1 "),
Some("GPT-5.1".to_string())
);
assert_eq!(normalize_provider_response_model(" "), None);
assert_eq!(normalize_provider_response_model(&"x".repeat(257)), None);
assert_eq!(
extract_provider_response_model_from_bodies(
Some(&json!({"model": "gpt-5"})),
None,
Some("openai:chat"),
Some(&json!({"model": 42})),
None,
Some("openai:chat"),
),
None
);
}
#[test]
fn response_model_uses_provider_format_specific_nested_paths() {
let request = json!({"model": "gemini-2.5-flash"});
let response = json!({
"response": {
"modelVersion": "gemini-2.5-flash-001",
"candidates": [{"content": {"role": "model"}}]
}
});
assert_eq!(
extract_provider_response_model_from_bodies(
Some(&request),
Some(UsageBodyCaptureState::Inline),
Some("gemini:generate_content"),
Some(&response),
Some(UsageBodyCaptureState::Inline),
Some("gemini:generate_content"),
),
Some("gemini-2.5-flash-001".to_string())
);
let wrapped_chunks = json!({
"chunks": [
{"response": {"modelVersion": "gemini-old"}},
{"response": {"modelVersion": "gemini-final"}}
]
});
assert_eq!(
extract_provider_response_model_from_bodies(
Some(&request),
Some(UsageBodyCaptureState::Inline),
Some("gemini:generate_content"),
Some(&wrapped_chunks),
Some(UsageBodyCaptureState::Inline),
Some("gemini:generate_content"),
),
Some("gemini-final".to_string())
);
let ambiguous = json!({
"metadata": {"model": "do-not-use"},
"candidates": [{"content": {"role": "model"}}]
});
assert_eq!(
extract_provider_response_model_from_bodies(
Some(&request),
Some(UsageBodyCaptureState::Inline),
Some("gemini:generate_content"),
Some(&ambiguous),
Some(UsageBodyCaptureState::Inline),
Some("gemini:generate_content"),
),
None
);
}
#[test]
fn provider_fact_normalization_preserves_non_ascii_case_and_byte_count() {
for normalize in [
crates/aether-data/contracts/src/repository/wallet/types.rs
+98 -1
View File
@@ -2463,6 +2463,59 @@ pub struct AdjustWalletBalanceInput {
pub balance_type: String,
pub operator_id: Option<String>,
pub description: Option<String>,
#[serde(default)]
pub clamp_deduction_to_available_balance: bool,
#[serde(default)]
pub batch_context: Option<AdminUserWalletBalanceBatchContext>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdminUserWalletBalanceBatchContext {
pub admin_user_id: String,
pub idempotency_key: String,
pub user_id: String,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct PrepareAdminUserWalletBalanceBatchInput {
pub admin_user_id: String,
pub idempotency_key: String,
pub request_fingerprint: String,
pub target_user_ids: Vec<String>,
pub missing_user_ids: Vec<String>,
pub warnings: Vec<serde_json::Value>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct StoredAdminUserWalletBalanceBatch {
pub admin_user_id: String,
pub idempotency_key: String,
pub request_fingerprint: String,
pub target_user_ids: Vec<String>,
pub missing_user_ids: Vec<String>,
pub warnings: Vec<serde_json::Value>,
pub user_outcomes: std::collections::BTreeMap<String, AdminUserWalletBalanceBatchUserOutcome>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub enum PrepareAdminUserWalletBalanceBatchOutcome {
Ready(StoredAdminUserWalletBalanceBatch),
Conflict,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
#[serde(tag = "status", content = "reason", rename_all = "snake_case")]
pub enum AdminUserWalletBalanceBatchUserOutcome {
Succeeded,
Failed(String),
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct AdjustWalletBalanceInBatchInput {
pub admin_user_id: String,
pub idempotency_key: String,
pub user_id: String,
pub adjustment: AdjustWalletBalanceInput,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
@@ -2894,7 +2947,51 @@ pub trait WalletWriteRepository: Send + Sync {
async fn adjust_wallet_balance(
&self,
input: AdjustWalletBalanceInput,
) -> Result<Option<(StoredWalletSnapshot, StoredAdminWalletTransaction)>, crate::DataLayerError>;
) -> Result<
Option<(StoredWalletSnapshot, Option<StoredAdminWalletTransaction>)>,
crate::DataLayerError,
>;
async fn prepare_admin_user_wallet_balance_batch(
&self,
_input: PrepareAdminUserWalletBalanceBatchInput,
) -> Result<PrepareAdminUserWalletBalanceBatchOutcome, crate::DataLayerError> {
Err(crate::DataLayerError::InvalidInput(
"idempotent admin wallet batches are not available".to_string(),
))
}
async fn get_admin_user_wallet_balance_batch(
&self,
_admin_user_id: &str,
_idempotency_key: &str,
_request_fingerprint: &str,
) -> Result<Option<PrepareAdminUserWalletBalanceBatchOutcome>, crate::DataLayerError> {
Err(crate::DataLayerError::InvalidInput(
"idempotent admin wallet batches are not available".to_string(),
))
}
async fn adjust_admin_user_wallet_balance_batch_user(
&self,
_input: AdjustWalletBalanceInBatchInput,
) -> Result<AdminUserWalletBalanceBatchUserOutcome, crate::DataLayerError> {
Err(crate::DataLayerError::InvalidInput(
"idempotent admin wallet batches are not available".to_string(),
))
}
async fn record_admin_user_wallet_balance_batch_failure(
&self,
_admin_user_id: &str,
_idempotency_key: &str,
_user_id: &str,
_reason: &str,
) -> Result<AdminUserWalletBalanceBatchUserOutcome, crate::DataLayerError> {
Err(crate::DataLayerError::InvalidInput(
"idempotent admin wallet batches are not available".to_string(),
))
}
async fn create_manual_wallet_recharge(
&self,
crates/aether-data/runtime/schema/generated/postgres/baseline/005_wallet_billing.sql
+14
View File
@@ -344,3 +344,17 @@ CREATE INDEX IF NOT EXISTS idx_redeem_codes_status ON public.redeem_codes USING
CREATE INDEX IF NOT EXISTS idx_redeem_codes_redeemed_user ON public.redeem_codes USING btree (redeemed_by_user_id, redeemed_at);
CREATE INDEX IF NOT EXISTS idx_redeem_codes_redeemed_order ON public.redeem_codes USING btree (redeemed_payment_order_id);
CREATE TABLE IF NOT EXISTS public.admin_user_wallet_balance_batches (
admin_user_id character varying(64) NOT NULL,
idempotency_key character varying(128) NOT NULL,
request_fingerprint character varying(64) NOT NULL,
target_user_ids jsonb NOT NULL,
missing_user_ids jsonb NOT NULL,
warnings jsonb NOT NULL,
user_outcomes jsonb NOT NULL,
created_at_unix_secs bigint NOT NULL,
updated_at_unix_secs bigint NOT NULL
);
ALTER TABLE ONLY public.admin_user_wallet_balance_batches ADD CONSTRAINT admin_user_wallet_balance_batches_pkey PRIMARY KEY (admin_user_id, idempotency_key);
crates/aether-data/runtime/schema/logical/005_wallet_billing.toml
+44
View File
@@ -1370,3 +1370,47 @@ columns = ["redeemed_by_user_id", "redeemed_at"]
[[table.redeem_codes.indexes]]
name = "idx_redeem_codes_redeemed_order"
columns = ["redeemed_payment_order_id"]
[table.admin_user_wallet_balance_batches]
domain = "wallet_billing"
order = 110
primary_key = ["admin_user_id", "idempotency_key"]
[[table.admin_user_wallet_balance_batches.columns]]
name = "admin_user_id"
type = "text_id"
length = 64
[[table.admin_user_wallet_balance_batches.columns]]
name = "idempotency_key"
type = "text"
length = 128
[[table.admin_user_wallet_balance_batches.columns]]
name = "request_fingerprint"
type = "text"
length = 64
[[table.admin_user_wallet_balance_batches.columns]]
name = "target_user_ids"
type = "json"
[[table.admin_user_wallet_balance_batches.columns]]
name = "missing_user_ids"
type = "json"
[[table.admin_user_wallet_balance_batches.columns]]
name = "warnings"
type = "json"
[[table.admin_user_wallet_balance_batches.columns]]
name = "user_outcomes"
type = "json"
[[table.admin_user_wallet_balance_batches.columns]]
name = "created_at_unix_secs"
type = "unix_seconds"
[[table.admin_user_wallet_balance_batches.columns]]
name = "updated_at_unix_secs"
type = "unix_seconds"
crates/aether-data/runtime/src/lifecycle/migrate/tests.rs
+3
View File
@@ -1592,6 +1592,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260921010000,
20260921020000,
20260921020100,
20260923000000,
20261001000000,
]
);
@@ -2563,12 +2564,14 @@ INSERT INTO public.stats_daily_api_key (
.expect("API key daily aggregate fixtures should be inserted");
let leaderboard_query = UsageLeaderboardQuery {
provider_names: None,
created_from_unix_secs: u64::try_from(stats_day.timestamp())
.expect("historical stats day should be nonnegative"),
created_until_unix_secs: u64::try_from((stats_day + chrono::Duration::days(1)).timestamp())
.expect("historical stats end should be nonnegative"),
group_by: UsageLeaderboardGroupBy::ApiKey,
user_id: Some("leaderboard-owner".to_string()),
user_ids: None,
provider_name: None,
model: None,
};
crates/aether-data/runtime/src/lifecycle/migrate/tests/legacy_overview_upgrade.rs
+8 -1
View File
@@ -107,7 +107,14 @@ WHERE version=20260919000000;
.iter()
.map(|migration| migration.version)
.collect::<Vec<_>>(),
vec![DIRTY_EVENTS, 20260921010000, 20260921020000, 20260921020100, 20261001000000]
vec![
DIRTY_EVENTS,
20260921010000,
20260921020000,
20260921020100,
20260923000000,
20261001000000,
]
);
assert_eq!(
rows_snapshot(&pool, "_sqlx_migrations").await,
crates/aether-data/runtime/src/repository/candidate_selection/memory.rs
+63
View File
@@ -346,6 +346,16 @@ fn key_auth_channel_matches(row: &StoredMinimalCandidateSelectionRow, api_format
"openai:chat" | "openai:responses" | "claude:messages" | "openai:image"
)
}
"xai" => {
matches!(auth_type.as_str(), "oauth" | "bearer" | "api_key")
&& matches!(
api_format.as_str(),
"openai:responses"
| "openai:responses:compact"
| "openai:image"
| "openai:video"
)
}
"windsurf" => {
matches!(auth_type.as_str(), "oauth" | "api_key" | "bearer")
&& api_format == "openai:chat"
@@ -591,6 +601,59 @@ mod tests {
assert_eq!(rows[0].global_model_name, "grok-4.20-0309-non-reasoning");
}
#[tokio::test]
async fn includes_xai_oauth_rows_for_responses_models() {
let mut row = sample_row("provider-xai", "openai:responses", "grok-4", 10);
row.provider_type = "xai".to_string();
row.provider_name = "xai".to_string();
row.key_auth_type = "oauth".to_string();
row.key_api_formats = Some(vec![
"openai:responses".to_string(),
"openai:responses:compact".to_string(),
]);
let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(vec![row]);
let rows = repository
.list_for_exact_api_format("openai:responses")
.await
.expect("list should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].provider_type, "xai");
assert_eq!(rows[0].global_model_name, "grok-4");
}
#[tokio::test]
async fn includes_xai_oauth_rows_for_image_and_video_models() {
let mut image = sample_row("provider-xai", "openai:image", "grok-imagine-image", 10);
image.provider_type = "xai".to_string();
image.provider_name = "xai".to_string();
image.key_auth_type = "oauth".to_string();
image.key_api_formats = Some(vec!["openai:image".to_string(), "openai:video".to_string()]);
let mut video = image.clone();
video.endpoint_id = "endpoint-video".to_string();
video.endpoint_api_format = "openai:video".to_string();
video.global_model_name = "grok-imagine-video".to_string();
video.model_provider_model_name = "grok-imagine-video".to_string();
let repository = InMemoryMinimalCandidateSelectionReadRepository::seed(vec![image, video]);
let image_rows = repository
.list_for_exact_api_format("openai:image")
.await
.expect("list should succeed");
assert_eq!(image_rows.len(), 1);
assert_eq!(image_rows[0].global_model_name, "grok-imagine-image");
let video_rows = repository
.list_for_exact_api_format("openai:video")
.await
.expect("list should succeed");
assert_eq!(video_rows.len(), 1);
assert_eq!(video_rows[0].global_model_name, "grok-imagine-video");
}
#[tokio::test]
async fn requested_model_filter_respects_endpoint_scoped_default_mapping() {
let mut selected = sample_row("provider-1", "openai:chat", "deepseek-v4-pro", 10);
crates/aether-data/runtime/src/repository/usage/memory.rs
+48
View File
@@ -600,6 +600,22 @@ fn usage_matches_summary_query(
return false;
}
}
if let Some(user_ids) = query.user_ids.as_deref() {
if !item
.user_id
.as_ref()
.is_some_and(|user_id| user_ids.contains(user_id))
{
return false;
}
}
if query
.provider_names
.as_ref()
.is_some_and(|names| !names.contains(&item.provider_name))
{
return false;
}
if let Some(provider_name) = query.provider_name.as_deref() {
if item.provider_name != provider_name {
return false;
@@ -627,6 +643,22 @@ fn usage_matches_time_series_query(
return false;
}
}
if let Some(user_ids) = query.user_ids.as_deref() {
if !item
.user_id
.as_ref()
.is_some_and(|user_id| user_ids.contains(user_id))
{
return false;
}
}
if query
.provider_names
.as_ref()
.is_some_and(|names| !names.contains(&item.provider_name))
{
return false;
}
if let Some(provider_name) = query.provider_name.as_deref() {
if item.provider_name != provider_name {
return false;
@@ -998,6 +1030,22 @@ fn usage_matches_leaderboard_query(
return false;
}
}
if let Some(user_ids) = query.user_ids.as_deref() {
if !item
.user_id
.as_ref()
.is_some_and(|user_id| user_ids.contains(user_id))
{
return false;
}
}
if query
.provider_names
.as_ref()
.is_some_and(|names| !names.contains(&item.provider_name))
{
return false;
}
if let Some(provider_name) = query.provider_name.as_deref() {
if item.provider_name != provider_name {
return false;
crates/aether-data/runtime/src/repository/usage/memory/tests.rs
+138 -1
View File
@@ -18,7 +18,7 @@ use aether_data_contracts::repository::usage::{
UsageAuditAggregationQuery, UsageAuditKeywordSearchQuery, UsageAuditListQuery,
UsageAuditSummaryQuery, UsageBodyCaptureState, UsageBodyField, UsageDashboardSummaryQuery,
UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageProviderPerformanceQuery,
UsageTimeSeriesGranularity,
UsageTimeSeriesGranularity, UsageTimeSeriesQuery,
};
use serde_json::json;
@@ -937,6 +937,7 @@ async fn unmetered_session_audit_counts_lifecycle_without_token_or_cost_contribu
let summary = repository
.summarize_usage_audits(&UsageAuditSummaryQuery {
provider_names: None,
created_from_unix_secs: 0,
created_until_unix_secs: 1_000,
..UsageAuditSummaryQuery::default()
@@ -2693,10 +2694,12 @@ async fn dashboard_and_leaderboard_total_tokens_use_effective_cache_aware_tokens
let leaderboard = repository
.summarize_usage_leaderboard(&UsageLeaderboardQuery {
provider_names: None,
created_from_unix_secs: 1_711_000_000,
created_until_unix_secs: 1_711_000_001,
group_by: UsageLeaderboardGroupBy::User,
user_id: None,
user_ids: None,
provider_name: None,
model: None,
})
@@ -2706,6 +2709,140 @@ async fn dashboard_and_leaderboard_total_tokens_use_effective_cache_aware_tokens
assert_eq!(leaderboard[0].total_tokens, 120);
}
#[tokio::test]
async fn usage_analytics_filters_by_multiple_user_ids() {
let user_one = sample_usage("req-user-1", 1_711_000_000);
let mut user_two = sample_usage("req-user-2", 1_711_000_000);
user_two.user_id = Some("user-2".to_string());
let mut user_three = sample_usage("req-user-3", 1_711_000_000);
user_three.user_id = Some("user-3".to_string());
let repository = InMemoryUsageReadRepository::seed(vec![user_one, user_two, user_three]);
let scoped_user_ids = vec!["user-1".to_string(), "user-2".to_string()];
let summary = repository
.summarize_usage_audits(&UsageAuditSummaryQuery {
provider_names: None,
created_from_unix_secs: 1_711_000_000,
created_until_unix_secs: 1_711_000_001,
user_ids: Some(scoped_user_ids.clone()),
..Default::default()
})
.await
.expect("summary should filter by multiple users");
assert_eq!(summary.total_requests, 2);
let buckets = repository
.summarize_usage_time_series(&UsageTimeSeriesQuery {
provider_names: None,
created_from_unix_secs: 1_711_000_000,
created_until_unix_secs: 1_711_000_001,
granularity: UsageTimeSeriesGranularity::Day,
tz_offset_minutes: 0,
user_id: None,
user_ids: Some(scoped_user_ids.clone()),
provider_name: None,
model: None,
})
.await
.expect("time series should filter by multiple users");
assert_eq!(
buckets
.iter()
.map(|bucket| bucket.total_requests)
.sum::<u64>(),
2
);
let leaderboard = repository
.summarize_usage_leaderboard(&UsageLeaderboardQuery {
provider_names: None,
created_from_unix_secs: 1_711_000_000,
created_until_unix_secs: 1_711_000_001,
group_by: UsageLeaderboardGroupBy::User,
user_id: None,
user_ids: Some(scoped_user_ids),
provider_name: None,
model: None,
})
.await
.expect("leaderboard should filter by multiple users");
assert_eq!(leaderboard.len(), 2);
assert!(leaderboard.iter().all(|item| item.group_key != "user-3"));
}
#[tokio::test]
async fn usage_analytics_intersects_provider_allowlist_and_user_scope() {
let mut a = sample_usage("allowed", 1_711_000_000);
a.provider_name = "Gemini".to_string();
a.user_id = Some("user-1".to_string());
let mut b = a.clone();
b.request_id = "other-provider".to_string();
b.provider_name = "Other".to_string();
let mut c = a.clone();
c.request_id = "other-user".to_string();
c.user_id = Some("user-2".to_string());
let repository = InMemoryUsageReadRepository::seed(vec![a, b, c]);
for (names, provider, count) in [
(
Some(vec!["Gemini".to_string(), "Gemini".to_string()]),
None,
1,
),
(Some(vec![]), None, 0),
(
Some(vec!["Gemini".to_string()]),
Some("Other".to_string()),
0,
),
(None, None, 2),
] {
let summary = repository
.summarize_usage_audits(&UsageAuditSummaryQuery {
created_from_unix_secs: 1_711_000_000,
created_until_unix_secs: 1_711_000_001,
user_ids: Some(vec!["user-1".to_string()]),
provider_names: names.clone(),
provider_name: provider.clone(),
..Default::default()
})
.await
.unwrap();
assert_eq!(summary.total_requests, count);
let buckets = repository
.summarize_usage_time_series(&UsageTimeSeriesQuery {
created_from_unix_secs: 1_711_000_000,
created_until_unix_secs: 1_711_000_001,
user_id: None,
user_ids: Some(vec!["user-1".to_string()]),
provider_names: names.clone(),
provider_name: provider.clone(),
model: None,
granularity: UsageTimeSeriesGranularity::Day,
tz_offset_minutes: 0,
})
.await
.unwrap();
assert_eq!(
buckets.iter().map(|row| row.total_requests).sum::<u64>(),
count
);
let rows = repository
.summarize_usage_leaderboard(&UsageLeaderboardQuery {
created_from_unix_secs: 1_711_000_000,
created_until_unix_secs: 1_711_000_001,
user_id: None,
user_ids: Some(vec!["user-1".to_string()]),
provider_names: names,
provider_name: provider,
model: None,
group_by: UsageLeaderboardGroupBy::User,
})
.await
.unwrap();
assert_eq!(rows.iter().map(|row| row.request_count).sum::<u64>(), count);
}
}
#[tokio::test]
async fn summarizes_provider_api_key_last_used_at_in_seconds() {
let repository = InMemoryUsageReadRepository::seed(vec![
crates/aether-data/runtime/src/repository/wallet/memory.rs
+7 -2
View File
@@ -2328,8 +2328,13 @@ impl WalletWriteRepository for InMemoryWalletRepository {
async fn adjust_wallet_balance(
&self,
_input: AdjustWalletBalanceInput,
) -> Result<Option<(StoredWalletSnapshot, super::StoredAdminWalletTransaction)>, DataLayerError>
{
) -> Result<
Option<(
StoredWalletSnapshot,
Option<super::StoredAdminWalletTransaction>,
)>,
DataLayerError,
> {
Ok(None)
}
crates/aether-data/runtime/src/repository/wallet/mod.rs
+23 -20
View File
@@ -16,27 +16,30 @@ pub use aether_data_contracts::repository::wallet::{
wallet_recharge_order_is_checkout_placeholder,
wallet_recharge_order_is_reclaimable_placeholder, wallet_recharge_replay_matches,
wallet_recharge_response_is_checkout_placeholder, wallet_refund_proof_is_success,
AdjustWalletBalanceInput, AdminPaymentCallbackRecord, AdminPaymentOrderListQuery,
AdminRedeemCodeBatchListQuery, AdminRedeemCodeListQuery, AdminWalletLedgerQuery,
AdminWalletListQuery, AdminWalletPaymentOrderRecord, AdminWalletRefundRecord,
AdminWalletRefundRequestListQuery, AdminWalletTransactionRecord, CanonicalWalletRefundFields,
CompareAndSwapPaymentOrderStripeClientSecretInput, CompleteAdminWalletRefundInput,
CreateAdminRedeemCodeBatchInput, CreateAdminRedeemCodeBatchResult,
CreateManualWalletRechargeInput, CreatePlanPurchaseOrderInput, CreatePlanPurchaseOrderOutcome,
CreateWalletRechargeOrderInput, CreateWalletRechargeOrderOutcome,
CreateWalletRefundRequestInput, CreateWalletRefundRequestOutcome,
CreatedAdminRedeemCodePlaintext, CreditAdminPaymentOrderInput, DeleteAdminRedeemCodeBatchInput,
AdjustWalletBalanceInBatchInput, AdjustWalletBalanceInput, AdminPaymentCallbackRecord,
AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery, AdminRedeemCodeListQuery,
AdminUserWalletBalanceBatchContext, AdminUserWalletBalanceBatchUserOutcome,
AdminWalletLedgerQuery, AdminWalletListQuery, AdminWalletPaymentOrderRecord,
AdminWalletRefundRecord, AdminWalletRefundRequestListQuery, AdminWalletTransactionRecord,
CanonicalWalletRefundFields, CompareAndSwapPaymentOrderStripeClientSecretInput,
CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput,
CreateAdminRedeemCodeBatchResult, CreateManualWalletRechargeInput,
CreatePlanPurchaseOrderInput, CreatePlanPurchaseOrderOutcome, CreateWalletRechargeOrderInput,
CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput,
CreateWalletRefundRequestOutcome, CreatedAdminRedeemCodePlaintext,
CreditAdminPaymentOrderInput, DeleteAdminRedeemCodeBatchInput,
DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput, FailAdminWalletRefundInput,
FailWalletRechargeCheckoutInput, InitializeAuthWalletOutcome, ProcessAdminWalletRefundInput,
ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome, ReclaimWalletRechargeCheckoutInput,
RedeemWalletCodeInput, RedeemWalletCodeOutcome, StoredAdminPaymentCallback,
StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder, StoredAdminPaymentOrderPage,
StoredAdminRedeemCode, StoredAdminRedeemCodeBatch, StoredAdminRedeemCodeBatchPage,
StoredAdminRedeemCodePage, StoredAdminWalletLedgerItem, StoredAdminWalletLedgerPage,
StoredAdminWalletListItem, StoredAdminWalletListPage, StoredAdminWalletRefund,
StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestItem,
StoredAdminWalletRefundRequestPage, StoredAdminWalletTransaction,
StoredAdminWalletTransactionPage, StoredWalletDailyUsageLedger,
FailWalletRechargeCheckoutInput, InitializeAuthWalletOutcome,
PrepareAdminUserWalletBalanceBatchInput, PrepareAdminUserWalletBalanceBatchOutcome,
ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome,
ReclaimWalletRechargeCheckoutInput, RedeemWalletCodeInput, RedeemWalletCodeOutcome,
StoredAdminPaymentCallback, StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder,
StoredAdminPaymentOrderPage, StoredAdminRedeemCode, StoredAdminRedeemCodeBatch,
StoredAdminRedeemCodeBatchPage, StoredAdminRedeemCodePage, StoredAdminUserWalletBalanceBatch,
StoredAdminWalletLedgerItem, StoredAdminWalletLedgerPage, StoredAdminWalletListItem,
StoredAdminWalletListPage, StoredAdminWalletRefund, StoredAdminWalletRefundPage,
StoredAdminWalletRefundRequestItem, StoredAdminWalletRefundRequestPage,
StoredAdminWalletTransaction, StoredAdminWalletTransactionPage, StoredWalletDailyUsageLedger,
StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, UpdateAdminWalletRefundGatewayInput,
UpdateWalletRechargeCheckoutInput, WalletLookupKey, WalletMutationOutcome,
WalletReadRepository, WalletReadSeed, WalletReadSnapshot, WalletRepository,