diff --git a/crates/collab/src/db/tables/user.rs b/crates/collab/src/db/tables/user.rs index 65f6604b298767162593f869fcbd85fd07434751..15cbc28b4459da9bdc92252b7ca304719b0287f3 100644 --- a/crates/collab/src/db/tables/user.rs +++ b/crates/collab/src/db/tables/user.rs @@ -43,6 +43,20 @@ pub enum Relation { Contributor, } +impl Model { + /// Returns the timestamp of when the user's account was created. + /// + /// This will be the earlier of the `created_at` and `github_user_created_at` timestamps. + pub fn account_created_at(&self) -> NaiveDateTime { + let mut account_created_at = self.created_at; + if let Some(github_created_at) = self.github_user_created_at { + account_created_at = account_created_at.min(github_created_at); + } + + account_created_at + } +} + impl Related for Entity { fn to() -> RelationDef { Relation::AccessToken.def() diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs index 46704e8efe4d9595639b52933aedfbaa2cdd8926..ca0d6f04203688b2ad48fa353944df7fdf2b531e 100644 --- a/crates/collab/src/llm/token.rs +++ b/crates/collab/src/llm/token.rs @@ -37,8 +37,7 @@ impl LlmTokenClaims { user: &user::Model, is_staff: bool, billing_preferences: Option, - has_llm_closed_beta_feature_flag: bool, - has_predict_edits_feature_flag: bool, + feature_flags: &Vec, has_llm_subscription: bool, plan: rpc::proto::Plan, system_id: Option, @@ -59,8 +58,12 @@ impl LlmTokenClaims { metrics_id: user.metrics_id, github_user_login: user.github_login.clone(), is_staff, - has_llm_closed_beta_feature_flag, - has_predict_edits_feature_flag, + has_llm_closed_beta_feature_flag: feature_flags + .iter() + .any(|flag| flag == "llm-closed-beta"), + has_predict_edits_feature_flag: feature_flags + .iter() + .any(|flag| flag == "predict-edits"), has_llm_subscription, max_monthly_spend_in_cents: billing_preferences .map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| { diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs index 94c82327d0fafc4de99774f12aa473f2ce9daa05..45a0d26f049f1079f1561a0740082d5770bb2b91 100644 --- a/crates/collab/src/rpc.rs +++ b/crates/collab/src/rpc.rs @@ -4047,8 +4047,8 @@ async fn get_llm_api_token( let flags = db.get_user_flags(session.user_id()).await?; let has_language_models_feature_flag = flags.iter().any(|flag| flag == "language-models"); - let has_llm_closed_beta_feature_flag = flags.iter().any(|flag| flag == "llm-closed-beta"); - let has_predict_edits_feature_flag = flags.iter().any(|flag| flag == "predict-edits"); + let has_bypass_account_age_check_feature_flag = + flags.iter().any(|flag| flag == "bypass-account-age-check"); if !session.is_staff() && !has_language_models_feature_flag { Err(anyhow!("permission denied"))? @@ -4067,13 +4067,9 @@ async fn get_llm_api_token( let has_llm_subscription = session.has_llm_subscription(&db).await?; let bypass_account_age_check = - has_llm_subscription || flags.iter().any(|flag| flag == "bypass-account-age-check"); + has_llm_subscription || has_bypass_account_age_check_feature_flag; if !bypass_account_age_check { - let mut account_created_at = user.created_at; - if let Some(github_created_at) = user.github_user_created_at { - account_created_at = account_created_at.min(github_created_at); - } - if Utc::now().naive_utc() - account_created_at < MIN_ACCOUNT_AGE_FOR_LLM_USE { + if Utc::now().naive_utc() - user.account_created_at() < MIN_ACCOUNT_AGE_FOR_LLM_USE { Err(anyhow!("account too young"))? } } @@ -4084,8 +4080,7 @@ async fn get_llm_api_token( &user, session.is_staff(), billing_preferences, - has_llm_closed_beta_feature_flag, - has_predict_edits_feature_flag, + &flags, has_llm_subscription, session.current_plan(&db).await?, session.system_id.clone(),