perf: replace lifetime spending scans with reconciled counters

.config/nextest.toml

@@ -43,10 +43,20 @@ filter = 'package(api) & binary(e2e_all) & test(/^admin_pricing_changes::test_(r
 test-group = "e2e-db"
 threads-required = "num-test-threads"
 
+[[profile.default.overrides]]
+filter = 'package(api) & binary(e2e_all) & test(/^admin_analytics::test_admin_platform_billing_summary$/)'
+test-group = "e2e-db"
+# Exact before/after platform totals require exclusive access to the shared DB.
+threads-required = "num-test-threads"
+
 [[profile.default.overrides]]
 filter = "package(api) & binary(e2e_all)"
 test-group = "e2e-db"
 
+[[profile.default.overrides]]
+filter = "package(database) & (binary(spend_counter_backfill) | binary(counter_readers))"
+test-group = "e2e-db"
+
 [[profile.default.scripts]]
 filter = "package(api) & binary(e2e_all)"
 setup = "e2e-db-bootstrap"

.github/workflows/test.yml

@@ -199,7 +199,7 @@ jobs:
         cache-key: e2e
 
     - name: Run e2e tests
-      run: cargo nextest run --test e2e_all
+      run: cargo nextest run --test e2e_all --test spend_counter_backfill --test counter_readers
       env:
         POSTGRES_PRIMARY_APP_ID: ${{ secrets.POSTGRES_PRIMARY_APP_ID }}
         DATABASE_HOST: localhost

Dockerfile

@@ -40,7 +40,7 @@ COPY crates/ ./crates/
 COPY .cargo/ ./.cargo/
 
 # Build the application in release mode
-RUN cargo build --release --locked --bin api
+RUN cargo build --release --locked --bin api --bin backfill-spend-counters
 
 
 # Runtime stage
@@ -86,6 +86,7 @@ WORKDIR /app
 
 # Copy the built binary
 COPY --from=builder /app/target/release/api /app/api
+COPY --from=builder /app/target/release/backfill-spend-counters /app/backfill-spend-counters
 
 # Copy the migration SQL files
 RUN mkdir -p /app/crates/database/src/migrations/sql

crates/api/src/main.rs

@@ -27,6 +27,9 @@ async fn main() {
             .await
             .expect("Usage reporting index prerequisites are not satisfied");
     }
+    database::ensure_spend_counters_ready(database.pool())
+        .await
+        .expect("Spend counters are incomplete; run the spend-counter backfill before starting Cloud API");
     let auth_components = init_auth_services(database.clone(), &config);
 
     // Initialize OpenTelemetry pipeline

crates/api/tests/e2e_all/api_keys.rs

@@ -1,4 +1,9 @@
 use crate::common::*;
+use database::models::RecordUsageRequest;
+use database::repositories::{
+    OrganizationServiceUsageRepository, OrganizationUsageRepository, RecordServiceUsageRequest,
+};
+use services::usage::InferenceType;
 
 // ============================================
 // API Key Creation and Management Tests
