collab: Don't try to transfer usage when a Zed Pro trial is canceled (#29843)

This PR fixes an issue where we would erroneously try to transfer
existing subscription usage when a Zed Pro trial was canceled.

Release Notes:

- N/A

crates/collab/src/api/billing.rs

@@ -1040,6 +1040,7 @@ async fn handle_customer_subscription_event(
                 billing_customer.user_id,
                 &existing_subscription,
                 subscription_kind,
+                subscription.status.into(),
                 new_period_start_at,
                 new_period_end_at,
             )

crates/collab/src/llm/db/queries/subscription_usages.rs

@@ -1,7 +1,7 @@
 use chrono::Timelike;
 use time::PrimitiveDateTime;
 
-use crate::db::billing_subscription::SubscriptionKind;
+use crate::db::billing_subscription::{StripeSubscriptionStatus, SubscriptionKind};
 use crate::db::{UserId, billing_subscription};
 
 use super::*;
@@ -120,12 +120,13 @@ impl LlmDatabase {
         user_id: UserId,
         existing_subscription: &billing_subscription::Model,
         new_subscription_kind: Option<SubscriptionKind>,
+        new_subscription_status: StripeSubscriptionStatus,
         new_period_start_at: DateTimeUtc,
         new_period_end_at: DateTimeUtc,
     ) -> Result<Option<subscription_usage::Model>> {
         self.transaction(|tx| async move {
-            match existing_subscription.kind {
+            match (existing_subscription.kind, new_subscription_status) {
-                Some(SubscriptionKind::ZedProTrial) => {
+                (Some(SubscriptionKind::ZedProTrial), StripeSubscriptionStatus::Active) => {
                     let trial_period_start_at = existing_subscription
                         .current_period_start_at()
                         .ok_or_else(|| anyhow!("No trial subscription period start"))?;

crates/collab/src/llm/db/tests/subscription_usage_tests.rs

@@ -1,7 +1,7 @@
 use chrono::{Duration, Utc};
 use pretty_assertions::assert_eq;
 
-use crate::db::billing_subscription::SubscriptionKind;
+use crate::db::billing_subscription::{StripeSubscriptionStatus, SubscriptionKind};
 use crate::db::{UserId, billing_subscription};
 use crate::llm::db::LlmDatabase;
 use crate::test_llm_db;
@@ -12,58 +12,108 @@ test_llm_db!(
 );
 
 async fn test_transfer_existing_subscription_usage(db: &mut LlmDatabase) {
-    let user_id = UserId(1);
+    // Test when an existing Zed Pro trial subscription is upgraded to Zed Pro.
+    {
+        let user_id = UserId(1);
 
-    let now = Utc::now();
+        let now = Utc::now();
 
-    let trial_period_start_at = now - Duration::days(14);
+        let trial_period_start_at = now - Duration::days(14);
-    let trial_period_end_at = now;
+        let trial_period_end_at = now;
 
-    let new_period_start_at = now;
+        let new_period_start_at = now;
-    let new_period_end_at = now + Duration::days(30);
+        let new_period_end_at = now + Duration::days(30);
 
-    let existing_subscription = billing_subscription::Model {
+        let existing_subscription = billing_subscription::Model {
-        kind: Some(SubscriptionKind::ZedProTrial),
+            kind: Some(SubscriptionKind::ZedProTrial),
-        stripe_current_period_start: Some(trial_period_start_at.timestamp()),
+            stripe_current_period_start: Some(trial_period_start_at.timestamp()),
-        stripe_current_period_end: Some(trial_period_end_at.timestamp()),
+            stripe_current_period_end: Some(trial_period_end_at.timestamp()),
-        ..Default::default()
+            ..Default::default()
-    };
+        };
 
-    let existing_usage = db
+        let existing_usage = db
-        .create_subscription_usage(
+            .create_subscription_usage(
-            user_id,
+                user_id,
-            trial_period_start_at,
+                trial_period_start_at,
-            trial_period_end_at,
+                trial_period_end_at,
-            SubscriptionKind::ZedProTrial,
+                SubscriptionKind::ZedProTrial,
-            25,
+                25,
-            1_000,
+                1_000,
-        )
+            )
-        .await
+            .await
-        .unwrap();
+            .unwrap();
 
-    let transferred_usage = db
+        let transferred_usage = db
-        .transfer_existing_subscription_usage(
+            .transfer_existing_subscription_usage(
-            user_id,
+                user_id,
-            &existing_subscription,
+                &existing_subscription,
-            Some(SubscriptionKind::ZedPro),
+                Some(SubscriptionKind::ZedPro),
-            new_period_start_at,
+                StripeSubscriptionStatus::Active,
-            new_period_end_at,
+                new_period_start_at,
-        )
+                new_period_end_at,
-        .await
+            )
-        .unwrap();
+            .await
+            .unwrap();
 
-    assert!(
+        assert!(
-        transferred_usage.is_some(),
+            transferred_usage.is_some(),
-        "subscription usage not transferred successfully"
+            "subscription usage not transferred successfully"
-    );
+        );
-    let transferred_usage = transferred_usage.unwrap();
+        let transferred_usage = transferred_usage.unwrap();
 
-    assert_eq!(
+        assert_eq!(
-        transferred_usage.model_requests,
+            transferred_usage.model_requests,
-        existing_usage.model_requests
+            existing_usage.model_requests
-    );
+        );
-    assert_eq!(
+        assert_eq!(
-        transferred_usage.edit_predictions,
+            transferred_usage.edit_predictions,
-        existing_usage.edit_predictions
+            existing_usage.edit_predictions
-    );
+        );
+    }
+
+    // Test when an existing Zed Pro trial subscription is canceled.
+    {
+        let user_id = UserId(2);
+
+        let now = Utc::now();
+
+        let trial_period_start_at = now - Duration::days(14);
+        let trial_period_end_at = now;
+
+        let existing_subscription = billing_subscription::Model {
+            kind: Some(SubscriptionKind::ZedProTrial),
+            stripe_current_period_start: Some(trial_period_start_at.timestamp()),
+            stripe_current_period_end: Some(trial_period_end_at.timestamp()),
+            ..Default::default()
+        };
+
+        let _existing_usage = db
+            .create_subscription_usage(
+                user_id,
+                trial_period_start_at,
+                trial_period_end_at,
+                SubscriptionKind::ZedProTrial,
+                25,
+                1_000,
+            )
+            .await
+            .unwrap();
+
+        let transferred_usage = db
+            .transfer_existing_subscription_usage(
+                user_id,
+                &existing_subscription,
+                Some(SubscriptionKind::ZedPro),
+                StripeSubscriptionStatus::Canceled,
+                trial_period_start_at,
+                trial_period_end_at,
+            )
+            .await
+            .unwrap();
+
+        assert!(
+            transferred_usage.is_none(),
+            "subscription usage was transferred when it should not have been"
+        );
+    }
 }