diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs index c365882efb..8df92145ce 100644 --- a/crates/collab/src/llm.rs +++ b/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, diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs index ca0d6f0420..d73cdfd3f1 100644 --- a/crates/collab/src/llm/token.rs +++ b/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, 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 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"), diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs index 45a0d26f04..8e63a80c55 100644 --- a/crates/collab/src/rpc.rs +++ b/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 {