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 | 161 ++++++++++++++++++++++++++------------- 1 file changed, 108 insertions(+), 53 deletions(-) diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs index 753e591914..7a454e11cf 100644 --- a/crates/collab/src/rpc.rs +++ b/crates/collab/src/rpc.rs @@ -2836,62 +2836,117 @@ 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, - }; - - 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, - }; - - 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 {}) - } - }), - }), - } - }), + 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 + }; + + zed_llm_client::UsageLimit::Limited(limit) + } + zed_llm_client::UsageLimit::Unlimited => zed_llm_client::UsageLimit::Unlimited, + } +} + +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 => { + 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<()> { let db = session.db().await;