collab: Send down new usage limits even when the user does not have any usage (#34217)

Marshall Bowers and Bennet Bo Fenner created 4 months ago

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 <[e-mail]>

Change summary

crates/collab/src/rpc.rs | 151 ++++++++++++++++++++++++++++-------------
1 file changed, 103 insertions(+), 48 deletions(-)

Detailed changes

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<String>,
+) -> 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<String>,
+) -> 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<String>,
+) -> 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<()> {