collab: Don't try to sync usage to Stripe for staff users (#29892)

Marshall Bowers created 7 months ago

This PR makes it so we don't try to sync billing usage to Stripe for
staff users.

Release Notes:

- N/A

Change summary

crates/collab/src/api/billing.rs      | 10 ++++++++++
crates/collab/src/db/queries/users.rs | 12 ++++++++++++
2 files changed, 22 insertions(+)

Detailed changes

crates/collab/src/api/billing.rs

@@ -1496,9 +1496,19 @@ async fn sync_model_request_usage_with_stripe(
     llm_db: &Arc<LlmDatabase>,
     stripe_billing: &Arc<StripeBilling>,
 ) -> anyhow::Result<()> {
+    let staff_users = app.db.get_staff_users().await?;
+    let staff_user_ids = staff_users
+        .iter()
+        .map(|user| user.id)
+        .collect::<HashSet<UserId>>();
+
     let usage_meters = llm_db
         .get_current_subscription_usage_meters(Utc::now())
         .await?;
+    let usage_meters = usage_meters
+        .into_iter()
+        .filter(|(_, usage)| !staff_user_ids.contains(&usage.user_id))
+        .collect::<Vec<_>>();
     let user_ids = usage_meters
         .iter()
         .map(|(_, usage)| usage.user_id)

crates/collab/src/db/queries/users.rs

@@ -65,6 +65,18 @@ impl Database {
         .await
     }
 
+    /// Returns all users flagged as staff.
+    pub async fn get_staff_users(&self) -> Result<Vec<user::Model>> {
+        self.transaction(|tx| async {
+            let tx = tx;
+            Ok(user::Entity::find()
+                .filter(user::Column::Admin.eq(true))
+                .all(&*tx)
+                .await?)
+        })
+        .await
+    }
+
     /// Returns a user by email address. There are no access checks here, so this should only be used internally.
     pub async fn get_user_by_email(&self, email: &str) -> Result<Option<User>> {
         self.transaction(|tx| async move {