title_bar: Show organization plans (#49769)

Marshall Bowers created 1 month ago

This PR updates the organization selector to also show the plan that is
available through each organization.

Also makes the `plan` method on the `UserStore` return the plan for the
current organization, if it exists.

We aren't yet populating the `plans_by_organization` collection.

Release Notes:

- N/A

Change summary

crates/client/src/user.rs                     | 16 ++++++++++
crates/cloud_api_types/src/cloud_api_types.rs |  4 +-
crates/title_bar/src/title_bar.rs             | 29 +++++++++++++++++++-
3 files changed, 44 insertions(+), 5 deletions(-)

Detailed changes

crates/client/src/user.rs

@@ -2,7 +2,9 @@ use super::{Client, Status, TypedEnvelope, proto};
 use anyhow::{Context as _, Result};
 use chrono::{DateTime, Utc};
 use cloud_api_client::websocket_protocol::MessageToClient;
-use cloud_api_client::{GetAuthenticatedUserResponse, Organization, Plan, PlanInfo};
+use cloud_api_client::{
+    GetAuthenticatedUserResponse, Organization, OrganizationId, Plan, PlanInfo,
+};
 use cloud_llm_client::{
     EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME, EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME, UsageLimit,
 };
@@ -111,6 +113,7 @@ pub struct UserStore {
     current_user: watch::Receiver<Option<Arc<User>>>,
     current_organization: Option<Arc<Organization>>,
     organizations: Vec<Arc<Organization>>,
+    plans_by_organization: HashMap<OrganizationId, Plan>,
     contacts: Vec<Arc<Contact>>,
     incoming_contact_requests: Vec<Arc<User>>,
     outgoing_contact_requests: Vec<Arc<User>>,
@@ -185,6 +188,7 @@ impl UserStore {
             current_user: current_user_rx,
             current_organization: None,
             organizations: Vec::new(),
+            plans_by_organization: HashMap::default(),
             plan_info: None,
             edit_prediction_usage: None,
             contacts: Default::default(),
@@ -698,6 +702,10 @@ impl UserStore {
         &self.organizations
     }
 
+    pub fn plan_for_organization(&self, organization_id: &OrganizationId) -> Option<Plan> {
+        self.plans_by_organization.get(organization_id).copied()
+    }
+
     pub fn plan(&self) -> Option<Plan> {
         #[cfg(debug_assertions)]
         if let Ok(plan) = std::env::var("ZED_SIMULATE_PLAN").as_ref() {
@@ -713,6 +721,12 @@ impl UserStore {
             };
         }
 
+        if let Some(organization) = &self.current_organization
+            && let Some(plan) = self.plan_for_organization(&organization.id)
+        {
+            return Some(plan);
+        }
+
         self.plan_info.as_ref().map(|info| info.plan())
     }
 

crates/cloud_api_types/src/cloud_api_types.rs

@@ -35,8 +35,8 @@ pub struct AuthenticatedUser {
     pub accepted_tos_at: Option<Timestamp>,
 }
 
-#[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize)]
-pub struct OrganizationId(Arc<str>);
+#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
+pub struct OrganizationId(pub Arc<str>);
 
 #[derive(Debug, PartialEq, Serialize, Deserialize)]
 pub struct Organization {

crates/title_bar/src/title_bar.rs

@@ -969,8 +969,18 @@ impl TitleBar {
                     ContextMenu::build(window, cx, |mut menu, _window, cx| {
                         menu = menu.header("Organizations").separator();
 
+                        let current_organization = user_store.read(cx).current_organization();
+
                         for organization in user_store.read(cx).organizations() {
                             let organization = organization.clone();
+                            let plan = user_store.read(cx).plan_for_organization(&organization.id);
+
+                            let is_current =
+                                current_organization
+                                    .as_ref()
+                                    .is_some_and(|current_organization| {
+                                        current_organization.id == organization.id
+                                    });
 
                             menu = menu.custom_entry(
                                 {
@@ -978,8 +988,23 @@ impl TitleBar {
                                     move |_window, _cx| {
                                         h_flex()
                                             .w_full()
-                                            .justify_between()
-                                            .child(Label::new(&organization.name))
+                                            .gap_1()
+                                            .child(
+                                                div()
+                                                    .flex_none()
+                                                    .when(!is_current, |parent| parent.invisible())
+                                                    .child(Icon::new(IconName::Check)),
+                                            )
+                                            .child(
+                                                h_flex()
+                                                    .w_full()
+                                                    .gap_3()
+                                                    .justify_between()
+                                                    .child(Label::new(&organization.name))
+                                                    .child(PlanChip::new(
+                                                        plan.unwrap_or(Plan::ZedFree),
+                                                    )),
+                                            )
                                             .into_any_element()
                                     }
                                 },