From 3fef3cc392e634914d9f978f8d36e6c19375ea34 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <git@maxdeviant.com>
Date: Wed, 16 Apr 2025 16:58:00 -0400
Subject: [PATCH] Use more types/constants from `zed_llm_client` (#28909)

This PR makes it so we use more types and constants from the
`zed_llm_client` crate to avoid duplicating information.

Also updates the current usage endpoint to use limits derived from the
`Plan`.

Release Notes:

- N/A
---
 Cargo.lock                                   |  4 ++--
 Cargo.toml                                   |  2 +-
 crates/collab/src/api/billing.rs             | 16 ++++++++++++++--
 crates/collab/src/llm/token.rs               |  9 +++++++--
 crates/language_models/src/provider/cloud.rs |  6 ++++--
 5 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 52398176d7..46041f802d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -18387,9 +18387,9 @@ dependencies = [
 
 [[package]]
 name = "zed_llm_client"
-version = "0.5.0"
+version = "0.5.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "57a5e1b5b3ace3fb55292a4c14036723bb8a01fac4aeaa3c2b63b51228412f94"
+checksum = "9ee4d410dbc030c3e6e3af78fc76296f6bebe20dcb6d7d3fa24bca306fc8c1ce"
 dependencies = [
  "serde",
  "serde_json",
diff --git a/Cargo.toml b/Cargo.toml
index 099ea3ddad..045a627024 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -605,7 +605,7 @@ wasmtime-wasi = "29"
 which = "6.0.0"
 wit-component = "0.221"
 workspace-hack = "0.1.0"
-zed_llm_client = "0.5.0"
+zed_llm_client = "0.5.1"
 zstd = "0.11"
 metal = "0.29"
 
diff --git a/crates/collab/src/api/billing.rs b/crates/collab/src/api/billing.rs
index bc2b508e84..ad8e1f662d 100644
--- a/crates/collab/src/api/billing.rs
+++ b/crates/collab/src/api/billing.rs
@@ -1020,8 +1020,20 @@ async fn get_current_usage(
         return Ok(Json(empty_usage));
     };
 
-    let model_requests_limit = Some(500);
-    let edit_prediction_limit = Some(2000);
+    let plan = match usage.plan {
+        SubscriptionKind::ZedPro => zed_llm_client::Plan::ZedPro,
+        SubscriptionKind::ZedProTrial => zed_llm_client::Plan::ZedProTrial,
+        SubscriptionKind::ZedFree => zed_llm_client::Plan::Free,
+    };
+
+    let model_requests_limit = match plan.model_requests_limit() {
+        zed_llm_client::UsageLimit::Limited(limit) => Some(limit),
+        zed_llm_client::UsageLimit::Unlimited => None,
+    };
+    let edit_prediction_limit = match plan.edit_predictions_limit() {
+        zed_llm_client::UsageLimit::Limited(limit) => Some(limit),
+        zed_llm_client::UsageLimit::Unlimited => None,
+    };
 
     Ok(Json(GetCurrentUsageResponse {
         model_requests: UsageCounts {
diff --git a/crates/collab/src/llm/token.rs b/crates/collab/src/llm/token.rs
index 067be93076..1f64551b28 100644
--- a/crates/collab/src/llm/token.rs
+++ b/crates/collab/src/llm/token.rs
@@ -10,6 +10,7 @@ use std::time::Duration;
 use thiserror::Error;
 use util::maybe;
 use uuid::Uuid;
+use zed_llm_client::Plan;
 
 #[derive(Clone, Debug, Default, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
@@ -28,7 +29,7 @@ pub struct LlmTokenClaims {
     pub has_llm_subscription: bool,
     pub max_monthly_spend_in_cents: u32,
     pub custom_llm_monthly_allowance_in_cents: Option<u32>,
-    pub plan: rpc::proto::Plan,
+    pub plan: Plan,
     #[serde(default)]
     pub subscription_period: Option<(NaiveDateTime, NaiveDateTime)>,
 }
@@ -77,7 +78,11 @@ impl LlmTokenClaims {
             custom_llm_monthly_allowance_in_cents: user
                 .custom_llm_monthly_allowance_in_cents
                 .map(|allowance| allowance as u32),
-            plan,
+            plan: match plan {
+                rpc::proto::Plan::Free => Plan::Free,
+                rpc::proto::Plan::ZedPro => Plan::ZedPro,
+                rpc::proto::Plan::ZedProTrial => Plan::ZedProTrial,
+            },
             subscription_period: maybe!({
                 let subscription = subscription?;
                 let period_start_at = subscription.current_period_start_at()?;
diff --git a/crates/language_models/src/provider/cloud.rs b/crates/language_models/src/provider/cloud.rs
index ee4c8540d8..8286d0a1f2 100644
--- a/crates/language_models/src/provider/cloud.rs
+++ b/crates/language_models/src/provider/cloud.rs
@@ -36,7 +36,8 @@ use thiserror::Error;
 use ui::{TintColor, prelude::*};
 use zed_llm_client::{
     CURRENT_PLAN_HEADER_NAME, CompletionBody, EXPIRED_LLM_TOKEN_HEADER_NAME,
-    MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME, SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME,
+    MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME, MODEL_REQUESTS_RESOURCE_HEADER_VALUE,
+    SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME,
 };
 
 use crate::AllLanguageModelSettings;
@@ -560,7 +561,7 @@ impl CloudLanguageModel {
                     .get(SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME)
                     .is_some()
             {
-                if let Some("model_requests") = response
+                if let Some(MODEL_REQUESTS_RESOURCE_HEADER_VALUE) = response
                     .headers()
                     .get(SUBSCRIPTION_LIMIT_RESOURCE_HEADER_NAME)
                     .and_then(|resource| resource.to_str().ok())
@@ -574,6 +575,7 @@ impl CloudLanguageModel {
                         let plan = match plan {
                             zed_llm_client::Plan::Free => Plan::Free,
                             zed_llm_client::Plan::ZedPro => Plan::ZedPro,
+                            zed_llm_client::Plan::ZedProTrial => Plan::ZedProTrial,
                         };
                         return Err(anyhow!(ModelRequestLimitReachedError { plan }));
                     }