Add logic for closed beta LLM models (#16482)

Release Notes:

- N/A

---------

Co-authored-by: Marshall <marshall@zed.dev>

crates/collab/src/llm/authorization.rs

@@ -12,11 +12,12 @@ pub fn authorize_access_to_language_model(
     model: &str,
 ) -> Result<()> {
     authorize_access_for_country(config, country_code, provider)?;
-    authorize_access_to_model(claims, provider, model)?;
+    authorize_access_to_model(config, claims, provider, model)?;
     Ok(())
 }
 
 fn authorize_access_to_model(
+    config: &Config,
     claims: &LlmTokenClaims,
     provider: LanguageModelProvider,
     model: &str,
@@ -25,13 +26,25 @@ fn authorize_access_to_model(
         return Ok(());
     }
 
-    match (provider, model) {
-        (LanguageModelProvider::Anthropic, "claude-3-5-sonnet") => Ok(()),
-        _ => Err(Error::http(
-            StatusCode::FORBIDDEN,
-            format!("access to model {model:?} is not included in your plan"),
-        ))?,
+    match provider {
+        LanguageModelProvider::Anthropic => {
+            if model == "claude-3-5-sonnet" {
+                return Ok(());
+            }
+
+            if claims.has_llm_closed_beta_feature_flag
+                && Some(model) == config.llm_closed_beta_model_name.as_deref()
+            {
+                return Ok(());
+            }
+        }
+        _ => {}
     }
+
+    Err(Error::http(
+        StatusCode::FORBIDDEN,
+        format!("access to model {model:?} is not included in your plan"),
+    ))
 }
 
 fn authorize_access_for_country(

crates/collab/src/llm/db/queries/usages.rs

@@ -82,12 +82,13 @@ impl LlmDatabase {
             let tokens_per_minute = self.usage_measure_ids[&UsageMeasure::TokensPerMinute];
 
             let mut results = Vec::new();
-            for (provider, model) in self.models.keys().cloned() {
+            for ((provider, model_name), model) in self.models.iter() {
                 let mut usages = usage::Entity::find()
                     .filter(
                         usage::Column::Timestamp
                             .gte(past_minute.naive_utc())
                             .and(usage::Column::IsStaff.eq(false))
+                            .and(usage::Column::ModelId.eq(model.id))
                             .and(
                                 usage::Column::MeasureId
                                     .eq(requests_per_minute)
@@ -125,8 +126,8 @@ impl LlmDatabase {
                 }
 
                 results.push(ApplicationWideUsage {
-                    provider,
-                    model,
+                    provider: *provider,
+                    model: model_name.clone(),
                     requests_this_minute,
                     tokens_this_minute,
                 })

crates/collab/src/llm/token.rs

@@ -20,6 +20,8 @@ pub struct LlmTokenClaims {
     #[serde(default)]
     pub github_user_login: Option<String>,
     pub is_staff: bool,
+    #[serde(default)]
+    pub has_llm_closed_beta_feature_flag: bool,
     pub plan: rpc::proto::Plan,
 }
 
@@ -30,6 +32,7 @@ impl LlmTokenClaims {
         user_id: UserId,
         github_user_login: String,
         is_staff: bool,
+        has_llm_closed_beta_feature_flag: bool,
         plan: rpc::proto::Plan,
         config: &Config,
     ) -> Result<String> {
@@ -46,6 +49,7 @@ impl LlmTokenClaims {
             user_id: user_id.to_proto(),
             github_user_login: Some(github_user_login),
             is_staff,
+            has_llm_closed_beta_feature_flag,
             plan,
         };