Check for predict-edits feature flag, remove is_staff check (#23165)

Release Notes:

- N/A

---------

Co-authored-by: Thorsten Ball <mrnugget@gmail.com>

crates/collab/src/llm.rs

@@ -440,8 +440,11 @@ async fn predict_edits(
     _country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,
     Json(params): Json<PredictEditsParams>,
 ) -> Result<impl IntoResponse> {
-    if !claims.is_staff {
-        return Err(anyhow!("not found"))?;
+    if !claims.is_staff && !claims.has_predict_edits_feature_flag {
+        return Err(Error::http(
+            StatusCode::FORBIDDEN,
+            "no access to Zed's edit prediction feature".to_string(),
+        ));
     }
 
     let api_url = state

crates/collab/src/llm/token.rs

@@ -22,6 +22,8 @@ pub struct LlmTokenClaims {
     pub github_user_login: String,
     pub is_staff: bool,
     pub has_llm_closed_beta_feature_flag: bool,
+    #[serde(default)]
+    pub has_predict_edits_feature_flag: bool,
     pub has_llm_subscription: bool,
     pub max_monthly_spend_in_cents: u32,
     pub custom_llm_monthly_allowance_in_cents: Option<u32>,
@@ -37,6 +39,7 @@ impl LlmTokenClaims {
         is_staff: bool,
         billing_preferences: Option<billing_preference::Model>,
         has_llm_closed_beta_feature_flag: bool,
+        has_predict_edits_feature_flag: bool,
         has_llm_subscription: bool,
         plan: rpc::proto::Plan,
         system_id: Option<String>,
@@ -58,6 +61,7 @@ impl LlmTokenClaims {
             github_user_login: user.github_login.clone(),
             is_staff,
             has_llm_closed_beta_feature_flag,
+            has_predict_edits_feature_flag,
             has_llm_subscription,
             max_monthly_spend_in_cents: billing_preferences
                 .map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {

crates/collab/src/rpc.rs

@@ -4025,6 +4025,7 @@ 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_llm_closed_beta_feature_flag = flags.iter().any(|flag| flag == "llm-closed-beta");
+    let has_predict_edits_feature_flag = flags.iter().any(|flag| flag == "predict-edits");
 
     if !session.is_staff() && !has_language_models_feature_flag {
         Err(anyhow!("permission denied"))?
@@ -4061,6 +4062,7 @@ async fn get_llm_api_token(
         session.is_staff(),
         billing_preferences,
         has_llm_closed_beta_feature_flag,
+        has_predict_edits_feature_flag,
         has_llm_subscription,
         session.current_plan(&db).await?,
         session.system_id.clone(),