collab: Clean up LLM token creation (#26955)

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
This commit is contained in:
Marshall Bowers 2025-03-17 18:25:43 -04:00 committed by GitHub
parent 2b2b9c1624
commit 1397e01735
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 26 additions and 14 deletions

View file

@ -43,6 +43,20 @@ pub enum Relation {
Contributor, 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 { impl Related<super::access_token::Entity> for Entity {
fn to() -> RelationDef { fn to() -> RelationDef {
Relation::AccessToken.def() Relation::AccessToken.def()

View file

@ -37,8 +37,7 @@ impl LlmTokenClaims {
user: &user::Model, user: &user::Model,
is_staff: bool, is_staff: bool,
billing_preferences: Option<billing_preference::Model>, billing_preferences: Option<billing_preference::Model>,
has_llm_closed_beta_feature_flag: bool, feature_flags: &Vec<String>,
has_predict_edits_feature_flag: bool,
has_llm_subscription: bool, has_llm_subscription: bool,
plan: rpc::proto::Plan, plan: rpc::proto::Plan,
system_id: Option<String>, system_id: Option<String>,
@ -59,8 +58,12 @@ impl LlmTokenClaims {
metrics_id: user.metrics_id, metrics_id: user.metrics_id,
github_user_login: user.github_login.clone(), github_user_login: user.github_login.clone(),
is_staff, is_staff,
has_llm_closed_beta_feature_flag, has_llm_closed_beta_feature_flag: feature_flags
has_predict_edits_feature_flag, .iter()
.any(|flag| flag == "llm-closed-beta"),
has_predict_edits_feature_flag: feature_flags
.iter()
.any(|flag| flag == "predict-edits"),
has_llm_subscription, has_llm_subscription,
max_monthly_spend_in_cents: billing_preferences max_monthly_spend_in_cents: billing_preferences
.map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| { .map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {

View file

@ -4047,8 +4047,8 @@ async fn get_llm_api_token(
let flags = db.get_user_flags(session.user_id()).await?; 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_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_bypass_account_age_check_feature_flag =
let has_predict_edits_feature_flag = flags.iter().any(|flag| flag == "predict-edits"); flags.iter().any(|flag| flag == "bypass-account-age-check");
if !session.is_staff() && !has_language_models_feature_flag { if !session.is_staff() && !has_language_models_feature_flag {
Err(anyhow!("permission denied"))? 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 has_llm_subscription = session.has_llm_subscription(&db).await?;
let bypass_account_age_check = 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 { if !bypass_account_age_check {
let mut account_created_at = user.created_at; if Utc::now().naive_utc() - user.account_created_at() < MIN_ACCOUNT_AGE_FOR_LLM_USE {
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 {
Err(anyhow!("account too young"))? Err(anyhow!("account too young"))?
} }
} }
@ -4084,8 +4080,7 @@ async fn get_llm_api_token(
&user, &user,
session.is_staff(), session.is_staff(),
billing_preferences, billing_preferences,
has_llm_closed_beta_feature_flag, &flags,
has_predict_edits_feature_flag,
has_llm_subscription, has_llm_subscription,
session.current_plan(&db).await?, session.current_plan(&db).await?,
session.system_id.clone(), session.system_id.clone(),