From cd5d7e82d0037dfb4dd5abb83dfac22a2dcc0f70 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Mon, 17 Mar 2025 19:54:44 -0400 Subject: [PATCH] collab: Make account age-related fields required in `LlmTokenClaims` (#26959) 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 --- 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(-) diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs index 8df92145ce..aa513c099f 100644 --- a/crates/collab/src/llm.rs +++ b/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"))? } } diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs index d73cdfd3f1..068e5f333e 100644 --- a/crates/collab/src/llm/token.rs +++ b/crates/collab/src/llm/token.rs @@ -20,17 +20,10 @@ pub struct LlmTokenClaims { pub system_id: Option, 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, + 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() diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs index 8e63a80c55..1d6c4bb01f 100644 --- a/crates/collab/src/rpc.rs +++ b/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(