collab: Use a separate Anthropic API key for Zed staff (#16128)

This PR makes it so Zed staff can use a separate Anthropic API key for
the LLM service.

We also added an `is_staff` column to the `usages` table so that we can
exclude staff usage from the "active users" metrics that influence the
rate limits.

Release Notes:

- N/A

---------

Co-authored-by: Max <max@zed.dev>

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

@@ -108,9 +108,11 @@ impl LlmDatabase {
         .await
     }
 
+    #[allow(clippy::too_many_arguments)]
     pub async fn record_usage(
         &self,
         user_id: i32,
+        is_staff: bool,
         provider: LanguageModelProvider,
         model_name: &str,
         input_token_count: usize,
@@ -132,6 +134,7 @@ impl LlmDatabase {
             let requests_this_minute = self
                 .update_usage_for_measure(
                     user_id,
+                    is_staff,
                     model.id,
                     &usages,
                     UsageMeasure::RequestsPerMinute,
@@ -143,6 +146,7 @@ impl LlmDatabase {
             let tokens_this_minute = self
                 .update_usage_for_measure(
                     user_id,
+                    is_staff,
                     model.id,
                     &usages,
                     UsageMeasure::TokensPerMinute,
@@ -154,6 +158,7 @@ impl LlmDatabase {
             let tokens_this_day = self
                 .update_usage_for_measure(
                     user_id,
+                    is_staff,
                     model.id,
                     &usages,
                     UsageMeasure::TokensPerDay,
@@ -165,6 +170,7 @@ impl LlmDatabase {
             let input_tokens_this_month = self
                 .update_usage_for_measure(
                     user_id,
+                    is_staff,
                     model.id,
                     &usages,
                     UsageMeasure::InputTokensPerMonth,
@@ -176,6 +182,7 @@ impl LlmDatabase {
             let output_tokens_this_month = self
                 .update_usage_for_measure(
                     user_id,
+                    is_staff,
                     model.id,
                     &usages,
                     UsageMeasure::OutputTokensPerMonth,
@@ -205,7 +212,11 @@ impl LlmDatabase {
             let day_since = now - Duration::days(5);
 
             let users_in_recent_minutes = usage::Entity::find()
-                .filter(usage::Column::Timestamp.gte(minute_since.naive_utc()))
+                .filter(
+                    usage::Column::Timestamp
+                        .gte(minute_since.naive_utc())
+                        .and(usage::Column::IsStaff.eq(false)),
+                )
                 .select_only()
                 .column(usage::Column::UserId)
                 .group_by(usage::Column::UserId)
@@ -213,7 +224,11 @@ impl LlmDatabase {
                 .await? as usize;
 
             let users_in_recent_days = usage::Entity::find()
-                .filter(usage::Column::Timestamp.gte(day_since.naive_utc()))
+                .filter(
+                    usage::Column::Timestamp
+                        .gte(day_since.naive_utc())
+                        .and(usage::Column::IsStaff.eq(false)),
+                )
                 .select_only()
                 .column(usage::Column::UserId)
                 .group_by(usage::Column::UserId)
@@ -232,6 +247,7 @@ impl LlmDatabase {
     async fn update_usage_for_measure(
         &self,
         user_id: i32,
+        is_staff: bool,
         model_id: ModelId,
         usages: &[usage::Model],
         usage_measure: UsageMeasure,
@@ -267,6 +283,7 @@ impl LlmDatabase {
 
         let mut model = usage::ActiveModel {
             user_id: ActiveValue::set(user_id),
+            is_staff: ActiveValue::set(is_staff),
             model_id: ActiveValue::set(model_id),
             measure_id: ActiveValue::set(measure_id),
             timestamp: ActiveValue::set(timestamp),

crates/collab/src/llm/db/tables/usage.rs

@@ -15,6 +15,7 @@ pub struct Model {
     pub measure_id: UsageMeasureId,
     pub timestamp: DateTime,
     pub buckets: Vec<i64>,
+    pub is_staff: bool,
 }
 
 #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]

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

@@ -29,12 +29,12 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
     let user_id = 123;
 
     let now = t0;
-    db.record_usage(user_id, provider, model, 1000, 0, now)
+    db.record_usage(user_id, false, provider, model, 1000, 0, now)
         .await
         .unwrap();
 
     let now = t0 + Duration::seconds(10);
-    db.record_usage(user_id, provider, model, 2000, 0, now)
+    db.record_usage(user_id, false, provider, model, 2000, 0, now)
         .await
         .unwrap();
 
@@ -66,7 +66,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
     );
 
     let now = t0 + Duration::seconds(60);
-    db.record_usage(user_id, provider, model, 3000, 0, now)
+    db.record_usage(user_id, false, provider, model, 3000, 0, now)
         .await
         .unwrap();
 
@@ -98,7 +98,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
         }
     );
 
-    db.record_usage(user_id, provider, model, 4000, 0, now)
+    db.record_usage(user_id, false, provider, model, 4000, 0, now)
         .await
         .unwrap();