Bypass account age check when feature flag is set (#32393)

Antonio Scandurra created 6 months ago

Release Notes:

- N/A

Change summary

crates/collab/src/llm.rs       |  6 ++++++
crates/collab/src/llm/token.rs |  4 ++--
crates/collab/src/rpc.rs       | 16 ++++++++++------
3 files changed, 18 insertions(+), 8 deletions(-)

Detailed changes

crates/collab/src/llm.rs

@@ -7,6 +7,12 @@ pub use token::*;
 
 pub const AGENT_EXTENDED_TRIAL_FEATURE_FLAG: &str = "agent-extended-trial";
 
+/// The name of the feature flag that bypasses the account age check.
+pub const BYPASS_ACCOUNT_AGE_CHECK_FEATURE_FLAG: &str = "bypass-account-age-check";
+
+/// The minimum account age an account must have in order to use the LLM service.
+pub const MIN_ACCOUNT_AGE_FOR_LLM_USE: chrono::Duration = chrono::Duration::days(30);
+
 /// The default value to use for maximum spend per month if the user did not
 /// explicitly set a maximum spend.
 ///

crates/collab/src/llm/token.rs

@@ -1,6 +1,6 @@
 use crate::db::billing_subscription::SubscriptionKind;
 use crate::db::{billing_customer, billing_subscription, user};
-use crate::llm::AGENT_EXTENDED_TRIAL_FEATURE_FLAG;
+use crate::llm::{AGENT_EXTENDED_TRIAL_FEATURE_FLAG, BYPASS_ACCOUNT_AGE_CHECK_FEATURE_FLAG};
 use crate::{Config, db::billing_preference};
 use anyhow::{Context as _, Result};
 use chrono::{NaiveDateTime, Utc};
@@ -84,7 +84,7 @@ impl LlmTokenClaims {
                 .any(|flag| flag == "llm-closed-beta"),
             bypass_account_age_check: feature_flags
                 .iter()
-                .any(|flag| flag == "bypass-account-age-check"),
+                .any(|flag| flag == BYPASS_ACCOUNT_AGE_CHECK_FEATURE_FLAG),
             can_use_web_search_tool: true,
             use_llm_request_queue: feature_flags.iter().any(|flag| flag == "llm-request-queue"),
             plan,

crates/collab/src/rpc.rs

@@ -4,7 +4,10 @@ use crate::api::billing::find_or_create_billing_customer;
 use crate::api::{CloudflareIpCountryHeader, SystemIdHeader};
 use crate::db::billing_subscription::SubscriptionKind;
 use crate::llm::db::LlmDatabase;
-use crate::llm::{AGENT_EXTENDED_TRIAL_FEATURE_FLAG, LlmTokenClaims};
+use crate::llm::{
+    AGENT_EXTENDED_TRIAL_FEATURE_FLAG, BYPASS_ACCOUNT_AGE_CHECK_FEATURE_FLAG, LlmTokenClaims,
+    MIN_ACCOUNT_AGE_FOR_LLM_USE,
+};
 use crate::stripe_client::StripeCustomerId;
 use crate::{
     AppState, Error, Result, auth,
@@ -2773,8 +2776,12 @@ async fn make_update_user_plan_message(
         (None, None)
     };
 
-    let account_too_young =
-        !matches!(plan, proto::Plan::ZedPro) && user.account_age() < MIN_ACCOUNT_AGE_FOR_LLM_USE;
+    let bypass_account_age_check = feature_flags
+        .iter()
+        .any(|flag| flag == BYPASS_ACCOUNT_AGE_CHECK_FEATURE_FLAG);
+    let account_too_young = !matches!(plan, proto::Plan::ZedPro)
+        && !bypass_account_age_check
+        && user.account_age() < MIN_ACCOUNT_AGE_FOR_LLM_USE;
 
     Ok(proto::UpdateUserPlan {
         plan: plan.into(),
@@ -4075,9 +4082,6 @@ async fn accept_terms_of_service(
     Ok(())
 }
 
-/// The minimum account age an account must have in order to use the LLM service.
-pub const MIN_ACCOUNT_AGE_FOR_LLM_USE: chrono::Duration = chrono::Duration::days(30);
-
 async fn get_llm_api_token(
     _request: proto::GetLlmToken,
     response: Response<proto::GetLlmToken>,