collab: Add `has_extended_trial` to `LlmTokenClaims` (#29622)

Marshall Bowers created 7 months ago

This PR adds the `has_extended_trial` field to the LLM token claims.

Release Notes:

- N/A

Change summary

crates/collab/src/llm.rs            | 2 ++
crates/collab/src/llm/token.rs      | 9 ++++++++-
crates/collab/src/stripe_billing.rs | 5 ++---
3 files changed, 12 insertions(+), 4 deletions(-)

Detailed changes

crates/collab/src/llm.rs

@@ -5,6 +5,8 @@ use crate::Cents;
 
 pub use token::*;
 
+pub const AGENT_EXTENDED_TRIAL_FEATURE_FLAG: &str = "agent-extended-trial";
+
 /// The maximum monthly spending an individual user can reach on the free tier
 /// before they have to pay.
 pub const FREE_TIER_MONTHLY_SPENDING_LIMIT: Cents = Cents::from_dollars(10);

crates/collab/src/llm/token.rs

@@ -1,7 +1,9 @@
 use crate::Cents;
 use crate::db::billing_subscription::SubscriptionKind;
 use crate::db::{billing_subscription, user};
-use crate::llm::{DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT};
+use crate::llm::{
+    AGENT_EXTENDED_TRIAL_FEATURE_FLAG, DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT,
+};
 use crate::{Config, db::billing_preference};
 use anyhow::{Result, anyhow};
 use chrono::{NaiveDateTime, Utc};
@@ -32,6 +34,8 @@ pub struct LlmTokenClaims {
     pub custom_llm_monthly_allowance_in_cents: Option<u32>,
     pub plan: Plan,
     #[serde(default)]
+    pub has_extended_trial: bool,
+    #[serde(default)]
     pub subscription_period: Option<(NaiveDateTime, NaiveDateTime)>,
     #[serde(default)]
     pub enable_model_request_overages: bool,
@@ -94,6 +98,9 @@ impl LlmTokenClaims {
                     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()?;

crates/collab/src/stripe_billing.rs

@@ -1,6 +1,7 @@
 use std::sync::Arc;
 
-use crate::{Cents, Result, llm};
+use crate::llm::{self, AGENT_EXTENDED_TRIAL_FEATURE_FLAG};
+use crate::{Cents, Result};
 use anyhow::{Context as _, anyhow};
 use chrono::{Datelike, Utc};
 use collections::HashMap;
@@ -492,8 +493,6 @@ impl StripeBilling {
         feature_flags: Vec<String>,
         success_url: &str,
     ) -> Result<String> {
-        const AGENT_EXTENDED_TRIAL_FEATURE_FLAG: &str = "agent-extended-trial";
-
         let eligible_for_extended_trial = feature_flags
             .iter()
             .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG);