collab: Defer account age check to `POST /completion` endpoint (#26956)

Marshall Bowers created 9 months ago

This PR defers the account age check to the `POST /completion` endpoint
instead of doing it when an LLM token is generated.

This will allow us to lift the account age restriction for using Edit
Prediction.

Note: We're still temporarily performing the account age check when
issuing the LLM token until this change is deployed and the LLM tokens
have had a chance to cycle.

Release Notes:

- N/A

Change summary

crates/collab/src/llm.rs       | 10 ++++++++++
crates/collab/src/llm/token.rs | 14 +++++++++++++-
crates/collab/src/rpc.rs       |  4 +++-
3 files changed, 26 insertions(+), 2 deletions(-)

Detailed changes

crates/collab/src/llm.rs

@@ -5,6 +5,7 @@ mod token;
 use crate::api::events::SnowflakeRow;
 use crate::api::CloudflareIpCountryHeader;
 use crate::build_kinesis_client;
+use crate::rpc::MIN_ACCOUNT_AGE_FOR_LLM_USE;
 use crate::{db::UserId, executor::Executor, Cents, Config, Error, Result};
 use anyhow::{anyhow, Context as _};
 use authorization::authorize_access_to_language_model;
@@ -217,6 +218,15 @@ async fn perform_completion(
         params.model,
     );
 
+    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"))?
+            }
+        }
+    }
+
     authorize_access_to_language_model(
         &state.config,
         &claims,

crates/collab/src/llm/token.rs

@@ -3,7 +3,7 @@ use crate::llm::{DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT};
 use crate::Cents;
 use crate::{db::billing_preference, Config};
 use anyhow::{anyhow, Result};
-use chrono::Utc;
+use chrono::{NaiveDateTime, Utc};
 use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation};
 use serde::{Deserialize, Serialize};
 use std::time::Duration;
@@ -20,9 +20,17 @@ 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 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,
@@ -57,10 +65,14 @@ impl LlmTokenClaims {
             system_id,
             metrics_id: user.metrics_id,
             github_user_login: user.github_login.clone(),
+            account_created_at: Some(user.account_created_at()),
             is_staff,
             has_llm_closed_beta_feature_flag: feature_flags
                 .iter()
                 .any(|flag| flag == "llm-closed-beta"),
+            bypass_account_age_check: feature_flags
+                .iter()
+                .any(|flag| flag == "bypass-account-age-check"),
             has_predict_edits_feature_flag: feature_flags
                 .iter()
                 .any(|flag| flag == "predict-edits"),

crates/collab/src/rpc.rs

@@ -4036,7 +4036,7 @@ async fn accept_terms_of_service(
 }
 
 /// The minimum account age an account must have in order to use the LLM service.
-const MIN_ACCOUNT_AGE_FOR_LLM_USE: chrono::Duration = chrono::Duration::days(30);
+pub const MIN_ACCOUNT_AGE_FOR_LLM_USE: chrono::Duration = chrono::Duration::days(30);
 
 async fn get_llm_api_token(
     _request: proto::GetLlmToken,
@@ -4066,6 +4066,8 @@ 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 {