collab: Make account age-related fields required in `LlmTokenClaims` (#26959)

Marshall Bowers created 9 months ago

This PR makes the account age-related fields required in
`LlmTokenClaims`.

We've also removed the account age check from the LLM token issuance
endpoint, instead having it solely be enforced in the `POST /completion`
endpoint.

This change will be safe to deploy at ~8:01PM EDT.

Release Notes:

- N/A

Change summary

crates/collab/src/llm.rs       |  6 ++----
crates/collab/src/llm/token.rs | 11 ++---------
crates/collab/src/rpc.rs       | 13 -------------
3 files changed, 4 insertions(+), 26 deletions(-)

Detailed changes

crates/collab/src/llm.rs

@@ -220,10 +220,8 @@ async fn perform_completion(
 
     let bypass_account_age_check = claims.has_llm_subscription || claims.bypass_account_age_check;
     if !bypass_account_age_check {
-        if let Some(account_created_at) = claims.account_created_at {
-            if Utc::now().naive_utc() - account_created_at < MIN_ACCOUNT_AGE_FOR_LLM_USE {
-                Err(anyhow!("account too young"))?
-            }
+        if Utc::now().naive_utc() - claims.account_created_at < MIN_ACCOUNT_AGE_FOR_LLM_USE {
+            Err(anyhow!("account too young"))?
         }
     }
 

crates/collab/src/llm/token.rs

@@ -20,17 +20,10 @@ pub struct LlmTokenClaims {
     pub system_id: Option<String>,
     pub metrics_id: Uuid,
     pub github_user_login: String,
-    // This field is temporarily optional so it can be added
-    // in a backwards-compatible way. We can make it required
-    // once all of the LLM tokens have cycled (~1 hour after
-    // this change has been deployed).
-    #[serde(default)]
-    pub account_created_at: Option<NaiveDateTime>,
+    pub account_created_at: NaiveDateTime,
     pub is_staff: bool,
     pub has_llm_closed_beta_feature_flag: bool,
-    #[serde(default)]
     pub bypass_account_age_check: bool,
-    #[serde(default)]
     pub has_predict_edits_feature_flag: bool,
     pub has_llm_subscription: bool,
     pub max_monthly_spend_in_cents: u32,
@@ -65,7 +58,7 @@ impl LlmTokenClaims {
             system_id,
             metrics_id: user.metrics_id,
             github_user_login: user.github_login.clone(),
-            account_created_at: Some(user.account_created_at()),
+            account_created_at: user.account_created_at(),
             is_staff,
             has_llm_closed_beta_feature_flag: feature_flags
                 .iter()

crates/collab/src/rpc.rs

@@ -4047,8 +4047,6 @@ 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_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"))?
@@ -4065,17 +4063,6 @@ async fn get_llm_api_token(
     }
 
     let has_llm_subscription = session.has_llm_subscription(&db).await?;
-
-    // This check is now handled in the `perform_completion` endpoint. We can remove the check here once the tokens have
-    // had ~1 hour to cycle.
-    let bypass_account_age_check =
-        has_llm_subscription || has_bypass_account_age_check_feature_flag;
-    if !bypass_account_age_check {
-        if Utc::now().naive_utc() - user.account_created_at() < MIN_ACCOUNT_AGE_FOR_LLM_USE {
-            Err(anyhow!("account too young"))?
-        }
-    }
-
     let billing_preferences = db.get_billing_preferences(user.id).await?;
 
     let token = LlmTokenClaims::create(