From 558bbfffaeb6b4c4da8933e58fa9e070e43ba1f0 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Thu, 31 Jul 2025 09:56:53 -0400 Subject: [PATCH] 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 | 1 + crates/client/src/cloud/user_store.rs | 81 ++++++++++++++++++++++++-- crates/client/src/user.rs | 2 + crates/collab/src/tests/test_server.rs | 3 +- crates/title_bar/Cargo.toml | 1 + crates/title_bar/src/title_bar.rs | 16 ++--- crates/workspace/src/workspace.rs | 6 +- crates/zed/src/main.rs | 3 +- 8 files changed, 94 insertions(+), 19 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b6ba5eaba1..34ca4f04df 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -16537,6 +16537,7 @@ dependencies = [ "call", "chrono", "client", + "cloud_llm_client", "collections", "db", "gpui", diff --git a/crates/client/src/cloud/user_store.rs b/crates/client/src/cloud/user_store.rs index ef4e92299a..a9b13ca23c 100644 --- a/crates/client/src/cloud/user_store.rs +++ b/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 gpui::{Context, Task}; +use chrono::{DateTime, Utc}; +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, authenticated_user: Option>, + plan_info: Option>, _maintain_authenticated_user_task: Task<()>, + _rpc_plan_updated_subscription: Subscription, } impl CloudUserStore { - pub fn new(cloud_client: Arc, cx: &mut Context) -> Self { + pub fn new( + cloud_client: Arc, + rpc_user_store: Entity, + cx: &mut Context, + ) -> 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> { self.authenticated_user.clone() } + + pub fn plan(&self) -> Option { + self.plan_info.as_ref().map(|plan| plan.plan) + } + + pub fn subscription_period(&self) -> Option<(DateTime, DateTime)> { + 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, + event: &RpcUserStoreEvent, + cx: &mut Context, + ) { + 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); + } + _ => {} + } + } } diff --git a/crates/client/src/user.rs b/crates/client/src/user.rs index a7dab2a8d3..e025ec0523 100644 --- a/crates/client/src/user.rs +++ b/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(()) diff --git a/crates/collab/src/tests/test_server.rs b/crates/collab/src/tests/test_server.rs index ab6bf1b912..00d1caa7c5 100644 --- a/crates/collab/src/tests/test_server.rs +++ b/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)); diff --git a/crates/title_bar/Cargo.toml b/crates/title_bar/Cargo.toml index 8e95c6f79f..cf178e2850 100644 --- a/crates/title_bar/Cargo.toml +++ b/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 diff --git a/crates/title_bar/src/title_bar.rs b/crates/title_bar/src/title_bar.rs index 6e03b52ef8..552ef915cb 100644 --- a/crates/title_bar/src/title_bar.rs +++ b/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) -> 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 plan = self.user_store.read(cx).current_plan().filter(|_| { + let has_subscription_period = cloud_user_store.subscription_period().is_some(); + 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) => { - ("Free", Color::Default, free_chip_bg) - } - Some(proto::Plan::ZedProTrial) => { - ("Pro Trial", Color::Accent, pro_chip_bg) - } - Some(proto::Plan::ZedPro) => ("Pro", Color::Accent, pro_chip_bg), + None | Some(Plan::ZedFree) => ("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), }; menu.custom_entry( diff --git a/crates/workspace/src/workspace.rs b/crates/workspace/src/workspace.rs index 700554b748..aad585e419 100644 --- a/crates/workspace/src/workspace.rs +++ b/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)); diff --git a/crates/zed/src/main.rs b/crates/zed/src/main.rs index 17ee7d2512..338840607b 100644 --- a/crates/zed/src/main.rs +++ b/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(