diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs index bf510c9c0d..acfe3ca95f 100644 --- a/crates/collab/src/llm/token.rs +++ b/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(); - Some((period_start_at.naive_utc(), period_end_at.naive_utc())) - }), + 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())) + }) + }, enable_model_request_overages: billing_preferences .as_ref() .map_or(false, |preferences| { diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs index f1fd537666..16a884bd10 100644 --- a/crates/collab/src/rpc.rs +++ b/crates/collab/src/rpc.rs @@ -180,6 +180,10 @@ impl Session { } pub async fn current_plan(&self, db: &MutexGuard<'_, DbHandle>) -> anyhow::Result { + 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?;