From 92f21ee39d86ff45f87425aa5dd291915ccecce6 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <git@maxdeviant.com>
Date: Thu, 24 Apr 2025 09:04:25 -0400
Subject: [PATCH] collab: Return current plan based on subscription status
 (#29341)

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

Release Notes:

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

diff --git a/crates/collab/src/rpc.rs b/crates/collab/src/rpc.rs
index e2dacc5389..3813ee67fd 100644
--- a/crates/collab/src/rpc.rs
+++ b/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 {