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

fix(usage): canonicalize cached token totals

Browse Source
This commit is contained in:
elky
2026-07-16 16:38:30 +08:00
parent e369e4aab1
commit 8fe4d24408
17 changed files with 1804 additions and 1124 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/adapters/postgres/src/usage/mod.rs
+48 -252
View File
@@ -3178,7 +3178,14 @@ 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(input_tokens + output_tokens), 0)::BIGINT AS recorded_total_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,
@@ -3210,7 +3217,14 @@ 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(input_tokens + output_tokens), 0)::BIGINT AS recorded_total_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,
@@ -3257,18 +3271,8 @@ SELECT
COALESCE(SUM(GREATEST(COALESCE("usage".input_tokens, 0), 0)), 0)::BIGINT AS input_tokens,
COALESCE(SUM(GREATEST(COALESCE("usage".output_tokens, 0), 0)), 0)::BIGINT AS output_tokens,
COALESCE(SUM(GREATEST(COALESCE("usage".total_tokens, 0), 0)), 0)::BIGINT AS recorded_total_tokens,
COALESCE(SUM(
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END
), 0)::BIGINT AS cache_creation_tokens,
COALESCE(SUM(GREATEST(COALESCE("usage".cache_creation_input_tokens, 0), 0)), 0)::BIGINT
AS cache_creation_tokens,
COALESCE(SUM(GREATEST(COALESCE("usage".cache_creation_input_tokens_5m, 0), 0)), 0)::BIGINT
AS cache_creation_ephemeral_5m_tokens,
COALESCE(SUM(GREATEST(COALESCE("usage".cache_creation_input_tokens_1h, 0), 0)), 0)::BIGINT
@@ -3945,80 +3949,10 @@ SELECT
COALESCE(SUM(GREATEST(COALESCE("usage".input_tokens, 0), 0)), 0)::BIGINT AS input_tokens,
COALESCE(SUM(GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0)), 0)::BIGINT
AS cache_read_tokens,
COALESCE(SUM(
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END
), 0)::BIGINT AS cache_creation_tokens,
COALESCE(SUM(
CASE
WHEN split_part(lower(COALESCE(COALESCE("usage".endpoint_api_format, "usage".api_format), '')), ':', 1)
IN ('claude', 'anthropic')
THEN GREATEST(COALESCE("usage".input_tokens, 0), 0)
+ CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END
+ GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0)
WHEN split_part(lower(COALESCE(COALESCE("usage".endpoint_api_format, "usage".api_format), '')), ':', 1)
IN ('openai', 'gemini', 'google')
THEN (
CASE
WHEN GREATEST(COALESCE("usage".input_tokens, 0), 0) <= 0 THEN 0
WHEN GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0) <= 0
THEN GREATEST(COALESCE("usage".input_tokens, 0), 0)
ELSE GREATEST(
GREATEST(COALESCE("usage".input_tokens, 0), 0)
- GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0),
0
)
END
) + GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0)
ELSE CASE
WHEN (
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END
) > 0
THEN GREATEST(COALESCE("usage".input_tokens, 0), 0)
+ (
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END
)
+ GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0)
ELSE GREATEST(COALESCE("usage".input_tokens, 0), 0)
+ GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0)
END
END
), 0)::BIGINT AS total_input_context,
COALESCE(SUM(GREATEST(COALESCE("usage".cache_creation_input_tokens, 0), 0)), 0)::BIGINT
AS cache_creation_tokens,
COALESCE(SUM(GREATEST(COALESCE("usage".total_input_context, 0), 0)), 0)::BIGINT
AS total_input_context,
COALESCE(SUM(COALESCE(CAST("usage".cache_read_cost_usd AS DOUBLE PRECISION), 0)), 0)
AS cache_read_cost_usd,
COALESCE(SUM(COALESCE(CAST("usage".cache_creation_cost_usd AS DOUBLE PRECISION), 0)), 0)
@@ -4887,7 +4821,7 @@ SELECT
{group_column} AS group_key,
COALESCE(SUM(total_requests), 0)::BIGINT AS request_count,
COALESCE(SUM(input_tokens), 0)::BIGINT AS input_tokens,
COALESCE(SUM(effective_input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0)::BIGINT AS total_tokens,
COALESCE(SUM(total_tokens), 0)::BIGINT AS total_tokens,
COALESCE(SUM(output_tokens), 0)::BIGINT AS output_tokens,
COALESCE(SUM(effective_input_tokens), 0)::BIGINT AS effective_input_tokens,
COALESCE(SUM(total_input_context), 0)::BIGINT AS total_input_context,
@@ -4944,18 +4878,11 @@ WITH filtered_usage AS (
SELECT
{group_key_expr} AS group_key,
GREATEST(COALESCE("usage".input_tokens, 0), 0) AS input_tokens,
GREATEST(COALESCE("usage".effective_input_tokens, 0), 0) AS effective_input_tokens,
GREATEST(COALESCE("usage".output_tokens, 0), 0) AS output_tokens,
GREATEST(COALESCE("usage".total_tokens, 0), 0) AS total_tokens,
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END AS cache_creation_tokens,
GREATEST(COALESCE("usage".total_input_context, 0), 0) AS total_input_context,
GREATEST(COALESCE("usage".cache_creation_input_tokens, 0), 0) AS cache_creation_tokens,
GREATEST(COALESCE("usage".cache_creation_input_tokens_5m, 0), 0)
AS cache_creation_ephemeral_5m_tokens,
GREATEST(COALESCE("usage".cache_creation_input_tokens_1h, 0), 0)
@@ -4995,8 +4922,7 @@ WITH filtered_usage AS (
AND "usage".response_time_ms IS NOT NULL
THEN 1
ELSE 0
END AS successful_response_time_samples,
COALESCE(COALESCE("usage".endpoint_api_format, "usage".api_format), '') AS normalized_api_format
END AS successful_response_time_samples
FROM usage_billing_facts AS "usage"
"#,
));
@@ -5054,8 +4980,10 @@ normalized_usage AS (
SELECT
group_key,
input_tokens,
effective_input_tokens,
total_tokens,
output_tokens,
total_input_context,
cache_creation_tokens,
cache_creation_ephemeral_5m_tokens,
cache_creation_ephemeral_1h_tokens,
@@ -5066,49 +4994,7 @@ normalized_usage AS (
response_time_ms,
response_time_samples,
successful_response_time_ms,
successful_response_time_samples,
CASE
WHEN input_tokens <= 0 THEN 0
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
= 'openai'
AND (cache_creation_tokens > 0 OR cache_read_tokens > 0)
THEN GREATEST(input_tokens - cache_creation_tokens - cache_read_tokens, 0)
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('gemini', 'google')
AND cache_read_tokens > 0
THEN GREATEST(input_tokens - cache_read_tokens, 0)
ELSE input_tokens
END AS effective_input_tokens,
CASE
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('claude', 'anthropic')
THEN input_tokens + cache_creation_tokens + cache_read_tokens
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
= 'openai'
THEN (
CASE
WHEN input_tokens <= 0 THEN 0
WHEN cache_creation_tokens > 0 OR cache_read_tokens > 0
THEN GREATEST(input_tokens - cache_creation_tokens - cache_read_tokens, 0)
ELSE input_tokens
END
) + cache_creation_tokens + cache_read_tokens
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('gemini', 'google')
THEN (
CASE
WHEN input_tokens <= 0 THEN 0
WHEN cache_read_tokens > 0
THEN GREATEST(input_tokens - cache_read_tokens, 0)
ELSE input_tokens
END
) + cache_read_tokens
ELSE CASE
WHEN cache_creation_tokens > 0
THEN input_tokens + cache_creation_tokens + cache_read_tokens
ELSE input_tokens + cache_read_tokens
END
END AS total_input_context
successful_response_time_samples
FROM filtered_usage
)
SELECT
@@ -6801,33 +6687,7 @@ SELECT
{group_key_expr} AS group_key,
MAX({legacy_name_expr}) AS legacy_name,
COUNT(*)::BIGINT AS request_count,
COALESCE(SUM(
CASE
WHEN GREATEST(COALESCE("usage".input_tokens, 0), 0) <= 0 THEN 0
WHEN GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0) <= 0
THEN GREATEST(COALESCE("usage".input_tokens, 0), 0)
WHEN split_part(lower(COALESCE(COALESCE("usage".endpoint_api_format, "usage".api_format), '')), ':', 1)
IN ('openai', 'gemini', 'google')
THEN GREATEST(
GREATEST(COALESCE("usage".input_tokens, 0), 0)
- GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0),
0
)
ELSE GREATEST(COALESCE("usage".input_tokens, 0), 0)
END
+ GREATEST(COALESCE("usage".output_tokens, 0), 0)
+ CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END
+ GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0)
), 0)::BIGINT AS total_tokens,
COALESCE(SUM(GREATEST(COALESCE("usage".total_tokens, 0), 0)), 0)::BIGINT AS total_tokens,
COALESCE(SUM(COALESCE(CAST("usage".total_cost_usd AS DOUBLE PRECISION), 0)), 0)
AS total_cost_usd
FROM usage_billing_facts AS "usage"
@@ -6983,10 +6843,7 @@ SELECT
user_id AS group_key,
MAX(NULLIF(BTRIM(username), '')) AS legacy_name,
COALESCE(SUM(total_requests), 0)::BIGINT AS request_count,
COALESCE(
SUM(effective_input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
0
)::BIGINT AS total_tokens,
COALESCE(SUM(total_tokens), 0)::BIGINT AS total_tokens,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost_usd
FROM stats_user_daily_provider
WHERE date >=
@@ -7012,10 +6869,7 @@ SELECT
user_id AS group_key,
MAX(NULLIF(BTRIM(username), '')) AS legacy_name,
COALESCE(SUM(total_requests), 0)::BIGINT AS request_count,
COALESCE(
SUM(effective_input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
0
)::BIGINT AS total_tokens,
COALESCE(SUM(total_tokens), 0)::BIGINT AS total_tokens,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost_usd
FROM stats_user_daily_model
WHERE date >=
@@ -7300,25 +7154,16 @@ WITH filtered_usage AS (
{provider_display_name_expr} AS provider_display_name,
{provider_identity_source_expr} AS provider_identity_source,
COALESCE("usage".api_format, 'unknown') AS api_format_group_key,
GREATEST(COALESCE("usage".input_tokens, 0), 0) AS input_tokens,
GREATEST(COALESCE("usage".effective_input_tokens, 0), 0) AS effective_input_tokens,
GREATEST(COALESCE("usage".output_tokens, 0), 0) AS output_tokens,
GREATEST(COALESCE("usage".total_tokens, 0), 0) AS total_tokens,
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END AS cache_creation_tokens,
GREATEST(COALESCE("usage".total_input_context, 0), 0) AS total_input_context,
GREATEST(COALESCE("usage".cache_creation_input_tokens, 0), 0) AS cache_creation_tokens,
GREATEST(COALESCE("usage".cache_creation_input_tokens_5m, 0), 0)
AS cache_creation_ephemeral_5m_tokens,
GREATEST(COALESCE("usage".cache_creation_input_tokens_1h, 0), 0)
AS cache_creation_ephemeral_1h_tokens,
GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0) AS cache_read_tokens,
COALESCE("usage".endpoint_api_format, "usage".api_format) AS normalized_api_format,
COALESCE(CAST("usage".total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
COALESCE(CAST("usage".actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
GREATEST(COALESCE("usage".response_time_ms, 0), 0) AS response_time_ms,
@@ -7343,6 +7188,8 @@ normalized_usage AS (
{secondary_name_expr} AS secondary_name,
total_tokens,
output_tokens,
effective_input_tokens,
total_input_context,
cache_creation_tokens,
cache_creation_ephemeral_5m_tokens,
cache_creation_ephemeral_1h_tokens,
@@ -7350,49 +7197,7 @@ normalized_usage AS (
total_cost_usd,
actual_total_cost_usd,
response_time_ms,
success_flag,
CASE
WHEN input_tokens <= 0 THEN 0
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
= 'openai'
AND (cache_creation_tokens > 0 OR cache_read_tokens > 0)
THEN GREATEST(input_tokens - cache_creation_tokens - cache_read_tokens, 0)
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('gemini', 'google')
AND cache_read_tokens > 0
THEN GREATEST(input_tokens - cache_read_tokens, 0)
ELSE input_tokens
END AS effective_input_tokens,
CASE
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('claude', 'anthropic')
THEN input_tokens + cache_creation_tokens + cache_read_tokens
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
= 'openai'
THEN (
CASE
WHEN input_tokens <= 0 THEN 0
WHEN cache_creation_tokens > 0 OR cache_read_tokens > 0
THEN GREATEST(input_tokens - cache_creation_tokens - cache_read_tokens, 0)
ELSE input_tokens
END
) + cache_creation_tokens + cache_read_tokens
WHEN split_part(lower(COALESCE(normalized_api_format, '')), ':', 1)
IN ('gemini', 'google')
THEN (
CASE
WHEN input_tokens <= 0 THEN 0
WHEN cache_read_tokens > 0
THEN GREATEST(input_tokens - cache_read_tokens, 0)
ELSE input_tokens
END
) + cache_read_tokens
ELSE CASE
WHEN cache_creation_tokens > 0
THEN input_tokens + cache_creation_tokens + cache_read_tokens
ELSE input_tokens + cache_read_tokens
END
END AS total_input_context
success_flag
FROM filtered_usage
),
aggregated_usage AS (
@@ -7552,14 +7357,7 @@ LIMIT $3
r#"SELECT
DATE("usage".created_at) AS day,
COUNT(*)::BIGINT AS requests,
COALESCE(SUM("usage".input_tokens + "usage".output_tokens
+ CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (COALESCE("usage".cache_creation_input_tokens_5m, 0) + COALESCE("usage".cache_creation_input_tokens_1h, 0)) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0) + COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END
+ COALESCE("usage".cache_read_input_tokens, 0)), 0)::BIGINT AS total_tokens,
COALESCE(SUM(GREATEST(COALESCE("usage".total_tokens, 0), 0)), 0)::BIGINT AS total_tokens,
COALESCE(SUM(CAST("usage".total_cost_usd AS DOUBLE PRECISION)), 0) AS total_cost_usd,
COALESCE(SUM(CAST("usage".actual_total_cost_usd AS DOUBLE PRECISION)), 0) AS actual_total_cost_usd
FROM usage_billing_facts AS "usage"
@@ -7803,14 +7601,9 @@ ORDER BY api_key_id ASC
SELECT
api_key_id,
COALESCE(
SUM(
COALESCE(
total_tokens,
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
)
),
SUM(GREATEST(COALESCE(total_tokens, 0), 0)),
0
) AS total_tokens
)::BIGINT AS total_tokens
FROM usage_billing_facts AS "usage"
WHERE api_key_id = ANY(
"#,
@@ -11281,9 +11074,12 @@ fn usage_total_input_context(
.saturating_add(cache_creation_tokens)
.saturating_add(cache_read_tokens),
),
"openai" | "gemini" | "google" => {
Some(effective_input_tokens.saturating_add(cache_read_tokens))
}
"openai" => Some(
effective_input_tokens
.saturating_add(cache_creation_tokens)
.saturating_add(cache_read_tokens),
),
"gemini" | "google" => Some(effective_input_tokens.saturating_add(cache_read_tokens)),
_ => Some(
input_tokens
.saturating_add(cache_creation_tokens)
crates/aether-data/adapters/postgres/src/usage/queries/find_by_id_sql.sql
+80 -15
View File
@@ -24,27 +24,92 @@ SELECT
CAST(COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens) AS INTEGER) AS output_tokens,
CAST(COALESCE(
CASE
WHEN usage_settlement_snapshots.billing_input_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_output_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_read_tokens IS NOT NULL
THEN COALESCE(usage_settlement_snapshots.billing_input_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_output_tokens, 0)
+ COALESCE(
usage_settlement_snapshots.billing_cache_creation_tokens,
COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0),
WHEN usage_settlement_snapshots.billing_effective_input_tokens IS NOT NULL
THEN GREATEST(usage_settlement_snapshots.billing_effective_input_tokens, 0)
+ GREATEST(
COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens, 0),
0
)
+ GREATEST(
COALESCE(
usage_settlement_snapshots.billing_cache_creation_tokens,
CASE
WHEN usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
THEN COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0)
END,
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END,
0
),
0
)
+ GREATEST(
COALESCE(
usage_settlement_snapshots.billing_cache_read_tokens,
"usage".cache_read_input_tokens,
0
),
0
)
WHEN usage_settlement_snapshots.billing_total_input_context IS NOT NULL
THEN GREATEST(usage_settlement_snapshots.billing_total_input_context, 0)
+ GREATEST(
COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens, 0),
0
)
+ COALESCE(usage_settlement_snapshots.billing_cache_read_tokens, 0)
END,
"usage".total_tokens
NULLIF(GREATEST(COALESCE("usage".total_tokens, 0), 0), 0),
CASE
WHEN split_part(lower(COALESCE(COALESCE("usage".endpoint_api_format, "usage".api_format), '')), ':', 1)
IN ('openai', 'gemini', 'google')
THEN GREATEST(COALESCE("usage".input_tokens, 0), 0)
+ GREATEST(COALESCE("usage".output_tokens, 0), 0)
ELSE GREATEST(COALESCE("usage".input_tokens, 0), 0)
+ GREATEST(COALESCE("usage".output_tokens, 0), 0)
+ GREATEST(
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END,
0
)
+ GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0)
END
) AS INTEGER) AS total_tokens,
CAST(COALESCE(
usage_settlement_snapshots.billing_cache_creation_tokens,
"usage".cache_creation_input_tokens,
CASE
WHEN usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
THEN COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0)
END,
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END,
0
) AS INTEGER) AS cache_creation_input_tokens,
CAST(COALESCE(
crates/aether-data/adapters/postgres/src/usage/queries/find_by_request_id_sql.sql
+80 -15
View File
@@ -24,27 +24,92 @@ SELECT
CAST(COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens) AS INTEGER) AS output_tokens,
CAST(COALESCE(
CASE
WHEN usage_settlement_snapshots.billing_input_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_output_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_read_tokens IS NOT NULL
THEN COALESCE(usage_settlement_snapshots.billing_input_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_output_tokens, 0)
+ COALESCE(
usage_settlement_snapshots.billing_cache_creation_tokens,
COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0),
WHEN usage_settlement_snapshots.billing_effective_input_tokens IS NOT NULL
THEN GREATEST(usage_settlement_snapshots.billing_effective_input_tokens, 0)
+ GREATEST(
COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens, 0),
0
)
+ GREATEST(
COALESCE(
usage_settlement_snapshots.billing_cache_creation_tokens,
CASE
WHEN usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
THEN COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0)
END,
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END,
0
),
0
)
+ GREATEST(
COALESCE(
usage_settlement_snapshots.billing_cache_read_tokens,
"usage".cache_read_input_tokens,
0
),
0
)
WHEN usage_settlement_snapshots.billing_total_input_context IS NOT NULL
THEN GREATEST(usage_settlement_snapshots.billing_total_input_context, 0)
+ GREATEST(
COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens, 0),
0
)
+ COALESCE(usage_settlement_snapshots.billing_cache_read_tokens, 0)
END,
"usage".total_tokens
NULLIF(GREATEST(COALESCE("usage".total_tokens, 0), 0), 0),
CASE
WHEN split_part(lower(COALESCE(COALESCE("usage".endpoint_api_format, "usage".api_format), '')), ':', 1)
IN ('openai', 'gemini', 'google')
THEN GREATEST(COALESCE("usage".input_tokens, 0), 0)
+ GREATEST(COALESCE("usage".output_tokens, 0), 0)
ELSE GREATEST(COALESCE("usage".input_tokens, 0), 0)
+ GREATEST(COALESCE("usage".output_tokens, 0), 0)
+ GREATEST(
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END,
0
)
+ GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0)
END
) AS INTEGER) AS total_tokens,
CAST(COALESCE(
usage_settlement_snapshots.billing_cache_creation_tokens,
"usage".cache_creation_input_tokens,
CASE
WHEN usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
THEN COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0)
END,
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END,
0
) AS INTEGER) AS cache_creation_input_tokens,
CAST(COALESCE(
crates/aether-data/adapters/postgres/src/usage/queries/list_recent_usage_audits_prefix.sql
+80 -15
View File
@@ -24,27 +24,92 @@ SELECT
CAST(COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens) AS INTEGER) AS output_tokens,
CAST(COALESCE(
CASE
WHEN usage_settlement_snapshots.billing_input_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_output_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_read_tokens IS NOT NULL
THEN COALESCE(usage_settlement_snapshots.billing_input_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_output_tokens, 0)
+ COALESCE(
usage_settlement_snapshots.billing_cache_creation_tokens,
COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0),
WHEN usage_settlement_snapshots.billing_effective_input_tokens IS NOT NULL
THEN GREATEST(usage_settlement_snapshots.billing_effective_input_tokens, 0)
+ GREATEST(
COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens, 0),
0
)
+ GREATEST(
COALESCE(
usage_settlement_snapshots.billing_cache_creation_tokens,
CASE
WHEN usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
THEN COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0)
END,
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END,
0
),
0
)
+ GREATEST(
COALESCE(
usage_settlement_snapshots.billing_cache_read_tokens,
"usage".cache_read_input_tokens,
0
),
0
)
WHEN usage_settlement_snapshots.billing_total_input_context IS NOT NULL
THEN GREATEST(usage_settlement_snapshots.billing_total_input_context, 0)
+ GREATEST(
COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens, 0),
0
)
+ COALESCE(usage_settlement_snapshots.billing_cache_read_tokens, 0)
END,
"usage".total_tokens
NULLIF(GREATEST(COALESCE("usage".total_tokens, 0), 0), 0),
CASE
WHEN split_part(lower(COALESCE(COALESCE("usage".endpoint_api_format, "usage".api_format), '')), ':', 1)
IN ('openai', 'gemini', 'google')
THEN GREATEST(COALESCE("usage".input_tokens, 0), 0)
+ GREATEST(COALESCE("usage".output_tokens, 0), 0)
ELSE GREATEST(COALESCE("usage".input_tokens, 0), 0)
+ GREATEST(COALESCE("usage".output_tokens, 0), 0)
+ GREATEST(
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END,
0
)
+ GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0)
END
) AS INTEGER) AS total_tokens,
CAST(COALESCE(
usage_settlement_snapshots.billing_cache_creation_tokens,
"usage".cache_creation_input_tokens,
CASE
WHEN usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
THEN COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0)
END,
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END,
0
) AS INTEGER) AS cache_creation_input_tokens,
CAST(COALESCE(
crates/aether-data/adapters/postgres/src/usage/queries/list_usage_audits_prefix.sql
+80 -15
View File
@@ -24,27 +24,92 @@ SELECT
CAST(COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens) AS INTEGER) AS output_tokens,
CAST(COALESCE(
CASE
WHEN usage_settlement_snapshots.billing_input_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_output_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_read_tokens IS NOT NULL
THEN COALESCE(usage_settlement_snapshots.billing_input_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_output_tokens, 0)
+ COALESCE(
usage_settlement_snapshots.billing_cache_creation_tokens,
COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0),
WHEN usage_settlement_snapshots.billing_effective_input_tokens IS NOT NULL
THEN GREATEST(usage_settlement_snapshots.billing_effective_input_tokens, 0)
+ GREATEST(
COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens, 0),
0
)
+ GREATEST(
COALESCE(
usage_settlement_snapshots.billing_cache_creation_tokens,
CASE
WHEN usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
THEN COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0)
END,
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END,
0
),
0
)
+ GREATEST(
COALESCE(
usage_settlement_snapshots.billing_cache_read_tokens,
"usage".cache_read_input_tokens,
0
),
0
)
WHEN usage_settlement_snapshots.billing_total_input_context IS NOT NULL
THEN GREATEST(usage_settlement_snapshots.billing_total_input_context, 0)
+ GREATEST(
COALESCE(usage_settlement_snapshots.billing_output_tokens, "usage".output_tokens, 0),
0
)
+ COALESCE(usage_settlement_snapshots.billing_cache_read_tokens, 0)
END,
"usage".total_tokens
NULLIF(GREATEST(COALESCE("usage".total_tokens, 0), 0), 0),
CASE
WHEN split_part(lower(COALESCE(COALESCE("usage".endpoint_api_format, "usage".api_format), '')), ':', 1)
IN ('openai', 'gemini', 'google')
THEN GREATEST(COALESCE("usage".input_tokens, 0), 0)
+ GREATEST(COALESCE("usage".output_tokens, 0), 0)
ELSE GREATEST(COALESCE("usage".input_tokens, 0), 0)
+ GREATEST(COALESCE("usage".output_tokens, 0), 0)
+ GREATEST(
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END,
0
)
+ GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0)
END
) AS INTEGER) AS total_tokens,
CAST(COALESCE(
usage_settlement_snapshots.billing_cache_creation_tokens,
"usage".cache_creation_input_tokens,
CASE
WHEN usage_settlement_snapshots.billing_cache_creation_5m_tokens IS NOT NULL
OR usage_settlement_snapshots.billing_cache_creation_1h_tokens IS NOT NULL
THEN COALESCE(usage_settlement_snapshots.billing_cache_creation_5m_tokens, 0)
+ COALESCE(usage_settlement_snapshots.billing_cache_creation_1h_tokens, 0)
END,
CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END,
0
) AS INTEGER) AS cache_creation_input_tokens,
CAST(COALESCE(
crates/aether-data/adapters/postgres/src/usage/queries/rebuild_api_key_usage_stats_sql.sql
+1 -4
View File
@@ -4,10 +4,7 @@ WITH aggregated AS (
COUNT(*)::BIGINT AS total_requests,
COALESCE(SUM(
GREATEST(
COALESCE(
total_tokens,
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
),
COALESCE(total_tokens, 0),
0
)::BIGINT
), 0)::BIGINT AS total_tokens,
crates/aether-data/adapters/postgres/src/usage/queries/rebuild_provider_api_key_usage_stats_sql.sql
+1 -4
View File
@@ -27,10 +27,7 @@ WITH aggregated AS (
CASE
WHEN status IN ('pending', 'streaming') THEN 0
ELSE GREATEST(
COALESCE(
total_tokens,
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
),
COALESCE(total_tokens, 0),
0
)::BIGINT
END
crates/aether-data/adapters/postgres/src/usage/queries/summarize_total_tokens_by_api_key_ids_sql.sql
+2 -5
View File
@@ -2,13 +2,10 @@ SELECT
api_key_id,
COALESCE(
SUM(
COALESCE(
total_tokens,
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
)
GREATEST(COALESCE(total_tokens, 0), 0)
),
0
) AS total_tokens
)::BIGINT AS total_tokens
FROM usage_billing_facts AS "usage"
WHERE api_key_id = ANY($1::TEXT[])
GROUP BY api_key_id
crates/aether-data/adapters/postgres/src/usage/tests.rs
+327 -29
View File
@@ -7,11 +7,11 @@ use super::{
inflate_usage_json_value, prepare_request_metadata_for_body_storage,
prepare_usage_body_storage, resolved_read_usage_body_ref, resolved_write_usage_body_ref,
split_dashboard_daily_aggregate_range, split_dashboard_hourly_aggregate_range,
usage_body_capture_state_for_storage, usage_body_ref, usage_http_audit_body_refs,
usage_http_audit_capture_mode, usage_routing_snapshot_from_usage,
usage_settlement_pricing_snapshot_from_usage, AggregateRangeSplit, SqlxUsageReadRepository,
UsageHttpAuditRefs, UsageRoutingSnapshot, UsageSettlementPricingSnapshot,
MAX_INLINE_USAGE_BODY_BYTES,
usage_body_capture_state_for_storage, usage_body_ref, usage_effective_input_tokens,
usage_http_audit_body_refs, usage_http_audit_capture_mode, usage_routing_snapshot_from_usage,
usage_settlement_pricing_snapshot_from_usage, usage_total_input_context, AggregateRangeSplit,
SqlxUsageReadRepository, UsageHttpAuditRefs, UsageRoutingSnapshot,
UsageSettlementPricingSnapshot, MAX_INLINE_USAGE_BODY_BYTES,
};
use crate::{PostgresPoolConfig, PostgresPoolFactory};
use aether_data_contracts::repository::usage::{
@@ -600,8 +600,25 @@ fn usage_sql_does_not_require_updated_at_column() {
#[test]
fn usage_sql_summarizes_tokens_by_api_key_ids_in_database() {
assert!(super::SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL.contains("GROUP BY api_key_id"));
assert!(super::SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL.contains("ANY($1::TEXT[])"));
let sql = super::SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL;
assert!(sql.contains("GROUP BY api_key_id"));
assert!(sql.contains("ANY($1::TEXT[])"));
assert!(sql.contains("GREATEST(COALESCE(total_tokens, 0), 0)"));
assert!(sql.contains(")::BIGINT AS total_tokens"));
assert!(!sql.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
let source = include_str!("mod.rs");
let implementation = source
.split("pub async fn summarize_total_tokens_by_api_key_ids")
.nth(1)
.and_then(|tail| {
tail.split("pub async fn summarize_usage_totals_by_user_ids")
.next()
})
.expect("API-key total-token summary implementation should be present");
assert!(implementation.contains("SUM(GREATEST(COALESCE(total_tokens, 0), 0))"));
assert!(implementation.contains(")::BIGINT AS total_tokens"));
assert!(!implementation.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
}
#[test]
@@ -674,31 +691,24 @@ fn usage_counter_pending_health_does_not_scan_processed_history() {
#[test]
fn usage_sql_rebuild_matches_online_api_key_usage_semantics() {
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')"));
let sql = super::REBUILD_API_KEY_USAGE_STATS_SQL;
assert!(sql.contains("COUNT(*)::BIGINT"));
assert!(sql.contains("GREATEST(\n COALESCE(total_tokens, 0),"));
assert!(!sql.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
assert!(sql.contains("AND BTRIM(api_key_id) <> ''"));
assert!(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("));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL.contains("total_tokens,"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("AND BTRIM(provider_api_key_id) <> ''"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("WHEN status NOT IN ('pending', 'streaming')"));
assert!(super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL
.contains("WHEN status IN ('pending', 'streaming') THEN 0"));
let sql = super::REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL;
assert!(sql.contains("COUNT(*)::BIGINT"));
assert!(sql.contains("NULLIF(BTRIM(error_message), '') IS NULL"));
assert!(sql.contains("GREATEST(\n COALESCE(total_tokens, 0),"));
assert!(!sql.contains("COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)"));
assert!(sql.contains("AND BTRIM(provider_api_key_id) <> ''"));
assert!(sql.contains("WHEN status NOT IN ('pending', 'streaming')"));
assert!(sql.contains("WHEN status IN ('pending', 'streaming') THEN 0"));
}
#[test]
@@ -885,6 +895,82 @@ fn usage_sql_summarize_usage_leaderboard_supports_daily_aggregates() {
);
}
#[test]
fn usage_sql_aggregate_reads_use_materialized_total_tokens_only_when_available() {
fn select_projections_before_table<'a>(source: &'a str, table: &str) -> Vec<&'a str> {
let marker = format!("FROM {table}\n");
let mut remaining = source;
let mut projections = Vec::new();
while let Some((before_table, after_table)) = remaining.split_once(marker.as_str()) {
projections.push(
before_table
.rsplit_once("SELECT")
.map(|(_, projection)| projection)
.unwrap_or_else(|| panic!("SELECT projection for {table} should be present")),
);
remaining = after_table;
}
projections
}
let source = include_str!("mod.rs");
let breakdown = source
.split("async fn summarize_usage_breakdown_from_daily_aggregates")
.nth(1)
.and_then(|tail| tail.split("async fn summarize_usage_breakdown_raw").next())
.expect("daily aggregate usage breakdown query should be present");
for table in [
"stats_user_daily_model",
"stats_user_daily_provider",
"stats_user_daily_api_format",
] {
assert!(breakdown.contains(table));
}
assert!(breakdown.contains("COALESCE(SUM(total_tokens), 0)::BIGINT AS total_tokens"));
assert!(!breakdown.contains(
"SUM(effective_input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens)"
));
let leaderboard = source
.split("async fn summarize_usage_leaderboard_from_daily_aggregates")
.nth(1)
.and_then(|tail| {
tail.split("pub async fn summarize_usage_leaderboard")
.next()
})
.expect("daily aggregate usage leaderboard queries should be present");
for (table, expected_queries) in [
("stats_user_daily_provider", 1),
("stats_user_daily_model", 2),
] {
let projections = select_projections_before_table(leaderboard, table);
assert_eq!(
projections.len(),
expected_queries,
"unexpected {table} query count"
);
assert!(projections.iter().all(|projection| projection
.contains("COALESCE(SUM(total_tokens), 0)::BIGINT AS total_tokens")));
}
for (table, component_expression) in [
(
"stats_user_daily",
"SUM(effective_input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens)",
),
(
"stats_daily_model",
"SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens)",
),
("stats_daily_api_key", "stats_daily_api_key.input_tokens"),
] {
let projections = select_projections_before_table(leaderboard, table);
assert_eq!(projections.len(), 1, "unexpected {table} query count");
assert!(projections[0].contains(component_expression));
assert!(!projections[0].contains("SUM(total_tokens)"));
}
}
#[test]
fn usage_sql_aggregate_usage_audits_supports_daily_model_and_provider_aggregates() {
let source = include_str!("mod.rs");
@@ -991,6 +1077,189 @@ fn usage_sql_raw_aggregates_use_canonical_billing_facts() {
assert!(!source.contains("apply_provider_api_key_codex_window_usage_delta_in_tx"));
}
#[test]
fn usage_sql_raw_token_aggregates_do_not_renormalize_canonical_billing_facts() {
let source = include_str!("mod.rs");
for (start, end, label) in [
(
"async fn summarize_usage_audits_raw",
"pub async fn summarize_usage_audits",
"usage audit summary",
),
(
"async fn summarize_usage_cache_affinity_hit_summary_raw",
"async fn summarize_usage_cache_affinity_hit_summary_from_daily_aggregates",
"cache affinity",
),
(
"async fn summarize_usage_breakdown_raw",
"pub async fn summarize_usage_breakdown",
"usage breakdown",
),
(
"async fn summarize_usage_leaderboard_raw",
"async fn summarize_usage_leaderboard_from_daily_aggregates",
"usage leaderboard",
),
(
"async fn aggregate_usage_audits_raw",
"pub async fn aggregate_usage_audits",
"usage audit aggregation",
),
(
"async fn summarize_usage_daily_heatmap_raw_from_range",
"async fn summarize_usage_daily_heatmap_from_daily_aggregates",
"usage daily heatmap",
),
] {
let body = source
.split(start)
.nth(1)
.and_then(|tail| tail.split(end).next())
.unwrap_or_else(|| panic!("{label} raw query should be present"));
assert!(body.contains("FROM usage_billing_facts"));
assert!(
!body.contains("split_part(lower("),
"{label} must not normalize canonical billing facts by API format"
);
assert!(
!body.contains("input_tokens - cache_"),
"{label} must not subtract cache tokens from canonical input tokens"
);
}
let cache_affinity = source
.split("async fn summarize_usage_cache_affinity_hit_summary_raw")
.nth(1)
.and_then(|tail| {
tail.split("async fn summarize_usage_cache_affinity_hit_summary_from_daily_aggregates")
.next()
})
.expect("cache affinity raw query should be present");
assert!(cache_affinity.contains("\"usage\".input_tokens"));
assert!(!cache_affinity.contains("\"usage\".effective_input_tokens"));
assert!(cache_affinity.contains("\"usage\".total_input_context"));
let audit_summary = source
.split("async fn summarize_usage_audits_raw")
.nth(1)
.and_then(|tail| tail.split("pub async fn summarize_usage_audits").next())
.expect("usage audit summary raw query should be present");
assert!(audit_summary.contains("SUM(GREATEST(COALESCE(\"usage\".cache_creation_input_tokens"));
assert!(!audit_summary.contains("WHEN COALESCE(\"usage\".cache_creation_input_tokens, 0) = 0"));
for (start, end, label) in [
(
"async fn summarize_usage_breakdown_raw",
"pub async fn summarize_usage_breakdown",
"usage breakdown",
),
(
"async fn aggregate_usage_audits_raw",
"pub async fn aggregate_usage_audits",
"usage audit aggregation",
),
] {
let body = source
.split(start)
.nth(1)
.and_then(|tail| tail.split(end).next())
.unwrap_or_else(|| panic!("{label} raw query should be present"));
assert!(body.contains("\"usage\".effective_input_tokens"));
assert!(body.contains("\"usage\".total_input_context"));
}
let leaderboard = source
.split("async fn summarize_usage_leaderboard_raw")
.nth(1)
.and_then(|tail| {
tail.split("async fn summarize_usage_leaderboard_from_daily_aggregates")
.next()
})
.expect("usage leaderboard raw query should be present");
assert!(leaderboard.contains("SUM(GREATEST(COALESCE(\"usage\".total_tokens"));
let daily_heatmap = source
.split("async fn summarize_usage_daily_heatmap_raw_from_range")
.nth(1)
.and_then(|tail| {
tail.split("async fn summarize_usage_daily_heatmap_from_daily_aggregates")
.next()
})
.expect("usage daily heatmap raw query should be present");
assert!(daily_heatmap.contains("SUM(GREATEST(COALESCE(\"usage\".total_tokens"));
assert!(!daily_heatmap.contains("\"usage\".input_tokens + \"usage\".output_tokens"));
}
#[test]
fn usage_sql_canonical_openai_cache_case_preserves_effective_and_total_tokens() {
let upstream_input_tokens = 166_103_i64;
let effective_input_tokens = 1_495_i64;
let cache_creation_tokens = 0_i64;
let cache_read_tokens = 164_608_i64;
let output_tokens = 94_i64;
let total_input_context = 166_103_i64;
let total_tokens = 166_197_i64;
assert_eq!(
effective_input_tokens + cache_creation_tokens + cache_read_tokens,
total_input_context
);
assert_eq!(total_input_context, upstream_input_tokens);
assert_eq!(
effective_input_tokens + cache_creation_tokens + cache_read_tokens + output_tokens,
total_tokens
);
let source = include_str!("mod.rs");
assert!(source.contains("\"usage\".effective_input_tokens"));
assert!(source.contains("\"usage\".total_input_context"));
assert!(!source.contains("split_part(lower("));
let aggregate_audit_summary = source
.split("async fn summarize_usage_audits_from_daily_aggregates")
.nth(1)
.and_then(|tail| tail.split("async fn summarize_usage_audits_raw").next())
.expect("daily aggregate usage audit summary should be present");
assert_eq!(
aggregate_audit_summary
.matches("WHEN effective_input_tokens = 0 AND total_input_context = 0")
.count(),
2
);
assert_eq!(
aggregate_audit_summary
.matches("+ output_tokens + cache_creation_tokens + cache_read_tokens")
.count(),
2
);
assert!(!aggregate_audit_summary.contains("SUM(input_tokens + output_tokens)"));
let raw_audit_summary = source
.split("async fn summarize_usage_audits_raw")
.nth(1)
.and_then(|tail| tail.split("pub async fn summarize_usage_audits").next())
.expect("raw usage audit summary should be present");
assert!(raw_audit_summary.contains("SUM(GREATEST(COALESCE(\"usage\".total_tokens, 0), 0))"));
}
#[test]
fn usage_openai_total_input_context_includes_cache_creation_tokens() {
let effective_input_tokens =
usage_effective_input_tokens(Some(1_000), Some(100), Some(400), "openai");
assert_eq!(effective_input_tokens, Some(500));
assert_eq!(
usage_total_input_context(
Some(1_000),
effective_input_tokens,
Some(100),
Some(400),
"openai",
),
Some(1_000)
);
}
#[test]
fn dashboard_covering_index_keeps_canonical_billing_reads_heap_light() {
let migration = include_str!(
@@ -1134,6 +1403,22 @@ fn usage_billing_facts_projects_upstream_stream_mode() {
);
}
#[test]
fn usage_billing_facts_total_tokens_uses_canonical_effective_input() {
let migration =
include_str!("../../migrations/20260716000000_fix_usage_billing_facts_total_tokens.sql");
let bootstrap =
include_str!("../../../../runtime/schema/bootstrap/postgres/100_usage_capture.sql");
for sql in [migration, bootstrap] {
assert!(sql.contains("WHEN settlement.billing_effective_input_tokens IS NOT NULL"));
assert!(sql.contains("GREATEST(settlement.billing_effective_input_tokens, 0)"));
assert!(sql.contains("WHEN settlement.billing_total_input_context IS NOT NULL"));
assert!(sql.contains("NULLIF(GREATEST(COALESCE(usage_rows.total_tokens, 0), 0), 0)"));
assert!(!sql.contains("THEN COALESCE(settlement.billing_input_tokens, 0)"));
}
}
#[test]
fn usage_sql_reads_http_audits_for_single_record_fetches() {
assert!(super::FIND_BY_REQUEST_ID_SQL.contains("LEFT JOIN usage_http_audits"));
@@ -1157,11 +1442,24 @@ fn usage_sql_reads_settlement_snapshots_for_single_record_fetches() {
assert!(super::FIND_BY_ID_SQL.contains("LEFT JOIN usage_settlement_snapshots"));
assert!(super::FIND_BY_REQUEST_ID_SQL.contains("settlement_billing_snapshot_schema_version"));
assert!(super::FIND_BY_ID_SQL.contains("settlement_price_per_request"));
for sql in [super::FIND_BY_REQUEST_ID_SQL, super::FIND_BY_ID_SQL] {
for sql in [
super::FIND_BY_REQUEST_ID_SQL,
super::FIND_BY_ID_SQL,
super::LIST_USAGE_AUDITS_PREFIX,
super::LIST_RECENT_USAGE_AUDITS_PREFIX,
] {
assert!(sql.contains("CAST(\"usage\".input_tokens AS INTEGER) AS input_tokens"));
assert!(sql.contains(
"usage_settlement_snapshots.billing_input_tokens AS settlement_billing_input_tokens"
));
assert!(sql.contains(
"WHEN usage_settlement_snapshots.billing_effective_input_tokens IS NOT NULL"
));
assert!(
sql.contains("WHEN usage_settlement_snapshots.billing_total_input_context IS NOT NULL")
);
assert!(sql.contains("NULLIF(GREATEST(COALESCE(\"usage\".total_tokens, 0), 0), 0)"));
assert!(!sql.contains("THEN COALESCE(usage_settlement_snapshots.billing_input_tokens, 0)"));
assert!(sql.contains("usage_settlement_snapshots.billing_cache_creation_5m_tokens"));
assert!(sql.contains(
"CAST(usage_settlement_snapshots.billing_total_cost_usd AS DOUBLE PRECISION)"