Subtract free tier usage when calculating monthly spend

zed-industries · Oct 17, 2024 · f9e717d · f9e717d
1 parent 6ba6dbc
commit f9e717d
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 50 deletions.
diff --git a/crates/collab/src/cents.rs b/crates/collab/src/cents.rs
@@ -10,6 +10,8 @@
     Copy,
     derive_more::Add,
     derive_more::AddAssign,
+    derive_more::Sub,
+    derive_more::SubAssign,
 )]
 pub struct Cents(pub u32);
 

diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs
@@ -469,7 +469,9 @@ async fn check_usage_limit(
                 ));
             }
 
-            if usage.spending_this_month >= Cents(claims.max_monthly_spend_in_cents) {
+            if (usage.spending_this_month - FREE_TIER_MONTHLY_SPENDING_LIMIT)
+                >= Cents(claims.max_monthly_spend_in_cents)
+            {
                 return Err(Error::Http(
                     StatusCode::FORBIDDEN,
                     "Maximum spending limit reached for this month.".to_string(),

diff --git a/crates/collab/src/llm/db/queries/usages.rs b/crates/collab/src/llm/db/queries/usages.rs
@@ -412,7 +412,7 @@ impl LlmDatabase {
             if !is_staff
                 && spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT
                 && has_llm_subscription
-                && spending_this_month <= max_monthly_spend
+                && (spending_this_month - FREE_TIER_MONTHLY_SPENDING_LIMIT) <= max_monthly_spend
             {
                 billing_event::ActiveModel {
                     id: ActiveValue::not_set(),

diff --git a/crates/collab/src/llm/db/tests/billing_tests.rs b/crates/collab/src/llm/db/tests/billing_tests.rs
@@ -1,10 +1,7 @@
 use crate::{
     db::UserId,
     llm::{
-        db::{
-            queries::{providers::ModelParams, usages::Usage},
-            LlmDatabase, TokenUsage,
-        },
+        db::{queries::providers::ModelParams, LlmDatabase, TokenUsage},
         FREE_TIER_MONTHLY_SPENDING_LIMIT,
     },
     test_llm_db, Cents,
@@ -76,29 +73,9 @@ async fn test_billing_limit_exceeded(db: &mut LlmDatabase) {
 
     // Verify the recorded usage and spending
     let recorded_usage = db.get_usage(user_id, provider, model, now).await.unwrap();
-
     // Verify that we exceeded the free tier usage
-    assert!(
-        recorded_usage.spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT,
-        "Expected spending to exceed free tier limit"
-    );
-
-    assert_eq!(
-        recorded_usage,
-        Usage {
-            requests_this_minute: 1,
-            tokens_this_minute: tokens_to_use,
-            tokens_this_day: tokens_to_use,
-            tokens_this_month: TokenUsage {
-                input: tokens_to_use,
-                input_cache_creation: 0,
-                input_cache_read: 0,
-                output: 0,
-            },
-            spending_this_month: Cents::new(1050),
-            lifetime_spending: Cents::new(1050),
-        }
-    );
+    assert_eq!(recorded_usage.spending_this_month, Cents::new(1050));
+    assert!(recorded_usage.spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT);
 
     // Verify that there is one `billing_event` record
     let billing_events = db.get_billing_events().await.unwrap();
@@ -111,7 +88,35 @@ async fn test_billing_limit_exceeded(db: &mut LlmDatabase) {
     assert_eq!(billing_event.input_cache_read_tokens, 0);
     assert_eq!(billing_event.output_tokens, 0);
 
-    let tokens_to_exceed = 20_000_000; // This will cost $1.00 more, pushing us from $10.50 to $11.50, which is over the $11 monthly maximum limit
+    // Record usage that puts us at $20.50
+    let usage_2 = TokenUsage {
+        input: 200_000_000, // This will cost $10 more, pushing us from $10.50 to $20.50,
+        input_cache_creation: 0,
+        input_cache_read: 0,
+        output: 0,
+    };
+    db.record_usage(
+        user_id,
+        false,
+        provider,
+        model,
+        usage_2,
+        true,
+        max_monthly_spend,
+        now,
+    )
+    .await
+    .unwrap();
+
+    // Verify the updated usage and spending
+    let updated_usage = db.get_usage(user_id, provider, model, now).await.unwrap();
+    assert_eq!(updated_usage.spending_this_month, Cents::new(2050));
+
+    // Verify that there are now two billing events
+    let billing_events = db.get_billing_events().await.unwrap();
+    assert_eq!(billing_events.len(), 2);
+
+    let tokens_to_exceed = 20_000_000; // This will cost $1.00 more, pushing us from $20.50 to $21.50, which is over the $11 monthly maximum limit
     let usage_exceeding = TokenUsage {
         input: tokens_to_exceed,
         input_cache_creation: 0,
@@ -132,27 +137,12 @@ async fn test_billing_limit_exceeded(db: &mut LlmDatabase) {
     )
     .await
     .unwrap();
-
-    // Verify that there is still one billing record
-    let billing_events = db.get_billing_events().await.unwrap();
-    assert_eq!(billing_events.len(), 1);
-
     // Verify the updated usage and spending
     let updated_usage = db.get_usage(user_id, provider, model, now).await.unwrap();
-    assert_eq!(
-        updated_usage,
-        Usage {
-            requests_this_minute: 2,
-            tokens_this_minute: tokens_to_use + tokens_to_exceed,
-            tokens_this_day: tokens_to_use + tokens_to_exceed,
-            tokens_this_month: TokenUsage {
-                input: tokens_to_use + tokens_to_exceed,
-                input_cache_creation: 0,
-                input_cache_read: 0,
-                output: 0,
-            },
-            spending_this_month: Cents::new(1150),
-            lifetime_spending: Cents::new(1150),
-        }
-    );
+    assert_eq!(updated_usage.spending_this_month, Cents::new(2150));
+
+    // Verify that we never exceed the user max spending for the user
+    // and avoid charging them.
+    let billing_events = db.get_billing_events().await.unwrap();
+    assert_eq!(billing_events.len(), 2);
 }