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

Rootfix usage counter outbox

Browse Source
This commit is contained in:
fawney19
2026-05-19 00:30:24 +08:00
parent 19c8688eb1
commit 5eed329916
60 changed files with 4910 additions and 362 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
crates/aether-data/src/lifecycle/bootstrap/postgres.rs
View File

@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260516000000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260518000000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT

9
crates/aether-data/src/lifecycle/migrate/tests.rs
View File

@@ -306,6 +306,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260512110000,
20260515000000,
20260516000000,
20260518000000,
]
);
}
@@ -377,6 +378,13 @@ fn empty_database_snapshot_sql_includes_usage_body_blobs_and_audit_admin_role()
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains(
"ALTER TABLE public.stats_daily_model\n ADD COLUMN IF NOT EXISTS cache_creation_ephemeral_5m_tokens bigint DEFAULT '0'::bigint NOT NULL,"
));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL
.contains("CREATE TABLE IF NOT EXISTS public.usage_counter_deltas"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("ix_usage_counter_deltas_unprocessed"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_entitlement_usage_entitlement_date"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("idx_video_tasks_due_poll"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("request_count bigint DEFAULT 0"));
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("usage_count bigint DEFAULT 0 NOT NULL"));
}
#[test]
@@ -1123,6 +1131,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260512110000,
20260515000000,
20260516000000,
20260518000000,
]
);
}

77
crates/aether-data/src/repository/settlement/postgres.rs
View File

