From 24e47de02a6c1c46f077d9b808446c207c348569 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <git@maxdeviant.com>
Date: Tue, 29 Apr 2025 16:22:51 -0400
Subject: [PATCH] collab: Add `has_extended_trial` to `LlmTokenClaims` (#29622)

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

Release Notes:

- N/A
---
 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(-)

diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs
index 13d503e7d4..0a5b5bbf5a 100644
--- a/crates/collab/src/llm.rs
+++ b/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);
diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs
index b08230269c..e5e8a10718 100644
--- a/crates/collab/src/llm/token.rs
+++ b/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()?;
diff --git a/crates/collab/src/stripe_billing.rs b/crates/collab/src/stripe_billing.rs
index eb6572e13d..04876708e5 100644
--- a/crates/collab/src/stripe_billing.rs
+++ b/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);