rei-forks/fred
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

collab: Add has_extended_trial to LlmTokenClaims (#29622)

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

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-04-29 16:22:51 -04:00 committed by GitHub
parent 15a83b5a10
commit 24e47de02a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 12 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
crates/collab/src/llm.rs
View File

@ -5,6 +5,8 @@ use crate::Cents;
pub use token::*; 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 /// The maximum monthly spending an individual user can reach on the free tier
/// before they have to pay. /// before they have to pay.
pub const FREE_TIER_MONTHLY_SPENDING_LIMIT: Cents = Cents::from_dollars(10); pub const FREE_TIER_MONTHLY_SPENDING_LIMIT: Cents = Cents::from_dollars(10);

9
crates/collab/src/llm/token.rs
View File

@ -1,7 +1,9 @@
use crate::Cents; use crate::Cents;
use crate::db::billing_subscription::SubscriptionKind; use crate::db::billing_subscription::SubscriptionKind;
use crate::db::{billing_subscription, user}; 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 crate::{Config, db::billing_preference};
use anyhow::{Result, anyhow}; use anyhow::{Result, anyhow};
use chrono::{NaiveDateTime, Utc}; use chrono::{NaiveDateTime, Utc};
@ -32,6 +34,8 @@ pub struct LlmTokenClaims {
pub custom_llm_monthly_allowance_in_cents: Option<u32>, pub custom_llm_monthly_allowance_in_cents: Option<u32>,
pub plan: Plan, pub plan: Plan,
#[serde(default)] #[serde(default)]
pub has_extended_trial: bool,
#[serde(default)]
pub subscription_period: Option<(NaiveDateTime, NaiveDateTime)>, pub subscription_period: Option<(NaiveDateTime, NaiveDateTime)>,
#[serde(default)] #[serde(default)]
pub enable_model_request_overages: bool, pub enable_model_request_overages: bool,
@ -94,6 +98,9 @@ impl LlmTokenClaims {
SubscriptionKind::ZedPro => Plan::ZedPro, SubscriptionKind::ZedPro => Plan::ZedPro,
SubscriptionKind::ZedProTrial => Plan::ZedProTrial, SubscriptionKind::ZedProTrial => Plan::ZedProTrial,
}), }),
has_extended_trial: feature_flags
.iter()
.any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG),
subscription_period: maybe!({ subscription_period: maybe!({
let subscription = subscription?; let subscription = subscription?;
let period_start_at = subscription.current_period_start_at()?; let period_start_at = subscription.current_period_start_at()?;

5
crates/collab/src/stripe_billing.rs
View File

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