collab: Clean up LLM token creation (#26955)

Marshall Bowers created 9 months ago

This PR cleans up the LLM token creation a bit.

We now pass in the entire list of feature flags to the
`LlmTokenClaims::create` method to prevent having a bunch of confusable
`bool` parameters.

Release Notes:

- N/A

Change summary

crates/collab/src/db/tables/user.rs | 14 ++++++++++++++
crates/collab/src/llm/token.rs      | 11 +++++++----
crates/collab/src/rpc.rs            | 15 +++++----------
3 files changed, 26 insertions(+), 14 deletions(-)

Detailed changes

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<super::access_token::Entity> for Entity {
     fn to() -> RelationDef {
         Relation::AccessToken.def()

crates/collab/src/llm/token.rs

@@ -37,8 +37,7 @@ impl LlmTokenClaims {
         user: &user::Model,
         is_staff: bool,
         billing_preferences: Option<billing_preference::Model>,
-        has_llm_closed_beta_feature_flag: bool,
-        has_predict_edits_feature_flag: bool,
+        feature_flags: &Vec<String>,
         has_llm_subscription: bool,
         plan: rpc::proto::Plan,
         system_id: Option<String>,
@@ -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| {

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(),