collab: Set billing-related fields for Zed staff (#29887)

Marshall Bowers created 9 months ago

This PR sets the billing-related fields in the LLM token claims for Zed
staff.

Staff members are automatically in the Zed Pro plan with a subscription
periods that spans the entirety of each month.

Release Notes:

- N/A

Change summary

crates/collab/src/llm/token.rs | 55 ++++++++++++++++++++++++++---------
crates/collab/src/rpc.rs       |  4 ++
2 files changed, 44 insertions(+), 15 deletions(-)

Detailed changes

crates/collab/src/llm/token.rs

@@ -6,7 +6,7 @@ use crate::llm::{
 };
 use crate::{Config, db::billing_preference};
 use anyhow::{Result, anyhow};
-use chrono::{NaiveDateTime, Utc};
+use chrono::{Datelike, NaiveDate, NaiveDateTime, Utc};
 use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation};
 use serde::{Deserialize, Serialize};
 use std::time::Duration;
@@ -96,24 +96,49 @@ impl LlmTokenClaims {
                 .map(|allowance| allowance as u32),
             use_new_billing: feature_flags.iter().any(|flag| flag == "new-billing"),
             use_llm_request_queue: feature_flags.iter().any(|flag| flag == "llm-request-queue"),
-            plan: subscription
-                .as_ref()
-                .and_then(|subscription| subscription.kind)
-                .map_or(Plan::Free, |kind| match kind {
-                    SubscriptionKind::ZedFree => Plan::Free,
-                    SubscriptionKind::ZedPro => Plan::ZedPro,
-                    SubscriptionKind::ZedProTrial => Plan::ZedProTrial,
-                }),
+            plan: if is_staff {
+                Plan::ZedPro
+            } else {
+                subscription
+                    .as_ref()
+                    .and_then(|subscription| subscription.kind)
+                    .map_or(Plan::Free, |kind| match kind {
+                        SubscriptionKind::ZedFree => Plan::Free,
+                        SubscriptionKind::ZedPro => Plan::ZedPro,
+                        SubscriptionKind::ZedProTrial => Plan::ZedProTrial,
+                    })
+            },
             has_extended_trial: feature_flags
                 .iter()
                 .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG),
-            subscription_period: maybe!({
-                let subscription = subscription?;
-                let period_start_at = subscription.current_period_start_at()?;
-                let period_end_at = subscription.current_period_end_at()?;
+            subscription_period: if is_staff {
+                maybe!({
+                    let now = Utc::now();
+                    let year = now.year();
+                    let month = now.month();
+
+                    let first_day_of_this_month =
+                        NaiveDate::from_ymd_opt(year, month, 1)?.and_hms_opt(0, 0, 0)?;
+
+                    let next_month = if month == 12 { 1 } else { month + 1 };
+                    let next_month_year = if month == 12 { year + 1 } else { year };
+                    let first_day_of_next_month =
+                        NaiveDate::from_ymd_opt(next_month_year, next_month, 1)?
+                            .and_hms_opt(23, 59, 59)?;
+
+                    let last_day_of_this_month = first_day_of_next_month - chrono::Days::new(1);
+
+                    Some((first_day_of_this_month, last_day_of_this_month))
+                })
+            } else {
+                maybe!({
+                    let subscription = subscription?;
+                    let period_start_at = subscription.current_period_start_at()?;
+                    let period_end_at = subscription.current_period_end_at()?;
 
-                Some((period_start_at.naive_utc(), period_end_at.naive_utc()))
-            }),
+                    Some((period_start_at.naive_utc(), period_end_at.naive_utc()))
+                })
+            },
             enable_model_request_overages: billing_preferences
                 .as_ref()
                 .map_or(false, |preferences| {

crates/collab/src/rpc.rs

@@ -180,6 +180,10 @@ impl Session {
     }
 
     pub async fn current_plan(&self, db: &MutexGuard<'_, DbHandle>) -> anyhow::Result<proto::Plan> {
+        if self.is_staff() {
+            return Ok(proto::Plan::ZedPro);
+        }
+
         let user_id = self.user_id();
 
         let subscription = db.get_active_billing_subscription(user_id).await?;