collab: Return current plan based on subscription status (#29341)

Marshall Bowers created 8 months ago

This PR makes collab return the current plan based on subscription
status instead of based on the staff bit.

Release Notes:

- N/A

Change summary

crates/collab/src/rpc.rs | 25 +++++++++++++++++--------
1 file changed, 17 insertions(+), 8 deletions(-)

Detailed changes

crates/collab/src/rpc.rs

@@ -1,6 +1,7 @@
 mod connection_pool;
 
 use crate::api::{CloudflareIpCountryHeader, SystemIdHeader};
+use crate::db::billing_subscription::SubscriptionKind;
 use crate::llm::LlmTokenClaims;
 use crate::{
     AppState, Error, Result, auth,
@@ -178,15 +179,23 @@ impl Session {
         Ok(db.has_active_billing_subscription(user_id).await?)
     }
 
-    pub async fn current_plan(
-        &self,
-        _db: &MutexGuard<'_, DbHandle>,
-    ) -> anyhow::Result<proto::Plan> {
-        if self.is_staff() {
-            Ok(proto::Plan::ZedPro)
+    pub async fn current_plan(&self, db: &MutexGuard<'_, DbHandle>) -> anyhow::Result<proto::Plan> {
+        let user_id = self.user_id();
+
+        let subscription = db.get_active_billing_subscription(user_id).await?;
+        let subscription_kind = subscription.and_then(|subscription| subscription.kind);
+
+        let plan = if let Some(subscription_kind) = subscription_kind {
+            match subscription_kind {
+                SubscriptionKind::ZedPro => proto::Plan::ZedPro,
+                SubscriptionKind::ZedProTrial => proto::Plan::ZedProTrial,
+                SubscriptionKind::ZedFree => proto::Plan::Free,
+            }
         } else {
-            Ok(proto::Plan::Free)
-        }
+            proto::Plan::Free
+        };
+
+        Ok(plan)
     }
 
     fn user_id(&self) -> UserId {