Merge remote-tracking branch 'origin/pr/558'

crates/aether-data/src/repository/usage/mod.rs

@@ -636,15 +636,13 @@ 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()
         .map(str::trim)
         .filter(|value| !value.is_empty())?
         .to_string();
+    let is_in_flight = matches!(usage.status.as_str(), "pending" | "streaming");
     let is_success = provider_api_key_usage_is_success(
         usage.status.as_str(),
         usage.status_code,
@@ -661,8 +659,14 @@ pub(crate) fn provider_api_key_usage_contribution(
         request_count: 1,
         success_count: i64::from(is_success),
         error_count: i64::from(is_error),
-        total_tokens: i64::try_from(usage.total_tokens).unwrap_or(i64::MAX),
-        total_cost_usd: if usage.total_cost_usd.is_finite() {
+        total_tokens: if is_in_flight {
+            0
+        } else {
+            i64::try_from(usage.total_tokens).unwrap_or(i64::MAX)
+        },
+        total_cost_usd: if is_in_flight {
+            0.0
+        } else if usage.total_cost_usd.is_finite() {
             usage.total_cost_usd.max(0.0)
         } else {
             0.0
@@ -1025,7 +1029,7 @@ mod tests {
     }
 
     #[test]
-    fn provider_api_key_usage_contribution_tracks_terminal_requests_only() {
+    fn provider_api_key_usage_contribution_counts_in_flight_requests_once() {
         let usage = StoredRequestUsageAudit::new(
             "usage-1".to_string(),
             "request-1".to_string(),
@@ -1070,11 +1074,36 @@ mod tests {
 
         let mut streaming = usage.clone();
         streaming.status = "streaming".to_string();
-        assert!(provider_api_key_usage_contribution(&streaming).is_none());
+        let streaming_contribution =
+            provider_api_key_usage_contribution(&streaming).expect("streaming should count");
+        assert_eq!(streaming_contribution.request_count, 1);
+        assert_eq!(streaming_contribution.success_count, 0);
+        assert_eq!(streaming_contribution.error_count, 0);
+        assert_eq!(streaming_contribution.total_tokens, 0);
+        assert_eq!(streaming_contribution.total_cost_usd, 0.0);
+        assert_eq!(streaming_contribution.total_response_time_ms, 0);
 
-        let mut pending = usage;
+        let mut pending = usage.clone();
         pending.status = "pending".to_string();
-        assert!(provider_api_key_usage_contribution(&pending).is_none());
+        let pending_contribution =
+            provider_api_key_usage_contribution(&pending).expect("pending should count");
+        assert_eq!(pending_contribution.request_count, 1);
+        assert_eq!(pending_contribution.success_count, 0);
+        assert_eq!(pending_contribution.error_count, 0);
+        assert_eq!(pending_contribution.total_tokens, 0);
+        assert_eq!(pending_contribution.total_cost_usd, 0.0);
+        assert_eq!(pending_contribution.total_response_time_ms, 0);
+
+        let terminal_contribution =
+            provider_api_key_usage_contribution(&usage).expect("terminal should count");
+        let delta =
+            ProviderApiKeyUsageDelta::between(&pending_contribution, &terminal_contribution);
+        assert_eq!(delta.request_count, 0);
+        assert_eq!(delta.success_count, 1);
+        assert_eq!(delta.error_count, 0);
+        assert_eq!(delta.total_tokens, 20);
+        assert_eq!(delta.total_cost_usd, 0.25);
+        assert_eq!(delta.total_response_time_ms, 120);
     }
 
     #[test]

crates/aether-data/src/repository/usage/mysql.rs

@@ -854,19 +854,28 @@ WHERE provider_api_key_id IS NOT NULL AND provider_api_key_id <> ''
             let error_message: Option<String> = row.try_get("error_message").map_sql_err()?;
             let entry = stats.entry(key_id).or_default();
             entry.request_count += 1;
-            if provider_api_key_usage_is_success(&status, status_code_u16, error_message.as_deref())
-            {
+            let is_success = provider_api_key_usage_is_success(
+                &status,
+                status_code_u16,
+                error_message.as_deref(),
+            );
+            let is_in_flight = matches!(status.as_str(), "pending" | "streaming");
+            if is_success {
                 entry.success_count += 1;
             }
             if provider_api_key_usage_is_error(&status, status_code_u16, error_message.as_deref()) {
                 entry.error_count += 1;
             }
-            entry.total_tokens += row.try_get::<i64, _>("total_tokens").map_sql_err()?;
-            entry.total_cost_usd += row.try_get::<f64, _>("total_cost_usd").map_sql_err()?;
-            entry.total_response_time_ms += row
-                .try_get::<Option<i64>, _>("response_time_ms")
-                .map_sql_err()?
-                .unwrap_or_default();
+            if !is_in_flight {
+                entry.total_tokens += row.try_get::<i64, _>("total_tokens").map_sql_err()?;
+                entry.total_cost_usd += row.try_get::<f64, _>("total_cost_usd").map_sql_err()?;
+            }
+            if is_success {
+                entry.total_response_time_ms += row
+                    .try_get::<Option<i64>, _>("response_time_ms")
+                    .map_sql_err()?
+                    .unwrap_or_default();
+            }
             entry.last_used_at = entry.last_used_at.max(
                 row.try_get::<Option<i64>, _>("updated_at_unix_secs")
                     .map_sql_err()?,

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

@@ -24,15 +24,23 @@ WITH aggregated AS (
       END
     ), 0)::BIGINT AS error_count,
     COALESCE(SUM(
-      GREATEST(
-        COALESCE(
-          total_tokens,
-          COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
-        ),
-        0
-      )::BIGINT
+      CASE
+        WHEN status IN ('pending', 'streaming') THEN 0
+        ELSE GREATEST(
+          COALESCE(
+            total_tokens,
+            COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
+          ),
+          0
+        )::BIGINT
+      END
     ), 0)::BIGINT AS total_tokens,
-    COALESCE(SUM(COALESCE(total_cost_usd, 0)), 0)::NUMERIC(20,8) AS total_cost_usd,
+    COALESCE(SUM(
+      CASE
+        WHEN status IN ('pending', 'streaming') THEN 0
+        ELSE COALESCE(total_cost_usd, 0)
+      END
+    ), 0)::NUMERIC(20,8) AS total_cost_usd,
     COALESCE(SUM(
       CASE
         WHEN status IN ('completed', 'success', 'ok', 'billed', 'settled')
@@ -47,7 +55,6 @@ WITH aggregated AS (
   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/sqlite.rs

@@ -3749,8 +3749,14 @@ SELECT
   COUNT(*) AS request_count,
   COALESCE(SUM({success_flag_expr}), 0) AS success_count,
   COALESCE(SUM({error_flag_expr}), 0) AS error_count,
-  COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS total_tokens,
-  COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd,
+  COALESCE(SUM(CASE
+    WHEN status IN ('pending', 'streaming') THEN 0
+    ELSE MAX(COALESCE(total_tokens, 0), 0)
+  END), 0) AS total_tokens,
+  COALESCE(SUM(CASE
+    WHEN status IN ('pending', 'streaming') THEN 0
+    ELSE COALESCE(CAST(total_cost_usd AS REAL), 0)
+  END), 0) AS total_cost_usd,
   COALESCE(SUM(CASE
     WHEN {success_flag_expr} = 1 AND response_time_ms IS NOT NULL
     THEN MAX(COALESCE(response_time_ms, 0), 0)