@@ -125,6 +125,22 @@ DO UPDATE SET
updated_at = NOW()
"#;
const ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL: &str = r#"
INSERT INTO usage_counter_deltas (
id,
request_id,
kind,
target_id,
total_cost_usd_delta
) VALUES (
$1,
$2,
'provider_monthly',
$3,
$4
)
"#;
#[derive(Debug, Clone)]
pub struct SqlxSettlementRepository {
tx_runner: PostgresTransactionRunner,
@@ -193,6 +209,37 @@ where
Ok(())
}
async fn enqueue_provider_monthly_usage_delta<'e, E>(
executor: E,
request_id: &str,
provider_id: &str,
total_cost_usd_delta: f64,
) -> Result<(), DataLayerError>
where
E: sqlx::Executor<'e, Database = sqlx::Postgres>,
{
let request_id = request_id.trim();
let provider_id = provider_id.trim();
if request_id.is_empty() || provider_id.is_empty() || total_cost_usd_delta == 0.0 {
return Ok(());
}
if !total_cost_usd_delta.is_finite() {
return Err(DataLayerError::UnexpectedValue(format!(
"provider monthly usage delta is not finite for {provider_id}"
)));
}
sqlx::query(ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL)
.bind(uuid::Uuid::new_v4().to_string())
.bind(request_id)
.bind(provider_id)
.bind(total_cost_usd_delta)
.execute(executor)
.await
.map_postgres_err()?;
Ok(())
}
#[derive(Debug, Default)]
struct DailyQuotaDebitResult {
debited_usd: f64,
@@ -663,23 +710,13 @@ WHERE id = $1
.as_deref()
.filter(|value| !value.is_empty())
{
let quota_row = sqlx::query(
r#"
UPDATE providers
SET
monthly_used_usd = COALESCE(monthly_used_usd, 0) + $2,
updated_at = NOW()
WHERE id = $1
RETURNING CAST(monthly_used_usd AS DOUBLE PRECISION) AS monthly_used_usd
"#,
enqueue_provider_monthly_usage_delta(
&mut **tx,
&input.request_id,
provider_id,
input.actual_total_cost_usd,
)
.bind(provider_id)
.bind(input.actual_total_cost_usd)
.fetch_optional(&mut **tx)
.await
.map_postgres_err()?;
settlement.provider_monthly_used_usd =
quota_row.and_then(|row| row.try_get("monthly_used_usd").ok());
.await?;
}
}
@@ -732,6 +769,14 @@ mod tests {
assert!(!source.contains("UPDATE \"usage\"\nSET\n wallet_id = $2"));
}
#[test]
fn settlement_sql_enqueues_provider_monthly_usage_delta() {
let source = include_str!("postgres.rs");
assert!(super::ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL.contains("usage_counter_deltas"));
assert!(super::ENQUEUE_PROVIDER_MONTHLY_USAGE_DELTA_SQL.contains("'provider_monthly'"));
assert!(!source.contains("UPDATE providers\nSET\n monthly_used_usd"));
}
#[test]
fn settlement_sql_blocks_standalone_key_owner_wallet_fallback() {
let source = include_str!("postgres.rs");

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

@@ -363,15 +363,15 @@ mod sqlite;
#[allow(unused_imports)]
pub(crate) use aether_data_contracts::repository::usage::{
usage_request_metadata_client_family, PendingUsageCleanupSummary,
ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, StoredUsageAuditAggregation, StoredUsageAuditSummary,
StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,
StoredUsageCacheAffinityIntervalRow, StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary,
StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow,
StoredUsageDashboardProviderCount, StoredUsageDashboardSummary,
StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
usage_request_metadata_client_family, ApiKeyLastUsedDelta, ManagementTokenCounterDelta,
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, ProxyNodeCounterDelta,
StoredProviderApiKeyUsageSummary, StoredProviderApiKeyWindowUsageSummary,
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
StoredUsageAuditAggregation, StoredUsageAuditSummary, StoredUsageBreakdownSummaryRow,
StoredUsageCacheAffinityHitSummary, StoredUsageCacheAffinityIntervalRow,
StoredUsageCacheHitSummary, StoredUsageCostSavingsSummary, StoredUsageDailySummary,
StoredUsageDashboardDailyBreakdownRow, StoredUsageDashboardProviderCount,
StoredUsageDashboardSummary, StoredUsageErrorDistributionRow, StoredUsageLeaderboardSummary,
StoredUsagePerformancePercentilesRow, StoredUsageProviderPerformance,
StoredUsageProviderPerformanceProviderRow, StoredUsageProviderPerformanceSummary,
StoredUsageProviderPerformanceTimelineRow, StoredUsageSettledCostSummary,
@@ -380,7 +380,8 @@ pub(crate) use aether_data_contracts::repository::usage::{
UsageAuditListQuery, UsageAuditSummaryQuery, UsageBreakdownGroupBy, UsageBreakdownSummaryQuery,
UsageCacheAffinityHitSummaryQuery, UsageCacheAffinityIntervalGroupBy,
UsageCacheAffinityIntervalQuery, UsageCacheHitSummaryQuery, UsageCleanupPreviewCounts,
UsageCleanupSummary, UsageCleanupWindow, UsageCostSavingsSummaryQuery, UsageDailyHeatmapQuery,
UsageCleanupSummary, UsageCleanupWindow, UsageCostSavingsSummaryQuery,
UsageCounterFlushSummary, UsageCounterHealthSnapshot, UsageDailyHeatmapQuery,
UsageDashboardDailyBreakdownQuery, UsageDashboardProviderCountsQuery,
UsageDashboardSummaryQuery, UsageErrorDistributionQuery, UsageLeaderboardGroupBy,
UsageLeaderboardQuery, UsageMonitoringErrorCountQuery, UsageMonitoringErrorListQuery,
@@ -423,7 +424,10 @@ impl ApiKeyUsageDelta {
total_requests: after.total_requests - before.total_requests,
total_tokens: after.total_tokens - before.total_tokens,
total_cost_usd: after.total_cost_usd - before.total_cost_usd,
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
candidate_last_used_at_unix_secs: newer_last_used_at(
before.last_used_at_unix_secs,
after.last_used_at_unix_secs,
),
removed_last_used_at_unix_secs: None,
}
}
@@ -530,7 +534,10 @@ impl ProviderApiKeyUsageDelta {
total_tokens: after.total_tokens - before.total_tokens,
total_cost_usd: after.total_cost_usd - before.total_cost_usd,
total_response_time_ms: after.total_response_time_ms - before.total_response_time_ms,
candidate_last_used_at_unix_secs: after.last_used_at_unix_secs,
candidate_last_used_at_unix_secs: newer_last_used_at(
before.last_used_at_unix_secs,
after.last_used_at_unix_secs,
),
removed_last_used_at_unix_secs: None,
usage_created_at_unix_secs: after.usage_created_at_unix_secs,
}
@@ -633,6 +640,9 @@ pub(crate) fn provider_api_key_usage_is_error(
pub(crate) fn provider_api_key_usage_contribution(
usage: &StoredRequestUsageAudit,
) -> Option<ProviderApiKeyUsageContribution> {
if matches!(usage.status.as_str(), "pending" | "streaming") {
return None;
}
let key_id = usage
.provider_api_key_id
.as_deref()
@@ -694,6 +704,9 @@ pub(crate) fn model_usage_contribution(
pub(crate) fn api_key_usage_contribution(
usage: &StoredRequestUsageAudit,
) -> Option<ApiKeyUsageContribution> {
if matches!(usage.status.as_str(), "pending" | "streaming") {
return None;
}
let api_key_id = usage
.api_key_id
.as_deref()
@@ -714,14 +727,23 @@ pub(crate) fn api_key_usage_contribution(
})
}
fn newer_last_used_at(before: Option<u64>, after: Option<u64>) -> Option<u64> {
match (before, after) {
(Some(before), Some(after)) if after > before => Some(after),
(None, Some(after)) => Some(after),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::{
api_key_usage_contribution, incoming_usage_can_recover_terminal_failure,
model_usage_contribution, provider_api_key_usage_contribution,
provider_api_key_usage_is_error, provider_api_key_usage_is_success,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure, ModelUsageDelta,
StoredRequestUsageAudit, UpsertUsageRecord,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
ApiKeyUsageDelta, ModelUsageDelta, ProviderApiKeyUsageDelta, StoredRequestUsageAudit,
UpsertUsageRecord,
};
#[test]
@@ -996,6 +1018,95 @@ mod tests {
assert_eq!(contribution.total_tokens, 20);
assert_eq!(contribution.total_cost_usd, 0.25);
assert_eq!(contribution.last_used_at_unix_secs, Some(123));
let mut streaming = usage.clone();
streaming.status = "streaming".to_string();
assert!(api_key_usage_contribution(&streaming).is_none());
let mut pending = usage;
pending.status = "pending".to_string();
assert!(api_key_usage_contribution(&pending).is_none());
}
#[test]
fn provider_api_key_usage_contribution_tracks_terminal_requests_only() {
let usage = StoredRequestUsageAudit::new(
"usage-1".to_string(),
"request-1".to_string(),
Some("user-1".to_string()),
Some("api-key-1".to_string()),
None,
None,
"OpenAI".to_string(),
"gpt-5".to_string(),
None,
Some("provider-1".to_string()),
None,
Some("provider-key-1".to_string()),
None,
None,
None,
None,
None,
None,
None,
false,
false,
12,
8,
20,
0.25,
0.25,
Some(200),
None,
None,
Some(120),
None,
"completed".to_string(),
"settled".to_string(),
123,
124,
Some(125),
)
.expect("usage should build");
assert!(provider_api_key_usage_contribution(&usage).is_some());
let mut streaming = usage.clone();
streaming.status = "streaming".to_string();
assert!(provider_api_key_usage_contribution(&streaming).is_none());
let mut pending = usage;
pending.status = "pending".to_string();
assert!(provider_api_key_usage_contribution(&pending).is_none());
}
#[test]
fn usage_delta_between_does_not_emit_duplicate_last_used_candidate() {
let api_key_contribution = super::ApiKeyUsageContribution {
api_key_id: "api-key-1".to_string(),
total_requests: 1,
total_tokens: 20,
total_cost_usd: 0.25,
last_used_at_unix_secs: Some(123),
};
assert!(ApiKeyUsageDelta::between(&api_key_contribution, &api_key_contribution).is_noop());
let provider_contribution = super::ProviderApiKeyUsageContribution {
key_id: "provider-key-1".to_string(),
request_count: 1,
success_count: 1,
error_count: 0,
total_tokens: 20,
total_cost_usd: 0.25,
total_response_time_ms: 120,
last_used_at_unix_secs: Some(123),
usage_created_at_unix_secs: Some(123),
};
assert!(
ProviderApiKeyUsageDelta::between(&provider_contribution, &provider_contribution,)
.is_noop()
);
}
#[test]

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

File diff suppressed because it is too large Load Diff

1
crates/aether-data/src/repository/usage/postgres/queries/apply_api_key_usage_delta_sql.sql
View File

@@ -17,6 +17,7 @@ SET
SELECT MAX(created_at)
FROM "usage"
WHERE api_key_id = $1
AND status NOT IN ('pending', 'streaming')
)
ELSE last_used_at
END

170
crates/aether-data/src/repository/usage/postgres/queries/apply_provider_api_key_codex_window_usage_delta_sql.sql
View File

@@ -1,170 +0,0 @@
WITH target_key AS (
SELECT
id,
COALESCE(status_snapshot::jsonb, '{}'::jsonb) AS snapshot
FROM provider_api_keys
WHERE id = $1
AND jsonb_typeof((status_snapshot::jsonb) -> 'quota' -> 'windows') = 'array'
AND lower(BTRIM(COALESCE((status_snapshot::jsonb) -> 'quota' ->> 'provider_type', ''))) = 'codex'
FOR UPDATE
),
window_items AS (
SELECT
target_key.id,
window_rows.window_item,
window_rows.window_ordinality
FROM target_key
CROSS JOIN LATERAL jsonb_array_elements(target_key.snapshot -> 'quota' -> 'windows')
WITH ORDINALITY AS window_rows(window_item, window_ordinality)
),
parsed_windows AS (
SELECT
window_items.id,
window_items.window_item,
window_items.window_ordinality,
lower(BTRIM(COALESCE(window_items.window_item ->> 'code', ''))) AS window_code,
CASE
WHEN text_values.reset_at_text ~ '^[0-9]+$'
AND (
length(text_values.reset_at_text) < 19
OR (
length(text_values.reset_at_text) = 19
AND text_values.reset_at_text <= '9223372036854775807'
)
)
THEN text_values.reset_at_text::BIGINT
ELSE NULL
END AS reset_at,
CASE
WHEN text_values.window_minutes_text ~ '^[0-9]+$'
AND (
length(text_values.window_minutes_text) < 19
OR (
length(text_values.window_minutes_text) = 19
AND text_values.window_minutes_text <= '9223372036854775807'
)
)
THEN text_values.window_minutes_text::BIGINT
ELSE CASE lower(BTRIM(COALESCE(window_items.window_item ->> 'code', '')))
WHEN '5h' THEN 300
WHEN 'weekly' THEN 10080
ELSE NULL
END
END AS window_minutes,
CASE
WHEN text_values.usage_reset_at_text ~ '^[0-9]+$'
AND (
length(text_values.usage_reset_at_text) < 19
OR (
length(text_values.usage_reset_at_text) = 19
AND text_values.usage_reset_at_text <= '9223372036854775807'
)
)
THEN text_values.usage_reset_at_text::BIGINT
ELSE NULL
END AS usage_reset_at,
CASE
WHEN text_values.request_count_text ~ '^[0-9]+$'
AND (
length(text_values.request_count_text) < 19
OR (
length(text_values.request_count_text) = 19
AND text_values.request_count_text <= '9223372036854775807'
)
)
THEN text_values.request_count_text::BIGINT
ELSE 0
END AS current_request_count,
CASE
WHEN text_values.total_tokens_text ~ '^[0-9]+$'
AND (
length(text_values.total_tokens_text) < 19
OR (
length(text_values.total_tokens_text) = 19
AND text_values.total_tokens_text <= '9223372036854775807'
)
)
THEN text_values.total_tokens_text::BIGINT
ELSE 0
END AS current_total_tokens,
CASE
WHEN text_values.total_cost_usd_text ~ '^[-+]?[0-9]+([.][0-9]+)?$'
THEN text_values.total_cost_usd_text::DOUBLE PRECISION
ELSE 0
END AS current_total_cost_usd
FROM window_items
CROSS JOIN LATERAL (
SELECT
BTRIM(COALESCE(window_items.window_item ->> 'reset_at', '')) AS reset_at_text,
BTRIM(COALESCE(window_items.window_item ->> 'window_minutes', '')) AS window_minutes_text,
BTRIM(COALESCE(window_items.window_item ->> 'usage_reset_at', '')) AS usage_reset_at_text,
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'request_count', '')) AS request_count_text,
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'total_tokens', '')) AS total_tokens_text,
BTRIM(COALESCE(window_items.window_item -> 'usage' ->> 'total_cost_usd', '')) AS total_cost_usd_text
) AS text_values
),
window_usage AS (
SELECT
parsed_windows.*,
CASE
WHEN parsed_windows.window_minutes BETWEEN 0 AND 153722867280912930
THEN parsed_windows.window_minutes * 60
ELSE NULL
END AS window_seconds
FROM parsed_windows
),
updated_windows AS (
SELECT
window_usage.id,
jsonb_agg(
CASE
WHEN window_usage.window_code IN ('5h', 'weekly')
AND window_usage.reset_at IS NOT NULL
AND window_usage.window_seconds IS NOT NULL
AND window_usage.reset_at >= window_usage.window_seconds
AND window_usage.reset_at > $2
AND GREATEST(
window_usage.reset_at - window_usage.window_seconds,
COALESCE(window_usage.usage_reset_at, 0)
) <= $2
THEN jsonb_set(
window_usage.window_item,
'{usage}',
jsonb_build_object(
'request_count',
LEAST(
GREATEST(window_usage.current_request_count::NUMERIC + $3::NUMERIC, 0),
9223372036854775807
)::BIGINT,
'total_tokens',
LEAST(
GREATEST(window_usage.current_total_tokens::NUMERIC + $4::NUMERIC, 0),
9223372036854775807
)::BIGINT,
'total_cost_usd',
to_char(
GREATEST(COALESCE(window_usage.current_total_cost_usd, 0) + $5, 0),
'FM999999999999999990.00000000'
)
),
true
)
ELSE window_usage.window_item
END
ORDER BY window_usage.window_ordinality
) AS windows
FROM window_usage
GROUP BY window_usage.id
)
UPDATE provider_api_keys AS keys
SET
status_snapshot = jsonb_set(
target_key.snapshot,
'{quota,windows}',
updated_windows.windows,
true
)::json,
updated_at = NOW()
FROM target_key
JOIN updated_windows ON updated_windows.id = target_key.id
WHERE keys.id = target_key.id

1
crates/aether-data/src/repository/usage/postgres/queries/apply_provider_api_key_usage_delta_sql.sql
View File

@@ -20,6 +20,7 @@ SET
SELECT MAX(created_at)
FROM "usage"
WHERE provider_api_key_id = $1
AND status NOT IN ('pending', 'streaming')
)
ELSE last_used_at
END

3
crates/aether-data/src/repository/usage/postgres/queries/rebuild_api_key_usage_stats_sql.sql
View File

@@ -1,7 +1,7 @@
WITH aggregated AS (
SELECT
api_key_id,
COUNT(*)::INTEGER AS total_requests,
COUNT(*)::BIGINT AS total_requests,
COALESCE(SUM(
GREATEST(
COALESCE(
@@ -16,6 +16,7 @@ WITH aggregated AS (
FROM usage_billing_facts AS "usage"
WHERE api_key_id IS NOT NULL
AND BTRIM(api_key_id) <> ''
AND status NOT IN ('pending', 'streaming')
GROUP BY api_key_id
)
UPDATE api_keys

9
crates/aether-data/src/repository/usage/postgres/queries/rebuild_provider_api_key_usage_stats_sql.sql
View File

@@ -1,7 +1,7 @@
WITH aggregated AS (
SELECT
provider_api_key_id,
COUNT(*)::INTEGER AS request_count,
COUNT(*)::BIGINT AS request_count,
COALESCE(SUM(
CASE
WHEN status IN ('completed', 'success', 'ok', 'billed', 'settled')
@@ -10,7 +10,7 @@ WITH aggregated AS (
THEN 1
ELSE 0
END
), 0)::INTEGER AS success_count,
), 0)::BIGINT AS success_count,
COALESCE(SUM(
CASE
WHEN status NOT IN ('pending', 'streaming')
@@ -22,7 +22,7 @@ WITH aggregated AS (
THEN 1
ELSE 0
END
), 0)::INTEGER AS error_count,
), 0)::BIGINT AS error_count,
COALESCE(SUM(
GREATEST(
COALESCE(
@@ -42,11 +42,12 @@ WITH aggregated AS (
THEN GREATEST(response_time_ms, 0)
ELSE 0
END
), 0)::INTEGER AS total_response_time_ms,
), 0)::BIGINT AS total_response_time_ms,
MAX(created_at) AS last_used_at
FROM usage_billing_facts AS "usage"
WHERE provider_api_key_id IS NOT NULL
AND BTRIM(provider_api_key_id) <> ''
AND status NOT IN ('pending', 'streaming')
GROUP BY provider_api_key_id
)
UPDATE provider_api_keys

52
crates/aether-data/src/repository/usage/postgres/tests.rs
View File

@@ -234,25 +234,6 @@ fn usage_sql_summarizes_usage_by_provider_api_key_ids_in_database() {
assert!(super::SUMMARIZE_USAGE_BY_PROVIDER_API_KEY_IDS_SQL.contains("ANY($1::TEXT[])"));
}
#[test]
fn usage_sql_materializes_provider_key_window_usage_in_status_snapshot() {
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
.contains("UPDATE provider_api_keys AS keys"));
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("jsonb_set"));
assert!(
super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("'{quota,windows}'")
);
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("'usage'"));
assert!(
super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("WHEN '5h' THEN 300")
);
assert!(super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
.contains("WHEN 'weekly' THEN 10080"));
assert!(
!super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL.contains("usage_billing_facts")
);
}
#[test]
fn usage_sql_rebuilds_provider_key_window_usage_into_status_snapshot() {
assert!(super::REBUILD_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_STATS_SQL
@@ -280,18 +261,44 @@ fn usage_sql_serializes_request_id_upserts_before_reading_previous_usage() {
.contains("lock_usage_request_id_in_tx(tx, &usage.request_id).await?;"));
}
#[test]
fn usage_sql_moves_shared_counter_updates_behind_outbox() {
let source = include_str!("mod.rs");
assert!(super::INSERT_USAGE_COUNTER_DELTA_SQL.contains("usage_counter_deltas"));
assert!(super::CLAIM_USAGE_COUNTER_DELTAS_SQL.contains("FOR UPDATE SKIP LOCKED"));
assert!(super::MARK_USAGE_COUNTER_DELTAS_PROCESSED_SQL.contains("processed_at = NOW()"));
assert!(super::TRY_LOCK_USAGE_COUNTER_FLUSH_SQL.contains("pg_try_advisory_xact_lock"));
assert!(source.contains("enqueue_api_key_usage_delta_in_tx("));
assert!(source.contains("enqueue_provider_api_key_usage_delta_in_tx("));
assert!(source.contains("enqueue_model_usage_delta_in_tx("));
assert!(source.contains("apply_provider_api_key_main_usage_delta_in_tx("));
assert!(source.contains("USAGE_COUNTER_KIND_PROVIDER_MONTHLY"));
assert!(source.contains("apply_provider_monthly_usage_delta_in_tx("));
}
#[test]
fn usage_sql_exposes_counter_outbox_health() {
assert!(super::READ_USAGE_COUNTER_HEALTH_SQL.contains("pending_rows"));
assert!(super::READ_USAGE_COUNTER_HEALTH_SQL.contains("oldest_pending_created_at_unix_secs"));
assert!(super::READ_USAGE_COUNTER_HEALTH_SQL.contains("latest_processed_at_unix_secs"));
assert!(super::READ_PENDING_USAGE_COUNTER_DELTAS_BY_KIND_SQL.contains("GROUP BY kind"));
}
#[test]
fn usage_sql_rebuild_matches_online_api_key_usage_semantics() {
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("COUNT(*)::INTEGER"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("COUNT(*)::BIGINT"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("COALESCE("));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("total_tokens,"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL
.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL.contains("AND BTRIM(api_key_id) <> ''"));
assert!(super::REBUILD_API_KEY_USAGE_STATS_SQL
.contains("AND status NOT IN ('pending', 'streaming')"));
}
#[test]
fn usage_sql_rebuild_matches_online_provider_key_usage_semantics() {
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL.contains("COUNT(*)::BIGINT"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("NULLIF(BTRIM(error_message), '') IS NULL"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL.contains("COALESCE("));
@@ -300,6 +307,8 @@ fn usage_sql_rebuild_matches_online_provider_key_usage_semantics() {
.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("AND BTRIM(provider_api_key_id) <> ''"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("AND status NOT IN ('pending', 'streaming')"));
}
#[test]
@@ -503,8 +512,7 @@ fn usage_sql_raw_aggregates_use_canonical_billing_facts() {
.contains("FROM usage_billing_facts AS \"usage\""));
assert!(super::SUMMARIZE_USAGE_TOTALS_BY_USER_IDS_SQL
.contains("FROM usage_billing_facts AS \"usage\""));
assert!(!super::APPLY_PROVIDER_API_KEY_CODEX_WINDOW_USAGE_DELTA_SQL
.contains("usage_billing_facts AS \"usage\""));
assert!(!source.contains("apply_provider_api_key_codex_window_usage_delta_in_tx"));
}
#[test]