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

Rootfix usage counter outbox

Browse Source
This commit is contained in:
fawney19
2026-05-19 00:44:37 +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
crates/aether-data/src/repository/usage/postgres/mod.rs
+1068 -75
View File
File diff suppressed because it is too large Load Diff
crates/aether-data/src/repository/usage/postgres/queries/apply_api_key_usage_delta_sql.sql
+1
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
crates/aether-data/src/repository/usage/postgres/queries/apply_provider_api_key_codex_window_usage_delta_sql.sql
-170
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
crates/aether-data/src/repository/usage/postgres/queries/apply_provider_api_key_usage_delta_sql.sql
+1
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
crates/aether-data/src/repository/usage/postgres/queries/rebuild_api_key_usage_stats_sql.sql
+2 -1
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
crates/aether-data/src/repository/usage/postgres/queries/rebuild_provider_api_key_usage_stats_sql.sql
+5 -4
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
crates/aether-data/src/repository/usage/postgres/tests.rs
+30 -22
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]