title_bar: Show the plan from the CloudUserStore (#35401)

This PR updates the user menu in the title bar to show the plan from the
`CloudUserStore` instead of the `UserStore`.

We're still leveraging the RPC connection to listen for `UpdateUserPlan`
messages so that we can get live-updates from the server, but we are
merely using this as a signal to re-fetch the information from Cloud.

Release Notes:

- N/A

Cargo.lock

@@ -16537,6 +16537,7 @@ dependencies = [
  "call",
  "chrono",
  "client",
+ "cloud_llm_client",
  "collections",
  "db",
  "gpui",

crates/client/src/cloud/user_store.rs

@@ -2,19 +2,36 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use anyhow::Context as _;
-use cloud_api_client::{AuthenticatedUser, CloudApiClient};
+use chrono::{DateTime, Utc};
-use gpui::{Context, Task};
+use cloud_api_client::{AuthenticatedUser, CloudApiClient, GetAuthenticatedUserResponse, PlanInfo};
+use cloud_llm_client::Plan;
+use gpui::{Context, Entity, Subscription, Task};
 use util::{ResultExt as _, maybe};
 
+use crate::UserStore;
+use crate::user::Event as RpcUserStoreEvent;
+
 pub struct CloudUserStore {
+    cloud_client: Arc<CloudApiClient>,
     authenticated_user: Option<Arc<AuthenticatedUser>>,
+    plan_info: Option<Arc<PlanInfo>>,
     _maintain_authenticated_user_task: Task<()>,
+    _rpc_plan_updated_subscription: Subscription,
 }
 
 impl CloudUserStore {
-    pub fn new(cloud_client: Arc<CloudApiClient>, cx: &mut Context<Self>) -> Self {
+    pub fn new(
+        cloud_client: Arc<CloudApiClient>,
+        rpc_user_store: Entity<UserStore>,
+        cx: &mut Context<Self>,
+    ) -> Self {
+        let rpc_plan_updated_subscription =
+            cx.subscribe(&rpc_user_store, Self::handle_rpc_user_store_event);
+
         Self {
+            cloud_client: cloud_client.clone(),
             authenticated_user: None,
+            plan_info: None,
             _maintain_authenticated_user_task: cx.spawn(async move |this, cx| {
                 maybe!(async move {
                     loop {
@@ -36,14 +53,15 @@ impl CloudUserStore {
                                     .context("failed to fetch authenticated user");
                                 if let Some(response) = authenticated_user_result.log_err() {
                                     this.update(cx, |this, _cx| {
-                                        this.authenticated_user = Some(Arc::new(response.user));
+                                        this.update_authenticated_user(response);
                                     })
                                     .ok();
                                 }
                             }
                         } else {
                             this.update(cx, |this, _cx| {
-                                this.authenticated_user = None;
+                                this.authenticated_user.take();
+                                this.plan_info.take();
                             })
                             .ok();
                         }
@@ -56,6 +74,7 @@ impl CloudUserStore {
                 .await
                 .log_err();
             }),
+            _rpc_plan_updated_subscription: rpc_plan_updated_subscription,
         }
     }
 
@@ -66,4 +85,56 @@ impl CloudUserStore {
     pub fn authenticated_user(&self) -> Option<Arc<AuthenticatedUser>> {
         self.authenticated_user.clone()
     }
+
+    pub fn plan(&self) -> Option<Plan> {
+        self.plan_info.as_ref().map(|plan| plan.plan)
+    }
+
+    pub fn subscription_period(&self) -> Option<(DateTime<Utc>, DateTime<Utc>)> {
+        self.plan_info
+            .as_ref()
+            .and_then(|plan| plan.subscription_period)
+            .map(|subscription_period| {
+                (
+                    subscription_period.started_at.0,
+                    subscription_period.ended_at.0,
+                )
+            })
+    }
+
+    fn update_authenticated_user(&mut self, response: GetAuthenticatedUserResponse) {
+        self.authenticated_user = Some(Arc::new(response.user));
+        self.plan_info = Some(Arc::new(response.plan));
+    }
+
+    fn handle_rpc_user_store_event(
+        &mut self,
+        _: Entity<UserStore>,
+        event: &RpcUserStoreEvent,
+        cx: &mut Context<Self>,
+    ) {
+        match event {
+            RpcUserStoreEvent::PlanUpdated => {
+                cx.spawn(async move |this, cx| {
+                    let cloud_client =
+                        cx.update(|cx| this.read_with(cx, |this, _cx| this.cloud_client.clone()))??;
+
+                    let response = cloud_client
+                        .get_authenticated_user()
+                        .await
+                        .context("failed to fetch authenticated user")?;
+
+                    cx.update(|cx| {
+                        this.update(cx, |this, _cx| {
+                            this.update_authenticated_user(response);
+                        })
+                    })??;
+
+                    anyhow::Ok(())
+                })
+                .detach_and_log_err(cx);
+            }
+            _ => {}
+        }
+    }
 }

crates/client/src/user.rs

@@ -145,6 +145,7 @@ pub enum Event {
     ShowContacts,
     ParticipantIndicesChanged,
     PrivateUserInfoUpdated,
+    PlanUpdated,
 }
 
 #[derive(Clone, Copy)]
@@ -388,6 +389,7 @@ impl UserStore {
                     .map(EditPredictionUsage);
             }
 
+            cx.emit(Event::PlanUpdated);
             cx.notify();
         })?;
         Ok(())

crates/collab/src/tests/test_server.rs

@@ -282,7 +282,8 @@ impl TestServer {
             .register_hosting_provider(Arc::new(git_hosting_providers::Github::public_instance()));
 
         let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
-        let cloud_user_store = cx.new(|cx| CloudUserStore::new(client.cloud_client(), cx));
+        let cloud_user_store =
+            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store.clone(), cx));
         let workspace_store = cx.new(|cx| WorkspaceStore::new(client.clone(), cx));
         let language_registry = Arc::new(LanguageRegistry::test(cx.executor()));
         let session = cx.new(|cx| AppSession::new(Session::test(), cx));

crates/title_bar/Cargo.toml

@@ -32,6 +32,7 @@ auto_update.workspace = true
 call.workspace = true
 chrono.workspace = true
 client.workspace = true
+cloud_llm_client.workspace = true
 db.workspace = true
 gpui = { workspace = true, features = ["screen-capture"] }
 notifications.workspace = true

crates/title_bar/src/title_bar.rs

@@ -21,6 +21,7 @@ use crate::application_menu::{
 use auto_update::AutoUpdateStatus;
 use call::ActiveCall;
 use client::{Client, CloudUserStore, UserStore, zed_urls};
+use cloud_llm_client::Plan;
 use gpui::{
     Action, AnyElement, App, Context, Corner, Element, Entity, Focusable, InteractiveElement,
     IntoElement, MouseButton, ParentElement, Render, StatefulInteractiveElement, Styled,
@@ -28,7 +29,6 @@ use gpui::{
 };
 use onboarding_banner::OnboardingBanner;
 use project::Project;
-use rpc::proto;
 use settings::Settings as _;
 use settings_ui::keybindings;
 use std::sync::Arc;
@@ -634,8 +634,8 @@ impl TitleBar {
     pub fn render_user_menu_button(&mut self, cx: &mut Context<Self>) -> impl Element {
         let cloud_user_store = self.cloud_user_store.read(cx);
         if let Some(user) = cloud_user_store.authenticated_user() {
-            let has_subscription_period = self.user_store.read(cx).subscription_period().is_some();
+            let has_subscription_period = cloud_user_store.subscription_period().is_some();
-            let plan = self.user_store.read(cx).current_plan().filter(|_| {
+            let plan = cloud_user_store.plan().filter(|_| {
                 // Since the user might be on the legacy free plan we filter based on whether we have a subscription period.
                 has_subscription_period
             });
@@ -662,13 +662,9 @@ impl TitleBar {
                         let user_login = user.github_login.clone();
 
                         let (plan_name, label_color, bg_color) = match plan {
-                            None | Some(proto::Plan::Free) => {
+                            None | Some(Plan::ZedFree) => ("Free", Color::Default, free_chip_bg),
-                                ("Free", Color::Default, free_chip_bg)
+                            Some(Plan::ZedProTrial) => ("Pro Trial", Color::Accent, pro_chip_bg),
-                            }
+                            Some(Plan::ZedPro) => ("Pro", Color::Accent, pro_chip_bg),
-                            Some(proto::Plan::ZedProTrial) => {
-                                ("Pro Trial", Color::Accent, pro_chip_bg)
-                            }
-                            Some(proto::Plan::ZedPro) => ("Pro", Color::Accent, pro_chip_bg),
                         };
 
                         menu.custom_entry(

crates/workspace/src/workspace.rs

@@ -913,7 +913,8 @@ impl AppState {
         let client = Client::new(clock, http_client.clone(), cx);
         let session = cx.new(|cx| AppSession::new(Session::test(), cx));
         let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
-        let cloud_user_store = cx.new(|cx| CloudUserStore::new(client.cloud_client(), cx));
+        let cloud_user_store =
+            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store.clone(), cx));
         let workspace_store = cx.new(|cx| WorkspaceStore::new(client.clone(), cx));
 
         theme::init(theme::LoadThemes::JustBase, cx);
@@ -5738,7 +5739,8 @@ impl Workspace {
 
         let client = project.read(cx).client();
         let user_store = project.read(cx).user_store();
-        let cloud_user_store = cx.new(|cx| CloudUserStore::new(client.cloud_client(), cx));
+        let cloud_user_store =
+            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store.clone(), cx));
 
         let workspace_store = cx.new(|cx| WorkspaceStore::new(client.clone(), cx));
         let session = cx.new(|cx| AppSession::new(Session::test(), cx));

crates/zed/src/main.rs

@@ -457,7 +457,8 @@ pub fn main() {
         language::init(cx);
         languages::init(languages.clone(), node_runtime.clone(), cx);
         let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
-        let cloud_user_store = cx.new(|cx| CloudUserStore::new(client.cloud_client(), cx));
+        let cloud_user_store =
+            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store.clone(), cx));
         let workspace_store = cx.new(|cx| WorkspaceStore::new(client.clone(), cx));
 
         language_extension::init(