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

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

View file

@ -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,