i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-06 17:37: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/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 {