@@ -128,76 +133,73 @@ async fn test_list_workspace_api_keys_orders_by_usage() {
         .unwrap()
         .get(0);
 
-    for (api_key_id, total_cost) in [
-        (service_spend_key_id, 100_000_000_i64),
-        (inference_spend_key_id, 300_000_000_i64),
+    drop(client);
+    let inference_repository = OrganizationUsageRepository::new(database.pool().clone());
+    let mut inference = RecordUsageRequest {
+        organization_id,
+        workspace_id,
+        api_key_id: service_spend_key_id,
+        model_id,
+        model_name: model_name.clone(),
+        input_tokens: 10,
+        output_tokens: 10,
+        input_cost: 1,
+        output_cost: 1,
+        total_cost: 100_000_000,
+        inference_type: InferenceType::ChatCompletion.as_str().to_string(),
+        ttft_ms: None,
+        avg_itl_ms: None,
+        inference_id: Some(uuid::Uuid::new_v4()),
+        provider_request_id: None,
+        stop_reason: None,
+        response_id: None,
+        image_count: None,
+        cache_read_tokens: 0,
+        cache_write_tokens: 0,
+        billing_details: None,
+        service_tier: None,
+        context_band: None,
+        served_provider_tier: None,
+        served_provider_type: None,
+        served_via_fallback: false,
+    };
+    inference_repository
+        .record_usage(inference.clone())
+        .await
+        .unwrap();
+    inference.api_key_id = inference_spend_key_id;
+    inference.total_cost = 300_000_000;
+    inference.inference_id = Some(uuid::Uuid::new_v4());
+    inference_repository.record_usage(inference).await.unwrap();
+    OrganizationServiceUsageRepository::new(database.pool().clone())
+        .record_usage(&RecordServiceUsageRequest {
+            organization_id,
+            workspace_id,
+            api_key_id: service_spend_key_id,
+            service_id,
+            quantity: 1,
+            total_cost: 400_000_000,
+            inference_id: None,
+        })
+        .await
+        .unwrap();
+
+    let client = database.pool().get().await.unwrap();
+
+    for (api_key_id, hours) in [
+        (service_spend_key_id, 3_i64),
+        (inference_spend_key_id, 2_i64),
+        (unused_key_id, 1_i64),
     ] {
         client
             .execute(
-                r#"
-                INSERT INTO organization_usage_log (
-                    id, organization_id, workspace_id, api_key_id,
-                    model_id, model_name, input_tokens, output_tokens,
-                    total_tokens, input_cost, output_cost, total_cost,
-                    inference_type, created_at
-                ) VALUES ($1, $2, $3, $4, $5, $6, 10, 10, 20, 1, 1, $7,
-                          'chat_completion', NOW())
-                "#,
-                &[
-                    &uuid::Uuid::new_v4(),
-                    &organization_id,
-                    &workspace_id,
-                    &api_key_id,
-                    &model_id,
-                    &model_name,
-                    &total_cost,
-                ],
+                "UPDATE api_keys SET created_at = NOW() + ($2::BIGINT * INTERVAL '1 hour') WHERE id = $1",
+                &[&api_key_id, &hours],
             )
             .await
             .unwrap();
     }
 
-    client
-        .execute(
-            r#"
-            INSERT INTO organization_service_usage_log (
-                id, organization_id, workspace_id, api_key_id,
-                service_id, quantity, total_cost, inference_id, created_at
-            ) VALUES ($1, $2, $3, $4, $5, 1, 400000000, NULL, NOW())
-            "#,
-            &[
-                &uuid::Uuid::new_v4(),
-                &organization_id,
-                &workspace_id,
-                &service_spend_key_id,
-                &service_id,
-            ],
-        )
-        .await
-        .unwrap();
-
-    client
-        .execute(
-            "UPDATE api_keys SET created_at = NOW() + INTERVAL '3 hours' WHERE id = $1",
-            &[&service_spend_key_id],
-        )
-        .await
-        .unwrap();
-    client
-        .execute(
-            "UPDATE api_keys SET created_at = NOW() + INTERVAL '2 hours' WHERE id = $1",
-            &[&inference_spend_key_id],
-        )
-        .await
-        .unwrap();
-    client
-        .execute(
-            "UPDATE api_keys SET created_at = NOW() + INTERVAL '1 hour' WHERE id = $1",
-            &[&unused_key_id],
-        )
-        .await
-        .unwrap();
-
     let default_response = server
         .get(format!("/v1/workspaces/{}/api-keys?limit=3", workspace.id).as_str())
         .add_header("Authorization", format!("Bearer {}", get_session_id()))

crates/api/tests/e2e_all/main.rs

@@ -87,6 +87,7 @@ mod score;
 mod serving_provider;
 mod session_logout;
 mod signature_verification;
+mod spend_counters;
 mod usage_chat_completions;
 mod usage_provider_attribution;
 mod usage_recording;