From f4106ad404097f00975c9f9a797b2f726b7e1269 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Thu, 10 Jul 2025 11:52:19 -0400 Subject: [PATCH] collab: Send down new usage limits even when the user does not have any usage (#34217) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR fixes an issue where the plan usage limits in Zed would not get updated immediately after the plan has changed. Previously we were only sending down the usage—which contains the limits—if there was a usage record in the database. This would be absent if the user had just changed their plan. We now always send down the usage in order to update the limits on the client side. Release Notes: - N/A Co-authored-by: Bennet Bo Fenner --- crates/collab/src/rpc.rs | 151 ++++++++++++++++++++++++++------------- 1 file changed, 103 insertions(+), 48 deletions(-) diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs index 753e591914f45ea962367130d1ecce9a4fd2620f..7a454e11cfced2fa7f9f1dc8c0263934830c7cad 100644 --- a/crates/collab/src/rpc.rs +++ b/crates/collab/src/rpc.rs @@ -2836,60 +2836,115 @@ async fn make_update_user_plan_message( account_too_young: Some(account_too_young), has_overdue_invoices: billing_customer .map(|billing_customer| billing_customer.has_overdue_invoices), - usage: usage.map(|usage| { - let plan = match plan { - proto::Plan::Free => zed_llm_client::Plan::ZedFree, - proto::Plan::ZedPro => zed_llm_client::Plan::ZedPro, - proto::Plan::ZedProTrial => zed_llm_client::Plan::ZedProTrial, + usage: Some( + usage + .map(|usage| subscription_usage_to_proto(plan, usage, &feature_flags)) + .unwrap_or_else(|| make_default_subscription_usage(plan, &feature_flags)), + ), + }) +} + +fn model_requests_limit( + plan: zed_llm_client::Plan, + feature_flags: &Vec, +) -> zed_llm_client::UsageLimit { + match plan.model_requests_limit() { + zed_llm_client::UsageLimit::Limited(limit) => { + let limit = if plan == zed_llm_client::Plan::ZedProTrial + && feature_flags + .iter() + .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG) + { + 1_000 + } else { + limit }; - let model_requests_limit = match plan.model_requests_limit() { - zed_llm_client::UsageLimit::Limited(limit) => { - let limit = if plan == zed_llm_client::Plan::ZedProTrial - && feature_flags - .iter() - .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG) - { - 1_000 - } else { - limit - }; + zed_llm_client::UsageLimit::Limited(limit) + } + zed_llm_client::UsageLimit::Unlimited => zed_llm_client::UsageLimit::Unlimited, + } +} - zed_llm_client::UsageLimit::Limited(limit) +fn subscription_usage_to_proto( + plan: proto::Plan, + usage: crate::llm::db::subscription_usage::Model, + feature_flags: &Vec, +) -> proto::SubscriptionUsage { + let plan = match plan { + proto::Plan::Free => zed_llm_client::Plan::ZedFree, + proto::Plan::ZedPro => zed_llm_client::Plan::ZedPro, + proto::Plan::ZedProTrial => zed_llm_client::Plan::ZedProTrial, + }; + + proto::SubscriptionUsage { + model_requests_usage_amount: usage.model_requests as u32, + model_requests_usage_limit: Some(proto::UsageLimit { + variant: Some(match model_requests_limit(plan, feature_flags) { + zed_llm_client::UsageLimit::Limited(limit) => { + proto::usage_limit::Variant::Limited(proto::usage_limit::Limited { + limit: limit as u32, + }) } - zed_llm_client::UsageLimit::Unlimited => zed_llm_client::UsageLimit::Unlimited, - }; + zed_llm_client::UsageLimit::Unlimited => { + proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {}) + } + }), + }), + edit_predictions_usage_amount: usage.edit_predictions as u32, + edit_predictions_usage_limit: Some(proto::UsageLimit { + variant: Some(match plan.edit_predictions_limit() { + zed_llm_client::UsageLimit::Limited(limit) => { + proto::usage_limit::Variant::Limited(proto::usage_limit::Limited { + limit: limit as u32, + }) + } + zed_llm_client::UsageLimit::Unlimited => { + proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {}) + } + }), + }), + } +} - proto::SubscriptionUsage { - model_requests_usage_amount: usage.model_requests as u32, - model_requests_usage_limit: Some(proto::UsageLimit { - variant: Some(match model_requests_limit { - zed_llm_client::UsageLimit::Limited(limit) => { - proto::usage_limit::Variant::Limited(proto::usage_limit::Limited { - limit: limit as u32, - }) - } - zed_llm_client::UsageLimit::Unlimited => { - proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {}) - } - }), - }), - edit_predictions_usage_amount: usage.edit_predictions as u32, - edit_predictions_usage_limit: Some(proto::UsageLimit { - variant: Some(match plan.edit_predictions_limit() { - zed_llm_client::UsageLimit::Limited(limit) => { - proto::usage_limit::Variant::Limited(proto::usage_limit::Limited { - limit: limit as u32, - }) - } - zed_llm_client::UsageLimit::Unlimited => { - proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {}) - } - }), - }), - } +fn make_default_subscription_usage( + plan: proto::Plan, + feature_flags: &Vec, +) -> proto::SubscriptionUsage { + let plan = match plan { + proto::Plan::Free => zed_llm_client::Plan::ZedFree, + proto::Plan::ZedPro => zed_llm_client::Plan::ZedPro, + proto::Plan::ZedProTrial => zed_llm_client::Plan::ZedProTrial, + }; + + proto::SubscriptionUsage { + model_requests_usage_amount: 0, + model_requests_usage_limit: Some(proto::UsageLimit { + variant: Some(match model_requests_limit(plan, feature_flags) { + zed_llm_client::UsageLimit::Limited(limit) => { + proto::usage_limit::Variant::Limited(proto::usage_limit::Limited { + limit: limit as u32, + }) + } + zed_llm_client::UsageLimit::Unlimited => { + proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {}) + } + }), }), - }) + edit_predictions_usage_amount: 0, + edit_predictions_usage_limit: Some(proto::UsageLimit { + variant: Some(match plan.edit_predictions_limit() { + zed_llm_client::UsageLimit::Limited(limit) => { + proto::usage_limit::Variant::Limited(proto::usage_limit::Limited { + limit: limit as u32, + }) + } + zed_llm_client::UsageLimit::Unlimited => { + proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {}) + } + }), + }), + } } async fn update_user_plan(session: &Session) -> Result<()> {