Start separating authentication from connection to collab (#35471)

Antonio Scandurra , Marshall Bowers , and Richard Feldman created

This pull request should be idempotent, but lays the groundwork for
avoiding to connect to collab in order to interact with AI features
provided by Zed.

Release Notes:

- N/A

---------

Co-authored-by: Marshall Bowers <git@maxdeviant.com>
Co-authored-by: Richard Feldman <oss@rtfeldman.com>

Change summary

Cargo.lock                                                            |  24 
crates/agent/Cargo.toml                                               |   1 
crates/agent/src/thread.rs                                            |  38 
crates/agent/src/thread_store.rs                                      |  18 
crates/agent_ui/src/active_thread.rs                                  |   7 
crates/agent_ui/src/agent_configuration.rs                            |   6 
crates/agent_ui/src/agent_diff.rs                                     |  13 
crates/agent_ui/src/agent_panel.rs                                    |  23 
crates/agent_ui/src/message_editor.rs                                 |  16 
crates/ai_onboarding/Cargo.toml                                       |   1 
crates/ai_onboarding/src/agent_panel_onboarding_content.rs            |   9 
crates/ai_onboarding/src/ai_onboarding.rs                             |  37 
crates/ai_onboarding/src/ai_upsell_card.rs                            |   6 
crates/assistant_tools/src/edit_agent/evals.rs                        |   6 
crates/channel/src/channel_store_tests.rs                             |  18 
crates/client/Cargo.toml                                              |   1 
crates/client/src/client.rs                                           | 187 
crates/client/src/cloud.rs                                            |   3 
crates/client/src/cloud/user_store.rs                                 | 211 
crates/client/src/test.rs                                             |  94 
crates/client/src/user.rs                                             | 230 
crates/collab/src/tests/integration_tests.rs                          |  10 
crates/collab/src/tests/notification_tests.rs                         |   4 
crates/collab/src/tests/test_server.rs                                |  59 
crates/collab_ui/src/collab_panel.rs                                  |   2 
crates/collab_ui/src/notification_panel.rs                            |  11 
crates/eval/src/eval.rs                                               |  13 
crates/eval/src/instance.rs                                           |   1 
crates/http_client/Cargo.toml                                         |   1 
crates/http_client/src/http_client.rs                                 |  71 
crates/inline_completion_button/src/inline_completion_button.rs       |  18 
crates/language_model/src/model/cloud_model.rs                        |   5 
crates/language_models/Cargo.toml                                     |   1 
crates/language_models/src/language_models.rs                         |  25 
crates/language_models/src/provider/cloud.rs                          |  63 
crates/onboarding/src/ai_setup_page.rs                                |   2 
crates/onboarding/src/onboarding.rs                                   |   7 
crates/project/src/project.rs                                         |   5 
crates/title_bar/src/title_bar.rs                                     |  26 
crates/workspace/src/workspace.rs                                     |  14 
crates/zed/src/main.rs                                                |  48 
crates/zed/src/zed.rs                                                 |   9 
crates/zed/src/zed/component_preview.rs                               |   3 
crates/zed/src/zed/component_preview/preview_support/active_thread.rs |  23 
crates/zed/src/zed/inline_completion_registry.rs                      |  28 
crates/zeta/src/zeta.rs                                               |  94 
46 files changed, 645 insertions(+), 847 deletions(-)

Detailed changes

Cargo.lock 🔗

@@ -114,7 +114,6 @@ dependencies = [
  "pretty_assertions",
  "project",
  "prompt_store",
- "proto",
  "rand 0.8.5",
  "ref-cast",
  "rope",
@@ -359,7 +358,6 @@ dependencies = [
  "component",
  "gpui",
  "language_model",
- "proto",
  "serde",
  "smallvec",
  "telemetry",
@@ -1076,17 +1074,6 @@ dependencies = [
  "tracing",
 ]
 
-[[package]]
-name = "async-recursion"
-version = "0.3.2"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d7d78656ba01f1b93024b7c3a0467f1608e4be67d725749fdcd7d2c7678fd7a2"
-dependencies = [
- "proc-macro2",
- "quote",
- "syn 1.0.109",
-]
-
 [[package]]
 name = "async-recursion"
 version = "1.1.1"
@@ -2972,7 +2959,6 @@ name = "client"
 version = "0.1.0"
 dependencies = [
  "anyhow",
- "async-recursion 0.3.2",
  "async-tungstenite",
  "base64 0.22.1",
  "chrono",
@@ -7814,6 +7800,7 @@ dependencies = [
  "http 1.3.1",
  "http-body 1.0.1",
  "log",
+ "parking_lot",
  "serde",
  "serde_json",
  "url",
@@ -9085,7 +9072,6 @@ dependencies = [
  "open_router",
  "partial-json-fixer",
  "project",
- "proto",
  "release_channel",
  "schemars",
  "serde",
@@ -9823,7 +9809,7 @@ name = "markdown_preview"
 version = "0.1.0"
 dependencies = [
  "anyhow",
- "async-recursion 1.1.1",
+ "async-recursion",
  "collections",
  "editor",
  "fs",
@@ -16192,7 +16178,7 @@ version = "0.1.0"
 dependencies = [
  "anyhow",
  "assistant_slash_command",
- "async-recursion 1.1.1",
+ "async-recursion",
  "breadcrumbs",
  "client",
  "collections",
@@ -19617,7 +19603,7 @@ version = "0.1.0"
 dependencies = [
  "any_vec",
  "anyhow",
- "async-recursion 1.1.1",
+ "async-recursion",
  "bincode",
  "call",
  "client",
@@ -20142,7 +20128,7 @@ dependencies = [
  "async-io",
  "async-lock",
  "async-process",
- "async-recursion 1.1.1",
+ "async-recursion",
  "async-task",
  "async-trait",
  "blocking",

crates/agent/Cargo.toml 🔗

@@ -47,7 +47,6 @@ paths.workspace = true
 postage.workspace = true
 project.workspace = true
 prompt_store.workspace = true
-proto.workspace = true
 ref-cast.workspace = true
 rope.workspace = true
 schemars.workspace = true

crates/agent/src/thread.rs 🔗

@@ -12,8 +12,8 @@ use agent_settings::{AgentProfileId, AgentSettings, CompletionMode};
 use anyhow::{Result, anyhow};
 use assistant_tool::{ActionLog, AnyToolCard, Tool, ToolWorkingSet};
 use chrono::{DateTime, Utc};
-use client::{CloudUserStore, ModelRequestUsage, RequestUsage};
-use cloud_llm_client::{CompletionIntent, CompletionRequestStatus, UsageLimit};
+use client::{ModelRequestUsage, RequestUsage};
+use cloud_llm_client::{CompletionIntent, CompletionRequestStatus, Plan, UsageLimit};
 use collections::HashMap;
 use feature_flags::{self, FeatureFlagAppExt};
 use futures::{FutureExt, StreamExt as _, future::Shared};
@@ -37,7 +37,6 @@ use project::{
     git_store::{GitStore, GitStoreCheckpoint, RepositoryState},
 };
 use prompt_store::{ModelContext, PromptBuilder};
-use proto::Plan;
 use schemars::JsonSchema;
 use serde::{Deserialize, Serialize};
 use settings::Settings;
@@ -374,7 +373,6 @@ pub struct Thread {
     completion_count: usize,
     pending_completions: Vec<PendingCompletion>,
     project: Entity<Project>,
-    cloud_user_store: Entity<CloudUserStore>,
     prompt_builder: Arc<PromptBuilder>,
     tools: Entity<ToolWorkingSet>,
     tool_use: ToolUseState,
@@ -445,7 +443,6 @@ pub struct ExceededWindowError {
 impl Thread {
     pub fn new(
         project: Entity<Project>,
-        cloud_user_store: Entity<CloudUserStore>,
         tools: Entity<ToolWorkingSet>,
         prompt_builder: Arc<PromptBuilder>,
         system_prompt: SharedProjectContext,
@@ -472,7 +469,6 @@ impl Thread {
             completion_count: 0,
             pending_completions: Vec::new(),
             project: project.clone(),
-            cloud_user_store,
             prompt_builder,
             tools: tools.clone(),
             last_restore_checkpoint: None,
@@ -506,7 +502,6 @@ impl Thread {
         id: ThreadId,
         serialized: SerializedThread,
         project: Entity<Project>,
-        cloud_user_store: Entity<CloudUserStore>,
         tools: Entity<ToolWorkingSet>,
         prompt_builder: Arc<PromptBuilder>,
         project_context: SharedProjectContext,
@@ -607,7 +602,6 @@ impl Thread {
             last_restore_checkpoint: None,
             pending_checkpoint: None,
             project: project.clone(),
-            cloud_user_store,
             prompt_builder,
             tools: tools.clone(),
             tool_use,
@@ -3260,15 +3254,18 @@ impl Thread {
     }
 
     fn update_model_request_usage(&self, amount: u32, limit: UsageLimit, cx: &mut Context<Self>) {
-        self.cloud_user_store.update(cx, |cloud_user_store, cx| {
-            cloud_user_store.update_model_request_usage(
-                ModelRequestUsage(RequestUsage {
-                    amount: amount as i32,
-                    limit,
-                }),
-                cx,
-            )
-        });
+        self.project
+            .read(cx)
+            .user_store()
+            .update(cx, |user_store, cx| {
+                user_store.update_model_request_usage(
+                    ModelRequestUsage(RequestUsage {
+                        amount: amount as i32,
+                        limit,
+                    }),
+                    cx,
+                )
+            });
     }
 
     pub fn deny_tool_use(
@@ -3886,7 +3883,6 @@ fn main() {{
                     thread.id.clone(),
                     serialized,
                     thread.project.clone(),
-                    thread.cloud_user_store.clone(),
                     thread.tools.clone(),
                     thread.prompt_builder.clone(),
                     thread.project_context.clone(),
@@ -5483,16 +5479,10 @@ fn main() {{
         let (workspace, cx) =
             cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
 
-        let (client, user_store) =
-            project.read_with(cx, |project, _cx| (project.client(), project.user_store()));
-        let cloud_user_store =
-            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store, cx));
-
         let thread_store = cx
             .update(|_, cx| {
                 ThreadStore::load(
                     project.clone(),
-                    cloud_user_store,
                     cx.new(|_| ToolWorkingSet::default()),
                     None,
                     Arc::new(PromptBuilder::new(None).unwrap()),

crates/agent/src/thread_store.rs 🔗

@@ -8,7 +8,6 @@ use agent_settings::{AgentProfileId, CompletionMode};
 use anyhow::{Context as _, Result, anyhow};
 use assistant_tool::{Tool, ToolId, ToolWorkingSet};
 use chrono::{DateTime, Utc};
-use client::CloudUserStore;
 use collections::HashMap;
 use context_server::ContextServerId;
 use futures::{
@@ -105,7 +104,6 @@ pub type TextThreadStore = assistant_context::ContextStore;
 
 pub struct ThreadStore {
     project: Entity<Project>,
-    cloud_user_store: Entity<CloudUserStore>,
     tools: Entity<ToolWorkingSet>,
     prompt_builder: Arc<PromptBuilder>,
     prompt_store: Option<Entity<PromptStore>>,
@@ -126,7 +124,6 @@ impl EventEmitter<RulesLoadingError> for ThreadStore {}
 impl ThreadStore {
     pub fn load(
         project: Entity<Project>,
-        cloud_user_store: Entity<CloudUserStore>,
         tools: Entity<ToolWorkingSet>,
         prompt_store: Option<Entity<PromptStore>>,
         prompt_builder: Arc<PromptBuilder>,
@@ -136,14 +133,8 @@ impl ThreadStore {
             let (thread_store, ready_rx) = cx.update(|cx| {
                 let mut option_ready_rx = None;
                 let thread_store = cx.new(|cx| {
-                    let (thread_store, ready_rx) = Self::new(
-                        project,
-                        cloud_user_store,
-                        tools,
-                        prompt_builder,
-                        prompt_store,
-                        cx,
-                    );
+                    let (thread_store, ready_rx) =
+                        Self::new(project, tools, prompt_builder, prompt_store, cx);
                     option_ready_rx = Some(ready_rx);
                     thread_store
                 });
@@ -156,7 +147,6 @@ impl ThreadStore {
 
     fn new(
         project: Entity<Project>,
-        cloud_user_store: Entity<CloudUserStore>,
         tools: Entity<ToolWorkingSet>,
         prompt_builder: Arc<PromptBuilder>,
         prompt_store: Option<Entity<PromptStore>>,
@@ -200,7 +190,6 @@ impl ThreadStore {
 
         let this = Self {
             project,
-            cloud_user_store,
             tools,
             prompt_builder,
             prompt_store,
@@ -418,7 +407,6 @@ impl ThreadStore {
         cx.new(|cx| {
             Thread::new(
                 self.project.clone(),
-                self.cloud_user_store.clone(),
                 self.tools.clone(),
                 self.prompt_builder.clone(),
                 self.project_context.clone(),
@@ -437,7 +425,6 @@ impl ThreadStore {
                 ThreadId::new(),
                 serialized,
                 self.project.clone(),
-                self.cloud_user_store.clone(),
                 self.tools.clone(),
                 self.prompt_builder.clone(),
                 self.project_context.clone(),
@@ -469,7 +456,6 @@ impl ThreadStore {
                         id.clone(),
                         thread,
                         this.project.clone(),
-                        this.cloud_user_store.clone(),
                         this.tools.clone(),
                         this.prompt_builder.clone(),
                         this.project_context.clone(),

crates/agent_ui/src/active_thread.rs 🔗

@@ -3820,7 +3820,6 @@ mod tests {
     use super::*;
     use agent::{MessageSegment, context::ContextLoadResult, thread_store};
     use assistant_tool::{ToolRegistry, ToolWorkingSet};
-    use client::CloudUserStore;
     use editor::EditorSettings;
     use fs::FakeFs;
     use gpui::{AppContext, TestAppContext, VisualTestContext};
@@ -4117,16 +4116,10 @@ mod tests {
         let (workspace, cx) =
             cx.add_window_view(|window, cx| Workspace::test_new(project.clone(), window, cx));
 
-        let (client, user_store) =
-            project.read_with(cx, |project, _cx| (project.client(), project.user_store()));
-        let cloud_user_store =
-            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store, cx));
-
         let thread_store = cx
             .update(|_, cx| {
                 ThreadStore::load(
                     project.clone(),
-                    cloud_user_store,
                     cx.new(|_| ToolWorkingSet::default()),
                     None,
                     Arc::new(PromptBuilder::new(None).unwrap()),

crates/agent_ui/src/agent_configuration.rs 🔗

@@ -7,6 +7,7 @@ use std::{sync::Arc, time::Duration};
 
 use agent_settings::AgentSettings;
 use assistant_tool::{ToolSource, ToolWorkingSet};
+use cloud_llm_client::Plan;
 use collections::HashMap;
 use context_server::ContextServerId;
 use extension::ExtensionManifest;
@@ -25,7 +26,6 @@ use project::{
     context_server_store::{ContextServerConfiguration, ContextServerStatus, ContextServerStore},
     project_settings::{ContextServerSettings, ProjectSettings},
 };
-use proto::Plan;
 use settings::{Settings, update_settings_file};
 use ui::{
     Chip, ContextMenu, Disclosure, Divider, DividerColor, ElevationIndex, Indicator, PopoverMenu,
@@ -180,7 +180,7 @@ impl AgentConfiguration {
         let current_plan = if is_zed_provider {
             self.workspace
                 .upgrade()
-                .and_then(|workspace| workspace.read(cx).user_store().read(cx).current_plan())
+                .and_then(|workspace| workspace.read(cx).user_store().read(cx).plan())
         } else {
             None
         };
@@ -508,7 +508,7 @@ impl AgentConfiguration {
                 .blend(cx.theme().colors().text_accent.opacity(0.2));
 
             let (plan_name, label_color, bg_color) = match plan {
-                Plan::Free => ("Free", Color::Default, free_chip_bg),
+                Plan::ZedFree => ("Free", Color::Default, free_chip_bg),
                 Plan::ZedProTrial => ("Pro Trial", Color::Accent, pro_chip_bg),
                 Plan::ZedPro => ("Pro", Color::Accent, pro_chip_bg),
             };

crates/agent_ui/src/agent_diff.rs 🔗

@@ -1896,7 +1896,6 @@ mod tests {
     use agent::thread_store::{self, ThreadStore};
     use agent_settings::AgentSettings;
     use assistant_tool::ToolWorkingSet;
-    use client::CloudUserStore;
     use editor::EditorSettings;
     use gpui::{TestAppContext, UpdateGlobal, VisualTestContext};
     use project::{FakeFs, Project};
@@ -1936,17 +1935,11 @@ mod tests {
             })
             .unwrap();
 
-        let (client, user_store) =
-            project.read_with(cx, |project, _cx| (project.client(), project.user_store()));
-        let cloud_user_store =
-            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store, cx));
-
         let prompt_store = None;
         let thread_store = cx
             .update(|cx| {
                 ThreadStore::load(
                     project.clone(),
-                    cloud_user_store,
                     cx.new(|_| ToolWorkingSet::default()),
                     prompt_store,
                     Arc::new(PromptBuilder::new(None).unwrap()),
@@ -2108,17 +2101,11 @@ mod tests {
             })
             .unwrap();
 
-        let (client, user_store) =
-            project.read_with(cx, |project, _cx| (project.client(), project.user_store()));
-        let cloud_user_store =
-            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store, cx));
-
         let prompt_store = None;
         let thread_store = cx
             .update(|cx| {
                 ThreadStore::load(
                     project.clone(),
-                    cloud_user_store,
                     cx.new(|_| ToolWorkingSet::default()),
                     prompt_store,
                     Arc::new(PromptBuilder::new(None).unwrap()),

crates/agent_ui/src/agent_panel.rs 🔗

@@ -43,8 +43,8 @@ use anyhow::{Result, anyhow};
 use assistant_context::{AssistantContext, ContextEvent, ContextSummary};
 use assistant_slash_command::SlashCommandWorkingSet;
 use assistant_tool::ToolWorkingSet;
-use client::{CloudUserStore, DisableAiSettings, UserStore, zed_urls};
-use cloud_llm_client::{CompletionIntent, UsageLimit};
+use client::{DisableAiSettings, UserStore, zed_urls};
+use cloud_llm_client::{CompletionIntent, Plan, UsageLimit};
 use editor::{Anchor, AnchorRangeExt as _, Editor, EditorEvent, MultiBuffer};
 use feature_flags::{self, FeatureFlagAppExt};
 use fs::Fs;
@@ -60,7 +60,6 @@ use language_model::{
 };
 use project::{Project, ProjectPath, Worktree};
 use prompt_store::{PromptBuilder, PromptStore, UserPromptId};
-use proto::Plan;
 use rules_library::{RulesLibrary, open_rules_library};
 use search::{BufferSearchBar, buffer_search};
 use settings::{Settings, update_settings_file};
@@ -427,7 +426,6 @@ impl ActiveView {
 pub struct AgentPanel {
     workspace: WeakEntity<Workspace>,
     user_store: Entity<UserStore>,
-    cloud_user_store: Entity<CloudUserStore>,
     project: Entity<Project>,
     fs: Arc<dyn Fs>,
     language_registry: Arc<LanguageRegistry>,
@@ -487,7 +485,6 @@ impl AgentPanel {
                     let project = workspace.project().clone();
                     ThreadStore::load(
                         project,
-                        workspace.app_state().cloud_user_store.clone(),
                         tools.clone(),
                         prompt_store.clone(),
                         prompt_builder.clone(),
@@ -555,7 +552,6 @@ impl AgentPanel {
         let thread = thread_store.update(cx, |this, cx| this.create_thread(cx));
         let fs = workspace.app_state().fs.clone();
         let user_store = workspace.app_state().user_store.clone();
-        let cloud_user_store = workspace.app_state().cloud_user_store.clone();
         let project = workspace.project();
         let language_registry = project.read(cx).languages().clone();
         let client = workspace.client().clone();
@@ -582,7 +578,6 @@ impl AgentPanel {
             MessageEditor::new(
                 fs.clone(),
                 workspace.clone(),
-                cloud_user_store.clone(),
                 message_editor_context_store.clone(),
                 prompt_store.clone(),
                 thread_store.downgrade(),
@@ -697,7 +692,6 @@ impl AgentPanel {
         let onboarding = cx.new(|cx| {
             AgentPanelOnboarding::new(
                 user_store.clone(),
-                cloud_user_store.clone(),
                 client,
                 |_window, cx| {
                     OnboardingUpsell::set_dismissed(true, cx);
@@ -710,7 +704,6 @@ impl AgentPanel {
             active_view,
             workspace,
             user_store,
-            cloud_user_store,
             project: project.clone(),
             fs: fs.clone(),
             language_registry,
@@ -853,7 +846,6 @@ impl AgentPanel {
             MessageEditor::new(
                 self.fs.clone(),
                 self.workspace.clone(),
-                self.cloud_user_store.clone(),
                 context_store.clone(),
                 self.prompt_store.clone(),
                 self.thread_store.downgrade(),
@@ -1127,7 +1119,6 @@ impl AgentPanel {
             MessageEditor::new(
                 self.fs.clone(),
                 self.workspace.clone(),
-                self.cloud_user_store.clone(),
                 context_store,
                 self.prompt_store.clone(),
                 self.thread_store.downgrade(),
@@ -1826,8 +1817,8 @@ impl AgentPanel {
     }
 
     fn render_toolbar(&self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
-        let cloud_user_store = self.cloud_user_store.read(cx);
-        let usage = cloud_user_store.model_request_usage();
+        let user_store = self.user_store.read(cx);
+        let usage = user_store.model_request_usage();
 
         let account_url = zed_urls::account_url(cx);
 
@@ -2298,10 +2289,10 @@ impl AgentPanel {
             | ActiveView::Configuration => return false,
         }
 
-        let plan = self.user_store.read(cx).current_plan();
+        let plan = self.user_store.read(cx).plan();
         let has_previous_trial = self.user_store.read(cx).trial_started_at().is_some();
 
-        matches!(plan, Some(Plan::Free)) && has_previous_trial
+        matches!(plan, Some(Plan::ZedFree)) && has_previous_trial
     }
 
     fn should_render_onboarding(&self, cx: &mut Context<Self>) -> bool {
@@ -2916,7 +2907,7 @@ impl AgentPanel {
     ) -> AnyElement {
         let error_message = match plan {
             Plan::ZedPro => "Upgrade to usage-based billing for more prompts.",
-            Plan::ZedProTrial | Plan::Free => "Upgrade to Zed Pro for more prompts.",
+            Plan::ZedProTrial | Plan::ZedFree => "Upgrade to Zed Pro for more prompts.",
         };
 
         let icon = Icon::new(IconName::XCircle)

crates/agent_ui/src/message_editor.rs 🔗

@@ -17,7 +17,6 @@ use agent::{
 use agent_settings::{AgentSettings, CompletionMode};
 use ai_onboarding::ApiKeysWithProviders;
 use buffer_diff::BufferDiff;
-use client::CloudUserStore;
 use cloud_llm_client::CompletionIntent;
 use collections::{HashMap, HashSet};
 use editor::actions::{MoveUp, Paste};
@@ -78,7 +77,6 @@ pub struct MessageEditor {
     editor: Entity<Editor>,
     workspace: WeakEntity<Workspace>,
     project: Entity<Project>,
-    cloud_user_store: Entity<CloudUserStore>,
     context_store: Entity<ContextStore>,
     prompt_store: Option<Entity<PromptStore>>,
     history_store: Option<WeakEntity<HistoryStore>>,
@@ -158,7 +156,6 @@ impl MessageEditor {
     pub fn new(
         fs: Arc<dyn Fs>,
         workspace: WeakEntity<Workspace>,
-        cloud_user_store: Entity<CloudUserStore>,
         context_store: Entity<ContextStore>,
         prompt_store: Option<Entity<PromptStore>>,
         thread_store: WeakEntity<ThreadStore>,
@@ -230,7 +227,6 @@ impl MessageEditor {
         Self {
             editor: editor.clone(),
             project: thread.read(cx).project().clone(),
-            cloud_user_store,
             thread,
             incompatible_tools_state: incompatible_tools.clone(),
             workspace,
@@ -1286,16 +1282,14 @@ impl MessageEditor {
             return None;
         }
 
-        let cloud_user_store = self.cloud_user_store.read(cx);
-        if cloud_user_store.is_usage_based_billing_enabled() {
+        let user_store = self.project.read(cx).user_store().read(cx);
+        if user_store.is_usage_based_billing_enabled() {
             return None;
         }
 
-        let plan = cloud_user_store
-            .plan()
-            .unwrap_or(cloud_llm_client::Plan::ZedFree);
+        let plan = user_store.plan().unwrap_or(cloud_llm_client::Plan::ZedFree);
 
-        let usage = cloud_user_store.model_request_usage()?;
+        let usage = user_store.model_request_usage()?;
 
         Some(
             div()
@@ -1758,7 +1752,6 @@ impl AgentPreview for MessageEditor {
     ) -> Option<AnyElement> {
         if let Some(workspace) = workspace.upgrade() {
             let fs = workspace.read(cx).app_state().fs.clone();
-            let cloud_user_store = workspace.read(cx).app_state().cloud_user_store.clone();
             let project = workspace.read(cx).project().clone();
             let weak_project = project.downgrade();
             let context_store = cx.new(|_cx| ContextStore::new(weak_project, None));
@@ -1771,7 +1764,6 @@ impl AgentPreview for MessageEditor {
                 MessageEditor::new(
                     fs,
                     workspace.downgrade(),
-                    cloud_user_store,
                     context_store,
                     None,
                     thread_store.downgrade(),

crates/ai_onboarding/Cargo.toml 🔗

@@ -20,7 +20,6 @@ cloud_llm_client.workspace = true
 component.workspace = true
 gpui.workspace = true
 language_model.workspace = true
-proto.workspace = true
 serde.workspace = true
 smallvec.workspace = true
 telemetry.workspace = true

crates/ai_onboarding/src/agent_panel_onboarding_content.rs 🔗

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use client::{Client, CloudUserStore, UserStore};
+use client::{Client, UserStore};
 use cloud_llm_client::Plan;
 use gpui::{Entity, IntoElement, ParentElement};
 use language_model::{LanguageModelRegistry, ZED_CLOUD_PROVIDER_ID};
@@ -10,7 +10,6 @@ use crate::{AgentPanelOnboardingCard, ApiKeysWithoutProviders, ZedAiOnboarding};
 
 pub struct AgentPanelOnboarding {
     user_store: Entity<UserStore>,
-    cloud_user_store: Entity<CloudUserStore>,
     client: Arc<Client>,
     configured_providers: Vec<(IconName, SharedString)>,
     continue_with_zed_ai: Arc<dyn Fn(&mut Window, &mut App)>,
@@ -19,7 +18,6 @@ pub struct AgentPanelOnboarding {
 impl AgentPanelOnboarding {
     pub fn new(
         user_store: Entity<UserStore>,
-        cloud_user_store: Entity<CloudUserStore>,
         client: Arc<Client>,
         continue_with_zed_ai: impl Fn(&mut Window, &mut App) + 'static,
         cx: &mut Context<Self>,
@@ -39,7 +37,6 @@ impl AgentPanelOnboarding {
 
         Self {
             user_store,
-            cloud_user_store,
             client,
             configured_providers: Self::compute_available_providers(cx),
             continue_with_zed_ai: Arc::new(continue_with_zed_ai),
@@ -60,8 +57,8 @@ impl AgentPanelOnboarding {
 
 impl Render for AgentPanelOnboarding {
     fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
-        let enrolled_in_trial = self.cloud_user_store.read(cx).plan() == Some(Plan::ZedProTrial);
-        let is_pro_user = self.cloud_user_store.read(cx).plan() == Some(Plan::ZedPro);
+        let enrolled_in_trial = self.user_store.read(cx).plan() == Some(Plan::ZedProTrial);
+        let is_pro_user = self.user_store.read(cx).plan() == Some(Plan::ZedPro);
 
         AgentPanelOnboardingCard::new()
             .child(

crates/ai_onboarding/src/ai_onboarding.rs 🔗

@@ -9,6 +9,7 @@ pub use agent_api_keys_onboarding::{ApiKeysWithProviders, ApiKeysWithoutProvider
 pub use agent_panel_onboarding_card::AgentPanelOnboardingCard;
 pub use agent_panel_onboarding_content::AgentPanelOnboarding;
 pub use ai_upsell_card::AiUpsellCard;
+use cloud_llm_client::Plan;
 pub use edit_prediction_onboarding_content::EditPredictionOnboarding;
 pub use young_account_banner::YoungAccountBanner;
 
@@ -79,7 +80,7 @@ impl From<client::Status> for SignInStatus {
 pub struct ZedAiOnboarding {
     pub sign_in_status: SignInStatus,
     pub has_accepted_terms_of_service: bool,
-    pub plan: Option<proto::Plan>,
+    pub plan: Option<Plan>,
     pub account_too_young: bool,
     pub continue_with_zed_ai: Arc<dyn Fn(&mut Window, &mut App)>,
     pub sign_in: Arc<dyn Fn(&mut Window, &mut App)>,
@@ -99,8 +100,8 @@ impl ZedAiOnboarding {
 
         Self {
             sign_in_status: status.into(),
-            has_accepted_terms_of_service: store.current_user_has_accepted_terms().unwrap_or(false),
-            plan: store.current_plan(),
+            has_accepted_terms_of_service: store.has_accepted_terms_of_service(),
+            plan: store.plan(),
             account_too_young: store.account_too_young(),
             continue_with_zed_ai,
             accept_terms_of_service: Arc::new({
@@ -113,11 +114,9 @@ impl ZedAiOnboarding {
             sign_in: Arc::new(move |_window, cx| {
                 cx.spawn({
                     let client = client.clone();
-                    async move |cx| {
-                        client.authenticate_and_connect(true, cx).await;
-                    }
+                    async move |cx| client.sign_in_with_optional_connect(true, cx).await
                 })
-                .detach();
+                .detach_and_log_err(cx);
             }),
             dismiss_onboarding: None,
         }
@@ -411,9 +410,9 @@ impl RenderOnce for ZedAiOnboarding {
         if matches!(self.sign_in_status, SignInStatus::SignedIn) {
             if self.has_accepted_terms_of_service {
                 match self.plan {
-                    None | Some(proto::Plan::Free) => self.render_free_plan_state(cx),
-                    Some(proto::Plan::ZedProTrial) => self.render_trial_state(cx),
-                    Some(proto::Plan::ZedPro) => self.render_pro_plan_state(cx),
+                    None | Some(Plan::ZedFree) => self.render_free_plan_state(cx),
+                    Some(Plan::ZedProTrial) => self.render_trial_state(cx),
+                    Some(Plan::ZedPro) => self.render_pro_plan_state(cx),
                 }
             } else {
                 self.render_accept_terms_of_service()
@@ -433,7 +432,7 @@ impl Component for ZedAiOnboarding {
         fn onboarding(
             sign_in_status: SignInStatus,
             has_accepted_terms_of_service: bool,
-            plan: Option<proto::Plan>,
+            plan: Option<Plan>,
             account_too_young: bool,
         ) -> AnyElement {
             ZedAiOnboarding {
@@ -468,25 +467,15 @@ impl Component for ZedAiOnboarding {
                     ),
                     single_example(
                         "Free Plan",
-                        onboarding(SignInStatus::SignedIn, true, Some(proto::Plan::Free), false),
+                        onboarding(SignInStatus::SignedIn, true, Some(Plan::ZedFree), false),
                     ),
                     single_example(
                         "Pro Trial",
-                        onboarding(
-                            SignInStatus::SignedIn,
-                            true,
-                            Some(proto::Plan::ZedProTrial),
-                            false,
-                        ),
+                        onboarding(SignInStatus::SignedIn, true, Some(Plan::ZedProTrial), false),
                     ),
                     single_example(
                         "Pro Plan",
-                        onboarding(
-                            SignInStatus::SignedIn,
-                            true,
-                            Some(proto::Plan::ZedPro),
-                            false,
-                        ),
+                        onboarding(SignInStatus::SignedIn, true, Some(Plan::ZedPro), false),
                     ),
                 ])
                 .into_any_element(),

crates/ai_onboarding/src/ai_upsell_card.rs 🔗

@@ -24,11 +24,9 @@ impl AiUpsellCard {
             sign_in: Arc::new(move |_window, cx| {
                 cx.spawn({
                     let client = client.clone();
-                    async move |cx| {
-                        client.authenticate_and_connect(true, cx).await;
-                    }
+                    async move |cx| client.sign_in_with_optional_connect(true, cx).await
                 })
-                .detach();
+                .detach_and_log_err(cx);
             }),
         }
     }

crates/assistant_tools/src/edit_agent/evals.rs 🔗

@@ -7,7 +7,7 @@ use crate::{
 };
 use Role::*;
 use assistant_tool::ToolRegistry;
-use client::{Client, CloudUserStore, UserStore};
+use client::{Client, UserStore};
 use collections::HashMap;
 use fs::FakeFs;
 use futures::{FutureExt, future::LocalBoxFuture};
@@ -1470,14 +1470,12 @@ impl EditAgentTest {
             client::init_settings(cx);
             let client = Client::production(cx);
             let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
-            let cloud_user_store =
-                cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store.clone(), cx));
 
             settings::init(cx);
             Project::init_settings(cx);
             language::init(cx);
             language_model::init(client.clone(), cx);
-            language_models::init(user_store.clone(), cloud_user_store, client.clone(), cx);
+            language_models::init(user_store.clone(), client.clone(), cx);
             crate::init(client.http_client(), cx);
         });
 

crates/channel/src/channel_store_tests.rs 🔗

@@ -259,20 +259,6 @@ async fn test_channel_messages(cx: &mut TestAppContext) {
         assert_channels(&channel_store, &[(0, "the-channel".to_string())], cx);
     });
 
-    let get_users = server.receive::<proto::GetUsers>().await.unwrap();
-    assert_eq!(get_users.payload.user_ids, vec![5]);
-    server.respond(
-        get_users.receipt(),
-        proto::UsersResponse {
-            users: vec![proto::User {
-                id: 5,
-                github_login: "nathansobo".into(),
-                avatar_url: "http://avatar.com/nathansobo".into(),
-                name: None,
-            }],
-        },
-    );
-
     // Join a channel and populate its existing messages.
     let channel = channel_store.update(cx, |store, cx| {
         let channel_id = store.ordered_channels().next().unwrap().1.id;
@@ -334,7 +320,7 @@ async fn test_channel_messages(cx: &mut TestAppContext) {
                 .map(|message| (message.sender.github_login.clone(), message.body.clone()))
                 .collect::<Vec<_>>(),
             &[
-                ("nathansobo".into(), "a".into()),
+                ("user-5".into(), "a".into()),
                 ("maxbrunsfeld".into(), "b".into())
             ]
         );
@@ -437,7 +423,7 @@ async fn test_channel_messages(cx: &mut TestAppContext) {
                 .map(|message| (message.sender.github_login.clone(), message.body.clone()))
                 .collect::<Vec<_>>(),
             &[
-                ("nathansobo".into(), "y".into()),
+                ("user-5".into(), "y".into()),
                 ("maxbrunsfeld".into(), "z".into())
             ]
         );

crates/client/Cargo.toml 🔗

@@ -17,7 +17,6 @@ test-support = ["clock/test-support", "collections/test-support", "gpui/test-sup
 
 [dependencies]
 anyhow.workspace = true
-async-recursion = "0.3"
 async-tungstenite = { workspace = true, features = ["tokio", "tokio-rustls-manual-roots"] }
 base64.workspace = true
 chrono = { workspace = true, features = ["serde"] }

crates/client/src/client.rs 🔗

@@ -1,14 +1,12 @@
 #[cfg(any(test, feature = "test-support"))]
 pub mod test;
 
-mod cloud;
 mod proxy;
 pub mod telemetry;
 pub mod user;
 pub mod zed_urls;
 
 use anyhow::{Context as _, Result, anyhow};
-use async_recursion::async_recursion;
 use async_tungstenite::tungstenite::{
     client::IntoClientRequest,
     error::Error as WebsocketError,
@@ -52,7 +50,6 @@ use tokio::net::TcpStream;
 use url::Url;
 use util::{ConnectionResult, ResultExt};
 
-pub use cloud::*;
 pub use rpc::*;
 pub use telemetry_events::Event;
 pub use user::*;
@@ -164,20 +161,8 @@ pub fn init(client: &Arc<Client>, cx: &mut App) {
         let client = client.clone();
         move |_: &SignIn, cx| {
             if let Some(client) = client.upgrade() {
-                cx.spawn(
-                    async move |cx| match client.authenticate_and_connect(true, &cx).await {
-                        ConnectionResult::Timeout => {
-                            log::error!("Initial authentication timed out");
-                        }
-                        ConnectionResult::ConnectionReset => {
-                            log::error!("Initial authentication connection reset");
-                        }
-                        ConnectionResult::Result(r) => {
-                            r.log_err();
-                        }
-                    },
-                )
-                .detach();
+                cx.spawn(async move |cx| client.sign_in_with_optional_connect(true, &cx).await)
+                    .detach_and_log_err(cx);
             }
         }
     });
@@ -286,6 +271,8 @@ pub enum Status {
     SignedOut,
     UpgradeRequired,
     Authenticating,
+    Authenticated,
+    AuthenticationError,
     Connecting,
     ConnectionError,
     Connected {
@@ -712,7 +699,7 @@ impl Client {
 
                     let mut delay = INITIAL_RECONNECTION_DELAY;
                     loop {
-                        match client.authenticate_and_connect(true, &cx).await {
+                        match client.connect(true, &cx).await {
                             ConnectionResult::Timeout => {
                                 log::error!("client connect attempt timed out")
                             }
@@ -882,40 +869,50 @@ impl Client {
             .is_some()
     }
 
-    #[async_recursion(?Send)]
-    pub async fn authenticate_and_connect(
+    pub async fn sign_in(
         self: &Arc<Self>,
         try_provider: bool,
         cx: &AsyncApp,
-    ) -> ConnectionResult<()> {
-        let was_disconnected = match *self.status().borrow() {
-            Status::SignedOut => true,
-            Status::ConnectionError
-            | Status::ConnectionLost
-            | Status::Authenticating { .. }
-            | Status::Reauthenticating { .. }
-            | Status::ReconnectionError { .. } => false,
-            Status::Connected { .. } | Status::Connecting { .. } | Status::Reconnecting { .. } => {
-                return ConnectionResult::Result(Ok(()));
-            }
-            Status::UpgradeRequired => {
-                return ConnectionResult::Result(
-                    Err(EstablishConnectionError::UpgradeRequired)
-                        .context("client auth and connect"),
-                );
-            }
-        };
-        if was_disconnected {
+    ) -> Result<Credentials> {
+        if self.status().borrow().is_signed_out() {
             self.set_status(Status::Authenticating, cx);
         } else {
-            self.set_status(Status::Reauthenticating, cx)
+            self.set_status(Status::Reauthenticating, cx);
+        }
+
+        let mut credentials = None;
+
+        let old_credentials = self.state.read().credentials.clone();
+        if let Some(old_credentials) = old_credentials {
+            self.cloud_client.set_credentials(
+                old_credentials.user_id as u32,
+                old_credentials.access_token.clone(),
+            );
+
+            // Fetch the authenticated user with the old credentials, to ensure they are still valid.
+            if self.cloud_client.get_authenticated_user().await.is_ok() {
+                credentials = Some(old_credentials);
+            }
         }
 
-        let mut read_from_provider = false;
-        let mut credentials = self.state.read().credentials.clone();
         if credentials.is_none() && try_provider {
-            credentials = self.credentials_provider.read_credentials(cx).await;
-            read_from_provider = credentials.is_some();
+            if let Some(stored_credentials) = self.credentials_provider.read_credentials(cx).await {
+                self.cloud_client.set_credentials(
+                    stored_credentials.user_id as u32,
+                    stored_credentials.access_token.clone(),
+                );
+
+                // Fetch the authenticated user with the stored credentials, and
+                // clear them from the credentials provider if that fails.
+                if self.cloud_client.get_authenticated_user().await.is_ok() {
+                    credentials = Some(stored_credentials);
+                } else {
+                    self.credentials_provider
+                        .delete_credentials(cx)
+                        .await
+                        .log_err();
+                }
+            }
         }
 
         if credentials.is_none() {
@@ -924,22 +921,86 @@ impl Client {
             futures::select_biased! {
                 authenticate = self.authenticate(cx).fuse() => {
                     match authenticate {
-                        Ok(creds) => credentials = Some(creds),
+                        Ok(creds) => {
+                            if IMPERSONATE_LOGIN.is_none() {
+                                self.credentials_provider
+                                    .write_credentials(creds.user_id, creds.access_token.clone(), cx)
+                                    .await
+                                    .log_err();
+                            }
+
+                            credentials = Some(creds);
+                        },
                         Err(err) => {
-                            self.set_status(Status::ConnectionError, cx);
-                            return ConnectionResult::Result(Err(err));
+                            self.set_status(Status::AuthenticationError, cx);
+                            return Err(err);
                         }
                     }
                 }
                 _ = status_rx.next().fuse() => {
-                    return ConnectionResult::Result(Err(anyhow!("authentication canceled")));
+                    return Err(anyhow!("authentication canceled"));
                 }
             }
         }
+
         let credentials = credentials.unwrap();
         self.set_id(credentials.user_id);
         self.cloud_client
             .set_credentials(credentials.user_id as u32, credentials.access_token.clone());
+        self.state.write().credentials = Some(credentials.clone());
+        self.set_status(Status::Authenticated, cx);
+
+        Ok(credentials)
+    }
+
+    /// Performs a sign-in and also connects to Collab.
+    ///
+    /// This is called in places where we *don't* need to connect in the future. We will replace these calls with calls
+    /// to `sign_in` when we're ready to remove auto-connection to Collab.
+    pub async fn sign_in_with_optional_connect(
+        self: &Arc<Self>,
+        try_provider: bool,
+        cx: &AsyncApp,
+    ) -> Result<()> {
+        let credentials = self.sign_in(try_provider, cx).await?;
+
+        let connect_result = match self.connect_with_credentials(credentials, cx).await {
+            ConnectionResult::Timeout => Err(anyhow!("connection timed out")),
+            ConnectionResult::ConnectionReset => Err(anyhow!("connection reset")),
+            ConnectionResult::Result(result) => result.context("client auth and connect"),
+        };
+        connect_result.log_err();
+
+        Ok(())
+    }
+
+    pub async fn connect(
+        self: &Arc<Self>,
+        try_provider: bool,
+        cx: &AsyncApp,
+    ) -> ConnectionResult<()> {
+        let was_disconnected = match *self.status().borrow() {
+            Status::SignedOut | Status::Authenticated => true,
+            Status::ConnectionError
+            | Status::ConnectionLost
+            | Status::Authenticating { .. }
+            | Status::AuthenticationError
+            | Status::Reauthenticating { .. }
+            | Status::ReconnectionError { .. } => false,
+            Status::Connected { .. } | Status::Connecting { .. } | Status::Reconnecting { .. } => {
+                return ConnectionResult::Result(Ok(()));
+            }
+            Status::UpgradeRequired => {
+                return ConnectionResult::Result(
+                    Err(EstablishConnectionError::UpgradeRequired)
+                        .context("client auth and connect"),
+                );
+            }
+        };
+        let credentials = match self.sign_in(try_provider, cx).await {
+            Ok(credentials) => credentials,
+            Err(err) => return ConnectionResult::Result(Err(err)),
+        };
 
         if was_disconnected {
             self.set_status(Status::Connecting, cx);
@@ -947,17 +1008,20 @@ impl Client {
             self.set_status(Status::Reconnecting, cx);
         }
 
+        self.connect_with_credentials(credentials, cx).await
+    }
+
+    async fn connect_with_credentials(
+        self: &Arc<Self>,
+        credentials: Credentials,
+        cx: &AsyncApp,
+    ) -> ConnectionResult<()> {
         let mut timeout =
             futures::FutureExt::fuse(cx.background_executor().timer(CONNECTION_TIMEOUT));
         futures::select_biased! {
             connection = self.establish_connection(&credentials, cx).fuse() => {
                 match connection {
                     Ok(conn) => {
-                        self.state.write().credentials = Some(credentials.clone());
-                        if !read_from_provider && IMPERSONATE_LOGIN.is_none() {
-                            self.credentials_provider.write_credentials(credentials.user_id, credentials.access_token, cx).await.log_err();
-                        }
-
                         futures::select_biased! {
                             result = self.set_connection(conn, cx).fuse() => {
                                 match result.context("client auth and connect") {
@@ -975,15 +1039,8 @@ impl Client {
                         }
                     }
                     Err(EstablishConnectionError::Unauthorized) => {
-                        self.state.write().credentials.take();
-                        if read_from_provider {
-                            self.credentials_provider.delete_credentials(cx).await.log_err();
-                            self.set_status(Status::SignedOut, cx);
-                            self.authenticate_and_connect(false, cx).await
-                        } else {
-                            self.set_status(Status::ConnectionError, cx);
-                            ConnectionResult::Result(Err(EstablishConnectionError::Unauthorized).context("client auth and connect"))
-                        }
+                        self.set_status(Status::ConnectionError, cx);
+                        ConnectionResult::Result(Err(EstablishConnectionError::Unauthorized).context("client auth and connect"))
                     }
                     Err(EstablishConnectionError::UpgradeRequired) => {
                         self.set_status(Status::UpgradeRequired, cx);
@@ -1733,7 +1790,7 @@ mod tests {
         });
         let auth_and_connect = cx.spawn({
             let client = client.clone();
-            |cx| async move { client.authenticate_and_connect(false, &cx).await }
+            |cx| async move { client.connect(false, &cx).await }
         });
         executor.run_until_parked();
         assert!(matches!(status.next().await, Some(Status::Connecting)));
@@ -1810,7 +1867,7 @@ mod tests {
 
         let _authenticate = cx.spawn({
             let client = client.clone();
-            move |cx| async move { client.authenticate_and_connect(false, &cx).await }
+            move |cx| async move { client.connect(false, &cx).await }
         });
         executor.run_until_parked();
         assert_eq!(*auth_count.lock(), 1);
@@ -1818,7 +1875,7 @@ mod tests {
 
         let _authenticate = cx.spawn({
             let client = client.clone();
-            |cx| async move { client.authenticate_and_connect(false, &cx).await }
+            |cx| async move { client.connect(false, &cx).await }
         });
         executor.run_until_parked();
         assert_eq!(*auth_count.lock(), 2);

crates/client/src/cloud/user_store.rs 🔗

@@ -1,211 +0,0 @@
-use std::sync::Arc;
-use std::time::Duration;
-
-use anyhow::Context as _;
-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::user::Event as RpcUserStoreEvent;
-use crate::{EditPredictionUsage, ModelRequestUsage, RequestUsage, UserStore};
-
-pub struct CloudUserStore {
-    cloud_client: Arc<CloudApiClient>,
-    authenticated_user: Option<Arc<AuthenticatedUser>>,
-    plan_info: Option<Arc<PlanInfo>>,
-    model_request_usage: Option<ModelRequestUsage>,
-    edit_prediction_usage: Option<EditPredictionUsage>,
-    _maintain_authenticated_user_task: Task<()>,
-    _rpc_plan_updated_subscription: Subscription,
-}
-
-impl CloudUserStore {
-    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,
-            model_request_usage: None,
-            edit_prediction_usage: None,
-            _maintain_authenticated_user_task: cx.spawn(async move |this, cx| {
-                maybe!(async move {
-                    loop {
-                        let Some(this) = this.upgrade() else {
-                            return anyhow::Ok(());
-                        };
-
-                        if cloud_client.has_credentials() {
-                            let already_fetched_authenticated_user = this
-                                .read_with(cx, |this, _cx| this.authenticated_user().is_some())
-                                .unwrap_or(false);
-
-                            if already_fetched_authenticated_user {
-                                // We already fetched the authenticated user; nothing to do.
-                            } else {
-                                let authenticated_user_result = cloud_client
-                                    .get_authenticated_user()
-                                    .await
-                                    .context("failed to fetch authenticated user");
-                                if let Some(response) = authenticated_user_result.log_err() {
-                                    this.update(cx, |this, _cx| {
-                                        this.update_authenticated_user(response);
-                                    })
-                                    .ok();
-                                }
-                            }
-                        } else {
-                            this.update(cx, |this, _cx| {
-                                this.authenticated_user.take();
-                                this.plan_info.take();
-                            })
-                            .ok();
-                        }
-
-                        cx.background_executor()
-                            .timer(Duration::from_millis(100))
-                            .await;
-                    }
-                })
-                .await
-                .log_err();
-            }),
-            _rpc_plan_updated_subscription: rpc_plan_updated_subscription,
-        }
-    }
-
-    pub fn is_authenticated(&self) -> bool {
-        self.authenticated_user.is_some()
-    }
-
-    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,
-                )
-            })
-    }
-
-    pub fn trial_started_at(&self) -> Option<DateTime<Utc>> {
-        self.plan_info
-            .as_ref()
-            .and_then(|plan| plan.trial_started_at)
-            .map(|trial_started_at| trial_started_at.0)
-    }
-
-    pub fn has_accepted_tos(&self) -> bool {
-        self.authenticated_user
-            .as_ref()
-            .map(|user| user.accepted_tos_at.is_some())
-            .unwrap_or_default()
-    }
-
-    /// Returns whether the user's account is too new to use the service.
-    pub fn account_too_young(&self) -> bool {
-        self.plan_info
-            .as_ref()
-            .map(|plan| plan.is_account_too_young)
-            .unwrap_or_default()
-    }
-
-    /// Returns whether the current user has overdue invoices and usage should be blocked.
-    pub fn has_overdue_invoices(&self) -> bool {
-        self.plan_info
-            .as_ref()
-            .map(|plan| plan.has_overdue_invoices)
-            .unwrap_or_default()
-    }
-
-    pub fn is_usage_based_billing_enabled(&self) -> bool {
-        self.plan_info
-            .as_ref()
-            .map(|plan| plan.is_usage_based_billing_enabled)
-            .unwrap_or_default()
-    }
-
-    pub fn model_request_usage(&self) -> Option<ModelRequestUsage> {
-        self.model_request_usage
-    }
-
-    pub fn update_model_request_usage(&mut self, usage: ModelRequestUsage, cx: &mut Context<Self>) {
-        self.model_request_usage = Some(usage);
-        cx.notify();
-    }
-
-    pub fn edit_prediction_usage(&self) -> Option<EditPredictionUsage> {
-        self.edit_prediction_usage
-    }
-
-    pub fn update_edit_prediction_usage(
-        &mut self,
-        usage: EditPredictionUsage,
-        cx: &mut Context<Self>,
-    ) {
-        self.edit_prediction_usage = Some(usage);
-        cx.notify();
-    }
-
-    fn update_authenticated_user(&mut self, response: GetAuthenticatedUserResponse) {
-        self.authenticated_user = Some(Arc::new(response.user));
-        self.model_request_usage = Some(ModelRequestUsage(RequestUsage {
-            limit: response.plan.usage.model_requests.limit,
-            amount: response.plan.usage.model_requests.used as i32,
-        }));
-        self.edit_prediction_usage = Some(EditPredictionUsage(RequestUsage {
-            limit: response.plan.usage.edit_predictions.limit,
-            amount: response.plan.usage.edit_predictions.used as i32,
-        }));
-        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/test.rs 🔗

@@ -1,8 +1,11 @@
 use crate::{Client, Connection, Credentials, EstablishConnectionError, UserStore};
 use anyhow::{Context as _, Result, anyhow};
 use chrono::Duration;
+use cloud_api_client::{AuthenticatedUser, GetAuthenticatedUserResponse, PlanInfo};
+use cloud_llm_client::{CurrentUsage, Plan, UsageData, UsageLimit};
 use futures::{StreamExt, stream::BoxStream};
 use gpui::{AppContext as _, BackgroundExecutor, Entity, TestAppContext};
+use http_client::{AsyncBody, Method, Request, http};
 use parking_lot::Mutex;
 use rpc::{
     ConnectionId, Peer, Receipt, TypedEnvelope,
@@ -39,6 +42,44 @@ impl FakeServer {
             executor: cx.executor(),
         };
 
+        client.http_client().as_fake().replace_handler({
+            let state = server.state.clone();
+            move |old_handler, req| {
+                let state = state.clone();
+                let old_handler = old_handler.clone();
+                async move {
+                    match (req.method(), req.uri().path()) {
+                        (&Method::GET, "/client/users/me") => {
+                            let credentials = parse_authorization_header(&req);
+                            if credentials
+                                != Some(Credentials {
+                                    user_id: client_user_id,
+                                    access_token: state.lock().access_token.to_string(),
+                                })
+                            {
+                                return Ok(http_client::Response::builder()
+                                    .status(401)
+                                    .body("Unauthorized".into())
+                                    .unwrap());
+                            }
+
+                            Ok(http_client::Response::builder()
+                                .status(200)
+                                .body(
+                                    serde_json::to_string(&make_get_authenticated_user_response(
+                                        client_user_id as i32,
+                                        format!("user-{client_user_id}"),
+                                    ))
+                                    .unwrap()
+                                    .into(),
+                                )
+                                .unwrap())
+                        }
+                        _ => old_handler(req).await,
+                    }
+                }
+            }
+        });
         client
             .override_authenticate({
                 let state = Arc::downgrade(&server.state);
@@ -105,7 +146,7 @@ impl FakeServer {
             });
 
         client
-            .authenticate_and_connect(false, &cx.to_async())
+            .connect(false, &cx.to_async())
             .await
             .into_response()
             .unwrap();
@@ -223,3 +264,54 @@ impl Drop for FakeServer {
         self.disconnect();
     }
 }
+
+pub fn parse_authorization_header(req: &Request<AsyncBody>) -> Option<Credentials> {
+    let mut auth_header = req
+        .headers()
+        .get(http::header::AUTHORIZATION)?
+        .to_str()
+        .ok()?
+        .split_whitespace();
+    let user_id = auth_header.next()?.parse().ok()?;
+    let access_token = auth_header.next()?;
+    Some(Credentials {
+        user_id,
+        access_token: access_token.to_string(),
+    })
+}
+
+pub fn make_get_authenticated_user_response(
+    user_id: i32,
+    github_login: String,
+) -> GetAuthenticatedUserResponse {
+    GetAuthenticatedUserResponse {
+        user: AuthenticatedUser {
+            id: user_id,
+            metrics_id: format!("metrics-id-{user_id}"),
+            avatar_url: "".to_string(),
+            github_login,
+            name: None,
+            is_staff: false,
+            accepted_tos_at: None,
+        },
+        feature_flags: vec![],
+        plan: PlanInfo {
+            plan: Plan::ZedPro,
+            subscription_period: None,
+            usage: CurrentUsage {
+                model_requests: UsageData {
+                    used: 0,
+                    limit: UsageLimit::Limited(500),
+                },
+                edit_predictions: UsageData {
+                    used: 250,
+                    limit: UsageLimit::Unlimited,
+                },
+            },
+            trial_started_at: None,
+            is_usage_based_billing_enabled: false,
+            is_account_too_young: false,
+            has_overdue_invoices: false,
+        },
+    }
+}

crates/client/src/user.rs 🔗

@@ -1,6 +1,7 @@
 use super::{Client, Status, TypedEnvelope, proto};
 use anyhow::{Context as _, Result, anyhow};
 use chrono::{DateTime, Utc};
+use cloud_api_client::{GetAuthenticatedUserResponse, PlanInfo};
 use cloud_llm_client::{
     EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME, EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME,
     MODEL_REQUESTS_USAGE_AMOUNT_HEADER_NAME, MODEL_REQUESTS_USAGE_LIMIT_HEADER_NAME, UsageLimit,
@@ -20,7 +21,7 @@ use std::{
     sync::{Arc, Weak},
 };
 use text::ReplicaId;
-use util::TryFutureExt as _;
+use util::{ResultExt, TryFutureExt as _};
 
 pub type UserId = u64;
 
@@ -110,12 +111,11 @@ pub struct UserStore {
     by_github_login: HashMap<SharedString, u64>,
     participant_indices: HashMap<u64, ParticipantIndex>,
     update_contacts_tx: mpsc::UnboundedSender<UpdateContacts>,
-    current_plan: Option<proto::Plan>,
-    trial_started_at: Option<DateTime<Utc>>,
-    is_usage_based_billing_enabled: Option<bool>,
-    account_too_young: Option<bool>,
+    model_request_usage: Option<ModelRequestUsage>,
+    edit_prediction_usage: Option<EditPredictionUsage>,
+    plan_info: Option<PlanInfo>,
     current_user: watch::Receiver<Option<Arc<User>>>,
-    accepted_tos_at: Option<Option<DateTime<Utc>>>,
+    accepted_tos_at: Option<Option<cloud_api_client::Timestamp>>,
     contacts: Vec<Arc<Contact>>,
     incoming_contact_requests: Vec<Arc<User>>,
     outgoing_contact_requests: Vec<Arc<User>>,
@@ -185,10 +185,9 @@ impl UserStore {
             users: Default::default(),
             by_github_login: Default::default(),
             current_user: current_user_rx,
-            current_plan: None,
-            trial_started_at: None,
-            is_usage_based_billing_enabled: None,
-            account_too_young: None,
+            plan_info: None,
+            model_request_usage: None,
+            edit_prediction_usage: None,
             accepted_tos_at: None,
             contacts: Default::default(),
             incoming_contact_requests: Default::default(),
@@ -218,53 +217,30 @@ impl UserStore {
                         return Ok(());
                     };
                     match status {
-                        Status::Connected { .. } => {
+                        Status::Authenticated | Status::Connected { .. } => {
                             if let Some(user_id) = client.user_id() {
-                                let fetch_user = if let Ok(fetch_user) =
-                                    this.update(cx, |this, cx| this.get_user(user_id, cx).log_err())
-                                {
-                                    fetch_user
-                                } else {
-                                    break;
-                                };
-                                let fetch_private_user_info =
-                                    client.request(proto::GetPrivateUserInfo {}).log_err();
-                                let (user, info) =
-                                    futures::join!(fetch_user, fetch_private_user_info);
-
+                                let response = client.cloud_client().get_authenticated_user().await;
+                                let mut current_user = None;
                                 cx.update(|cx| {
-                                    if let Some(info) = info {
-                                        let staff =
-                                            info.staff && !*feature_flags::ZED_DISABLE_STAFF;
-                                        cx.update_flags(staff, info.flags);
-                                        client.telemetry.set_authenticated_user_info(
-                                            Some(info.metrics_id.clone()),
-                                            staff,
-                                        );
-
+                                    if let Some(response) = response.log_err() {
+                                        let user = Arc::new(User {
+                                            id: user_id,
+                                            github_login: response.user.github_login.clone().into(),
+                                            avatar_uri: response.user.avatar_url.clone().into(),
+                                            name: response.user.name.clone(),
+                                        });
+                                        current_user = Some(user.clone());
                                         this.update(cx, |this, cx| {
-                                            let accepted_tos_at = {
-                                                #[cfg(debug_assertions)]
-                                                if std::env::var("ZED_IGNORE_ACCEPTED_TOS").is_ok()
-                                                {
-                                                    None
-                                                } else {
-                                                    info.accepted_tos_at
-                                                }
-
-                                                #[cfg(not(debug_assertions))]
-                                                info.accepted_tos_at
-                                            };
-
-                                            this.set_current_user_accepted_tos_at(accepted_tos_at);
-                                            cx.emit(Event::PrivateUserInfoUpdated);
+                                            this.by_github_login
+                                                .insert(user.github_login.clone(), user_id);
+                                            this.users.insert(user_id, user);
+                                            this.update_authenticated_user(response, cx)
                                         })
                                     } else {
                                         anyhow::Ok(())
                                     }
                                 })??;
-
-                                current_user_tx.send(user).await.ok();
+                                current_user_tx.send(current_user).await.ok();
 
                                 this.update(cx, |_, cx| cx.notify())?;
                             }
@@ -345,22 +321,22 @@ impl UserStore {
 
     async fn handle_update_plan(
         this: Entity<Self>,
-        message: TypedEnvelope<proto::UpdateUserPlan>,
+        _message: TypedEnvelope<proto::UpdateUserPlan>,
         mut cx: AsyncApp,
     ) -> Result<()> {
+        let client = this
+            .read_with(&cx, |this, _| this.client.upgrade())?
+            .context("client was dropped")?;
+
+        let response = client
+            .cloud_client()
+            .get_authenticated_user()
+            .await
+            .context("failed to fetch authenticated user")?;
+
         this.update(&mut cx, |this, cx| {
-            this.current_plan = Some(message.payload.plan());
-            this.trial_started_at = message
-                .payload
-                .trial_started_at
-                .and_then(|trial_started_at| DateTime::from_timestamp(trial_started_at as i64, 0));
-            this.is_usage_based_billing_enabled = message.payload.is_usage_based_billing_enabled;
-            this.account_too_young = message.payload.account_too_young;
-
-            cx.emit(Event::PlanUpdated);
-            cx.notify();
-        })?;
-        Ok(())
+            this.update_authenticated_user(response, cx);
+        })
     }
 
     fn update_contacts(&mut self, message: UpdateContacts, cx: &Context<Self>) -> Task<Result<()>> {
@@ -719,42 +695,131 @@ impl UserStore {
         self.current_user.borrow().clone()
     }
 
-    pub fn current_plan(&self) -> Option<proto::Plan> {
+    pub fn plan(&self) -> Option<cloud_llm_client::Plan> {
         #[cfg(debug_assertions)]
         if let Ok(plan) = std::env::var("ZED_SIMULATE_PLAN").as_ref() {
             return match plan.as_str() {
-                "free" => Some(proto::Plan::Free),
-                "trial" => Some(proto::Plan::ZedProTrial),
-                "pro" => Some(proto::Plan::ZedPro),
+                "free" => Some(cloud_llm_client::Plan::ZedFree),
+                "trial" => Some(cloud_llm_client::Plan::ZedProTrial),
+                "pro" => Some(cloud_llm_client::Plan::ZedPro),
                 _ => {
                     panic!("ZED_SIMULATE_PLAN must be one of 'free', 'trial', or 'pro'");
                 }
             };
         }
 
-        self.current_plan
+        self.plan_info.as_ref().map(|info| info.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,
+                )
+            })
     }
 
     pub fn trial_started_at(&self) -> Option<DateTime<Utc>> {
-        self.trial_started_at
+        self.plan_info
+            .as_ref()
+            .and_then(|plan| plan.trial_started_at)
+            .map(|trial_started_at| trial_started_at.0)
     }
 
-    pub fn usage_based_billing_enabled(&self) -> Option<bool> {
-        self.is_usage_based_billing_enabled
+    /// Returns whether the user's account is too new to use the service.
+    pub fn account_too_young(&self) -> bool {
+        self.plan_info
+            .as_ref()
+            .map(|plan| plan.is_account_too_young)
+            .unwrap_or_default()
     }
 
-    pub fn watch_current_user(&self) -> watch::Receiver<Option<Arc<User>>> {
-        self.current_user.clone()
+    /// Returns whether the current user has overdue invoices and usage should be blocked.
+    pub fn has_overdue_invoices(&self) -> bool {
+        self.plan_info
+            .as_ref()
+            .map(|plan| plan.has_overdue_invoices)
+            .unwrap_or_default()
     }
 
-    /// Returns whether the user's account is too new to use the service.
-    pub fn account_too_young(&self) -> bool {
-        self.account_too_young.unwrap_or(false)
+    pub fn is_usage_based_billing_enabled(&self) -> bool {
+        self.plan_info
+            .as_ref()
+            .map(|plan| plan.is_usage_based_billing_enabled)
+            .unwrap_or_default()
+    }
+
+    pub fn model_request_usage(&self) -> Option<ModelRequestUsage> {
+        self.model_request_usage
+    }
+
+    pub fn update_model_request_usage(&mut self, usage: ModelRequestUsage, cx: &mut Context<Self>) {
+        self.model_request_usage = Some(usage);
+        cx.notify();
+    }
+
+    pub fn edit_prediction_usage(&self) -> Option<EditPredictionUsage> {
+        self.edit_prediction_usage
     }
 
-    pub fn current_user_has_accepted_terms(&self) -> Option<bool> {
+    pub fn update_edit_prediction_usage(
+        &mut self,
+        usage: EditPredictionUsage,
+        cx: &mut Context<Self>,
+    ) {
+        self.edit_prediction_usage = Some(usage);
+        cx.notify();
+    }
+
+    fn update_authenticated_user(
+        &mut self,
+        response: GetAuthenticatedUserResponse,
+        cx: &mut Context<Self>,
+    ) {
+        let staff = response.user.is_staff && !*feature_flags::ZED_DISABLE_STAFF;
+        cx.update_flags(staff, response.feature_flags);
+        if let Some(client) = self.client.upgrade() {
+            client
+                .telemetry
+                .set_authenticated_user_info(Some(response.user.metrics_id.clone()), staff);
+        }
+
+        let accepted_tos_at = {
+            #[cfg(debug_assertions)]
+            if std::env::var("ZED_IGNORE_ACCEPTED_TOS").is_ok() {
+                None
+            } else {
+                response.user.accepted_tos_at
+            }
+
+            #[cfg(not(debug_assertions))]
+            response.user.accepted_tos_at
+        };
+
+        self.accepted_tos_at = Some(accepted_tos_at);
+        self.model_request_usage = Some(ModelRequestUsage(RequestUsage {
+            limit: response.plan.usage.model_requests.limit,
+            amount: response.plan.usage.model_requests.used as i32,
+        }));
+        self.edit_prediction_usage = Some(EditPredictionUsage(RequestUsage {
+            limit: response.plan.usage.edit_predictions.limit,
+            amount: response.plan.usage.edit_predictions.used as i32,
+        }));
+        self.plan_info = Some(response.plan);
+        cx.emit(Event::PrivateUserInfoUpdated);
+    }
+
+    pub fn watch_current_user(&self) -> watch::Receiver<Option<Arc<User>>> {
+        self.current_user.clone()
+    }
+
+    pub fn has_accepted_terms_of_service(&self) -> bool {
         self.accepted_tos_at
-            .map(|accepted_tos_at| accepted_tos_at.is_some())
+            .map_or(false, |accepted_tos_at| accepted_tos_at.is_some())
     }
 
     pub fn accept_terms_of_service(&self, cx: &Context<Self>) -> Task<Result<()>> {
@@ -766,23 +831,18 @@ impl UserStore {
         cx.spawn(async move |this, cx| -> anyhow::Result<()> {
             let client = client.upgrade().context("client not found")?;
             let response = client
-                .request(proto::AcceptTermsOfService {})
+                .cloud_client()
+                .accept_terms_of_service()
                 .await
                 .context("error accepting tos")?;
             this.update(cx, |this, cx| {
-                this.set_current_user_accepted_tos_at(Some(response.accepted_tos_at));
+                this.accepted_tos_at = Some(response.user.accepted_tos_at);
                 cx.emit(Event::PrivateUserInfoUpdated);
             })?;
             Ok(())
         })
     }
 
-    fn set_current_user_accepted_tos_at(&mut self, accepted_tos_at: Option<u64>) {
-        self.accepted_tos_at = Some(
-            accepted_tos_at.and_then(|timestamp| DateTime::from_timestamp(timestamp as i64, 0)),
-        );
-    }
-
     fn load_users(
         &self,
         request: impl RequestMessage<Response = UsersResponse>,

crates/collab/src/tests/integration_tests.rs 🔗

@@ -1286,7 +1286,7 @@ async fn test_calls_on_multiple_connections(
     client_b1.disconnect(&cx_b1.to_async());
     executor.advance_clock(RECEIVE_TIMEOUT);
     client_b1
-        .authenticate_and_connect(false, &cx_b1.to_async())
+        .connect(false, &cx_b1.to_async())
         .await
         .into_response()
         .unwrap();
@@ -1667,7 +1667,7 @@ async fn test_project_reconnect(
     // Client A reconnects. Their project is re-shared, and client B re-joins it.
     server.allow_connections();
     client_a
-        .authenticate_and_connect(false, &cx_a.to_async())
+        .connect(false, &cx_a.to_async())
         .await
         .into_response()
         .unwrap();
@@ -1796,7 +1796,7 @@ async fn test_project_reconnect(
     // Client B reconnects. They re-join the room and the remaining shared project.
     server.allow_connections();
     client_b
-        .authenticate_and_connect(false, &cx_b.to_async())
+        .connect(false, &cx_b.to_async())
         .await
         .into_response()
         .unwrap();
@@ -5738,7 +5738,7 @@ async fn test_contacts(
 
     server.allow_connections();
     client_c
-        .authenticate_and_connect(false, &cx_c.to_async())
+        .connect(false, &cx_c.to_async())
         .await
         .into_response()
         .unwrap();
@@ -6269,7 +6269,7 @@ async fn test_contact_requests(
         client.disconnect(&cx.to_async());
         client.clear_contacts(cx).await;
         client
-            .authenticate_and_connect(false, &cx.to_async())
+            .connect(false, &cx.to_async())
             .await
             .into_response()
             .unwrap();

crates/collab/src/tests/notification_tests.rs 🔗

@@ -3,6 +3,7 @@ use std::sync::Arc;
 use gpui::{BackgroundExecutor, TestAppContext};
 use notifications::NotificationEvent;
 use parking_lot::Mutex;
+use pretty_assertions::assert_eq;
 use rpc::{Notification, proto};
 
 use crate::tests::TestServer;
@@ -17,6 +18,9 @@ async fn test_notifications(
     let client_a = server.create_client(cx_a, "user_a").await;
     let client_b = server.create_client(cx_b, "user_b").await;
 
+    // Wait for authentication/connection to Collab to be established.
+    executor.run_until_parked();
+
     let notification_events_a = Arc::new(Mutex::new(Vec::new()));
     let notification_events_b = Arc::new(Mutex::new(Vec::new()));
     client_a.notification_store().update(cx_a, |_, cx| {

crates/collab/src/tests/test_server.rs 🔗

@@ -8,7 +8,7 @@ use crate::{
 use anyhow::anyhow;
 use call::ActiveCall;
 use channel::{ChannelBuffer, ChannelStore};
-use client::CloudUserStore;
+use client::test::{make_get_authenticated_user_response, parse_authorization_header};
 use client::{
     self, ChannelId, Client, Connection, Credentials, EstablishConnectionError, UserStore,
     proto::PeerId,
@@ -21,7 +21,7 @@ use fs::FakeFs;
 use futures::{StreamExt as _, channel::oneshot};
 use git::GitHostingProviderRegistry;
 use gpui::{AppContext as _, BackgroundExecutor, Entity, Task, TestAppContext, VisualTestContext};
-use http_client::FakeHttpClient;
+use http_client::{FakeHttpClient, Method};
 use language::LanguageRegistry;
 use node_runtime::NodeRuntime;
 use notifications::NotificationStore;
@@ -162,6 +162,8 @@ impl TestServer {
     }
 
     pub async fn create_client(&mut self, cx: &mut TestAppContext, name: &str) -> TestClient {
+        const ACCESS_TOKEN: &str = "the-token";
+
         let fs = FakeFs::new(cx.executor());
 
         cx.update(|cx| {
@@ -176,7 +178,7 @@ impl TestServer {
         });
 
         let clock = Arc::new(FakeSystemClock::new());
-        let http = FakeHttpClient::with_404_response();
+
         let user_id = if let Ok(Some(user)) = self.app_state.db.get_user_by_github_login(name).await
         {
             user.id
@@ -198,6 +200,47 @@ impl TestServer {
                 .expect("creating user failed")
                 .user_id
         };
+
+        let http = FakeHttpClient::create({
+            let name = name.to_string();
+            move |req| {
+                let name = name.clone();
+                async move {
+                    match (req.method(), req.uri().path()) {
+                        (&Method::GET, "/client/users/me") => {
+                            let credentials = parse_authorization_header(&req);
+                            if credentials
+                                != Some(Credentials {
+                                    user_id: user_id.to_proto(),
+                                    access_token: ACCESS_TOKEN.into(),
+                                })
+                            {
+                                return Ok(http_client::Response::builder()
+                                    .status(401)
+                                    .body("Unauthorized".into())
+                                    .unwrap());
+                            }
+
+                            Ok(http_client::Response::builder()
+                                .status(200)
+                                .body(
+                                    serde_json::to_string(&make_get_authenticated_user_response(
+                                        user_id.0, name,
+                                    ))
+                                    .unwrap()
+                                    .into(),
+                                )
+                                .unwrap())
+                        }
+                        _ => Ok(http_client::Response::builder()
+                            .status(404)
+                            .body("Not Found".into())
+                            .unwrap()),
+                    }
+                }
+            }
+        });
+
         let client_name = name.to_string();
         let mut client = cx.update(|cx| Client::new(clock, http.clone(), cx));
         let server = self.server.clone();
@@ -209,11 +252,10 @@ impl TestServer {
             .unwrap()
             .set_id(user_id.to_proto())
             .override_authenticate(move |cx| {
-                let access_token = "the-token".to_string();
                 cx.spawn(async move |_| {
                     Ok(Credentials {
                         user_id: user_id.to_proto(),
-                        access_token,
+                        access_token: ACCESS_TOKEN.into(),
                     })
                 })
             })
@@ -222,7 +264,7 @@ impl TestServer {
                     credentials,
                     &Credentials {
                         user_id: user_id.0 as u64,
-                        access_token: "the-token".into()
+                        access_token: ACCESS_TOKEN.into(),
                     }
                 );
 
@@ -282,15 +324,12 @@ 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(), 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));
         let app_state = Arc::new(workspace::AppState {
             client: client.clone(),
             user_store: user_store.clone(),
-            cloud_user_store,
             workspace_store,
             languages: language_registry,
             fs: fs.clone(),
@@ -323,7 +362,7 @@ impl TestServer {
         });
 
         client
-            .authenticate_and_connect(false, &cx.to_async())
+            .connect(false, &cx.to_async())
             .await
             .into_response()
             .unwrap();

crates/collab_ui/src/collab_panel.rs 🔗

@@ -2331,7 +2331,7 @@ impl CollabPanel {
                                 let client = this.client.clone();
                                 cx.spawn_in(window, async move |_, cx| {
                                     client
-                                        .authenticate_and_connect(true, &cx)
+                                        .connect(true, &cx)
                                         .await
                                         .into_response()
                                         .notify_async_err(cx);

crates/collab_ui/src/notification_panel.rs 🔗

@@ -634,13 +634,13 @@ impl Render for NotificationPanel {
                     .child(Icon::new(IconName::Envelope)),
             )
             .map(|this| {
-                if self.client.user_id().is_none() {
+                if !self.client.status().borrow().is_connected() {
                     this.child(
                         v_flex()
                             .gap_2()
                             .p_4()
                             .child(
-                                Button::new("sign_in_prompt_button", "Sign in")
+                                Button::new("connect_prompt_button", "Connect")
                                     .icon_color(Color::Muted)
                                     .icon(IconName::Github)
                                     .icon_position(IconPosition::Start)
@@ -652,10 +652,7 @@ impl Render for NotificationPanel {
                                             let client = client.clone();
                                             window
                                                 .spawn(cx, async move |cx| {
-                                                    match client
-                                                        .authenticate_and_connect(true, &cx)
-                                                        .await
-                                                    {
+                                                    match client.connect(true, &cx).await {
                                                         util::ConnectionResult::Timeout => {
                                                             log::error!("Connection timeout");
                                                         }
@@ -673,7 +670,7 @@ impl Render for NotificationPanel {
                             )
                             .child(
                                 div().flex().w_full().items_center().child(
-                                    Label::new("Sign in to view notifications.")
+                                    Label::new("Connect to view notifications.")
                                         .color(Color::Muted)
                                         .size(LabelSize::Small),
                                 ),

crates/eval/src/eval.rs 🔗

@@ -13,7 +13,7 @@ pub(crate) use tool_metrics::*;
 
 use ::fs::RealFs;
 use clap::Parser;
-use client::{Client, CloudUserStore, ProxySettings, UserStore};
+use client::{Client, ProxySettings, UserStore};
 use collections::{HashMap, HashSet};
 use extension::ExtensionHostProxy;
 use futures::future;
@@ -329,7 +329,6 @@ pub struct AgentAppState {
     pub languages: Arc<LanguageRegistry>,
     pub client: Arc<Client>,
     pub user_store: Entity<UserStore>,
-    pub cloud_user_store: Entity<CloudUserStore>,
     pub fs: Arc<dyn fs::Fs>,
     pub node_runtime: NodeRuntime,
 
@@ -384,8 +383,6 @@ pub fn init(cx: &mut App) -> Arc<AgentAppState> {
     let languages = Arc::new(languages);
 
     let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
-    let cloud_user_store =
-        cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store.clone(), cx));
 
     extension::init(cx);
 
@@ -425,12 +422,7 @@ pub fn init(cx: &mut App) -> Arc<AgentAppState> {
         languages.clone(),
     );
     language_model::init(client.clone(), cx);
-    language_models::init(
-        user_store.clone(),
-        cloud_user_store.clone(),
-        client.clone(),
-        cx,
-    );
+    language_models::init(user_store.clone(), client.clone(), cx);
     languages::init(languages.clone(), node_runtime.clone(), cx);
     prompt_store::init(cx);
     terminal_view::init(cx);
@@ -455,7 +447,6 @@ pub fn init(cx: &mut App) -> Arc<AgentAppState> {
         languages,
         client,
         user_store,
-        cloud_user_store,
         fs,
         node_runtime,
         prompt_builder,

crates/eval/src/instance.rs 🔗

@@ -221,7 +221,6 @@ impl ExampleInstance {
         let prompt_store = None;
         let thread_store = ThreadStore::load(
             project.clone(),
-            app_state.cloud_user_store.clone(),
             tools,
             prompt_store,
             app_state.prompt_builder.clone(),

crates/http_client/Cargo.toml 🔗

@@ -23,6 +23,7 @@ futures.workspace = true
 http.workspace = true
 http-body.workspace = true
 log.workspace = true
+parking_lot.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 url.workspace = true

crates/http_client/src/http_client.rs 🔗

@@ -9,12 +9,10 @@ pub use http::{self, Method, Request, Response, StatusCode, Uri};
 
 use futures::future::BoxFuture;
 use http::request::Builder;
+use parking_lot::Mutex;
 #[cfg(feature = "test-support")]
 use std::fmt;
-use std::{
-    any::type_name,
-    sync::{Arc, Mutex},
-};
+use std::{any::type_name, sync::Arc};
 pub use url::Url;
 
 #[derive(Default, Debug, Clone, PartialEq, Eq, Hash)]
@@ -86,6 +84,11 @@ pub trait HttpClient: 'static + Send + Sync {
     }
 
     fn proxy(&self) -> Option<&Url>;
+
+    #[cfg(feature = "test-support")]
+    fn as_fake(&self) -> &FakeHttpClient {
+        panic!("called as_fake on {}", type_name::<Self>())
+    }
 }
 
 /// An [`HttpClient`] that may have a proxy.
@@ -132,6 +135,11 @@ impl HttpClient for HttpClientWithProxy {
     fn type_name(&self) -> &'static str {
         self.client.type_name()
     }
+
+    #[cfg(feature = "test-support")]
+    fn as_fake(&self) -> &FakeHttpClient {
+        self.client.as_fake()
+    }
 }
 
 impl HttpClient for Arc<HttpClientWithProxy> {
@@ -153,6 +161,11 @@ impl HttpClient for Arc<HttpClientWithProxy> {
     fn type_name(&self) -> &'static str {
         self.client.type_name()
     }
+
+    #[cfg(feature = "test-support")]
+    fn as_fake(&self) -> &FakeHttpClient {
+        self.client.as_fake()
+    }
 }
 
 /// An [`HttpClient`] that has a base URL.
@@ -199,20 +212,13 @@ impl HttpClientWithUrl {
 
     /// Returns the base URL.
     pub fn base_url(&self) -> String {
-        self.base_url
-            .lock()
-            .map_or_else(|_| Default::default(), |url| url.clone())
+        self.base_url.lock().clone()
     }
 
     /// Sets the base URL.
     pub fn set_base_url(&self, base_url: impl Into<String>) {
         let base_url = base_url.into();
-        self.base_url
-            .lock()
-            .map(|mut url| {
-                *url = base_url;
-            })
-            .ok();
+        *self.base_url.lock() = base_url;
     }
 
     /// Builds a URL using the given path.
@@ -288,6 +294,11 @@ impl HttpClient for Arc<HttpClientWithUrl> {
     fn type_name(&self) -> &'static str {
         self.client.type_name()
     }
+
+    #[cfg(feature = "test-support")]
+    fn as_fake(&self) -> &FakeHttpClient {
+        self.client.as_fake()
+    }
 }
 
 impl HttpClient for HttpClientWithUrl {
@@ -309,6 +320,11 @@ impl HttpClient for HttpClientWithUrl {
     fn type_name(&self) -> &'static str {
         self.client.type_name()
     }
+
+    #[cfg(feature = "test-support")]
+    fn as_fake(&self) -> &FakeHttpClient {
+        self.client.as_fake()
+    }
 }
 
 pub fn read_proxy_from_env() -> Option<Url> {
@@ -360,10 +376,15 @@ impl HttpClient for BlockedHttpClient {
     fn type_name(&self) -> &'static str {
         type_name::<Self>()
     }
+
+    #[cfg(feature = "test-support")]
+    fn as_fake(&self) -> &FakeHttpClient {
+        panic!("called as_fake on {}", type_name::<Self>())
+    }
 }
 
 #[cfg(feature = "test-support")]
-type FakeHttpHandler = Box<
+type FakeHttpHandler = Arc<
     dyn Fn(Request<AsyncBody>) -> BoxFuture<'static, anyhow::Result<Response<AsyncBody>>>
         + Send
         + Sync
@@ -372,7 +393,7 @@ type FakeHttpHandler = Box<
 
 #[cfg(feature = "test-support")]
 pub struct FakeHttpClient {
-    handler: FakeHttpHandler,
+    handler: Mutex<Option<FakeHttpHandler>>,
     user_agent: HeaderValue,
 }
 
@@ -387,7 +408,7 @@ impl FakeHttpClient {
             base_url: Mutex::new("http://test.example".into()),
             client: HttpClientWithProxy {
                 client: Arc::new(Self {
-                    handler: Box::new(move |req| Box::pin(handler(req))),
+                    handler: Mutex::new(Some(Arc::new(move |req| Box::pin(handler(req))))),
                     user_agent: HeaderValue::from_static(type_name::<Self>()),
                 }),
                 proxy: None,
@@ -412,6 +433,18 @@ impl FakeHttpClient {
                 .unwrap())
         })
     }
+
+    pub fn replace_handler<Fut, F>(&self, new_handler: F)
+    where
+        Fut: futures::Future<Output = anyhow::Result<Response<AsyncBody>>> + Send + 'static,
+        F: Fn(FakeHttpHandler, Request<AsyncBody>) -> Fut + Send + Sync + 'static,
+    {
+        let mut handler = self.handler.lock();
+        let old_handler = handler.take().unwrap();
+        *handler = Some(Arc::new(move |req| {
+            Box::pin(new_handler(old_handler.clone(), req))
+        }));
+    }
 }
 
 #[cfg(feature = "test-support")]
@@ -427,7 +460,7 @@ impl HttpClient for FakeHttpClient {
         &self,
         req: Request<AsyncBody>,
     ) -> BoxFuture<'static, anyhow::Result<Response<AsyncBody>>> {
-        let future = (self.handler)(req);
+        let future = (self.handler.lock().as_ref().unwrap())(req);
         future
     }
 
@@ -442,4 +475,8 @@ impl HttpClient for FakeHttpClient {
     fn type_name(&self) -> &'static str {
         type_name::<Self>()
     }
+
+    fn as_fake(&self) -> &FakeHttpClient {
+        self
+    }
 }

crates/inline_completion_button/src/inline_completion_button.rs 🔗

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use client::{CloudUserStore, DisableAiSettings, zed_urls};
+use client::{DisableAiSettings, UserStore, zed_urls};
 use cloud_llm_client::UsageLimit;
 use copilot::{Copilot, Status};
 use editor::{
@@ -59,7 +59,7 @@ pub struct InlineCompletionButton {
     file: Option<Arc<dyn File>>,
     edit_prediction_provider: Option<Arc<dyn inline_completion::InlineCompletionProviderHandle>>,
     fs: Arc<dyn Fs>,
-    cloud_user_store: Entity<CloudUserStore>,
+    user_store: Entity<UserStore>,
     popover_menu_handle: PopoverMenuHandle<ContextMenu>,
 }
 
@@ -245,9 +245,9 @@ impl Render for InlineCompletionButton {
                     IconName::ZedPredictDisabled
                 };
 
-                if zeta::should_show_upsell_modal(&self.cloud_user_store, cx) {
-                    let tooltip_meta = if self.cloud_user_store.read(cx).is_authenticated() {
-                        if self.cloud_user_store.read(cx).has_accepted_tos() {
+                if zeta::should_show_upsell_modal(&self.user_store, cx) {
+                    let tooltip_meta = if self.user_store.read(cx).current_user().is_some() {
+                        if self.user_store.read(cx).has_accepted_terms_of_service() {
                             "Choose a Plan"
                         } else {
                             "Accept the Terms of Service"
@@ -371,7 +371,7 @@ impl Render for InlineCompletionButton {
 impl InlineCompletionButton {
     pub fn new(
         fs: Arc<dyn Fs>,
-        cloud_user_store: Entity<CloudUserStore>,
+        user_store: Entity<UserStore>,
         popover_menu_handle: PopoverMenuHandle<ContextMenu>,
         cx: &mut Context<Self>,
     ) -> Self {
@@ -390,9 +390,9 @@ impl InlineCompletionButton {
             language: None,
             file: None,
             edit_prediction_provider: None,
+            user_store,
             popover_menu_handle,
             fs,
-            cloud_user_store,
         }
     }
 
@@ -763,7 +763,7 @@ impl InlineCompletionButton {
                         })
                     })
                     .separator();
-            } else if self.cloud_user_store.read(cx).account_too_young() {
+            } else if self.user_store.read(cx).account_too_young() {
                 menu = menu
                     .custom_entry(
                         |_window, _cx| {
@@ -778,7 +778,7 @@ impl InlineCompletionButton {
                         cx.open_url(&zed_urls::account_url(cx))
                     })
                     .separator();
-            } else if self.cloud_user_store.read(cx).has_overdue_invoices() {
+            } else if self.user_store.read(cx).has_overdue_invoices() {
                 menu = menu
                     .custom_entry(
                         |_window, _cx| {

crates/language_model/src/model/cloud_model.rs 🔗

@@ -3,10 +3,11 @@ use std::sync::Arc;
 
 use anyhow::Result;
 use client::Client;
+use cloud_llm_client::Plan;
 use gpui::{
     App, AppContext as _, AsyncApp, Context, Entity, EventEmitter, Global, ReadGlobal as _,
 };
-use proto::{Plan, TypedEnvelope};
+use proto::TypedEnvelope;
 use smol::lock::{RwLock, RwLockUpgradableReadGuard, RwLockWriteGuard};
 use thiserror::Error;
 
@@ -30,7 +31,7 @@ pub struct ModelRequestLimitReachedError {
 impl fmt::Display for ModelRequestLimitReachedError {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         let message = match self.plan {
-            Plan::Free => "Model request limit reached. Upgrade to Zed Pro for more requests.",
+            Plan::ZedFree => "Model request limit reached. Upgrade to Zed Pro for more requests.",
             Plan::ZedPro => {
                 "Model request limit reached. Upgrade to usage-based billing for more requests."
             }

crates/language_models/Cargo.toml 🔗

@@ -44,7 +44,6 @@ ollama = { workspace = true, features = ["schemars"] }
 open_ai = { workspace = true, features = ["schemars"] }
 open_router = { workspace = true, features = ["schemars"] }
 partial-json-fixer.workspace = true
-proto.workspace = true
 release_channel.workspace = true
 schemars.workspace = true
 serde.workspace = true

crates/language_models/src/language_models.rs 🔗

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use ::settings::{Settings, SettingsStore};
-use client::{Client, CloudUserStore, UserStore};
+use client::{Client, UserStore};
 use collections::HashSet;
 use gpui::{App, Context, Entity};
 use language_model::{LanguageModelProviderId, LanguageModelRegistry};
@@ -26,22 +26,11 @@ use crate::provider::vercel::VercelLanguageModelProvider;
 use crate::provider::x_ai::XAiLanguageModelProvider;
 pub use crate::settings::*;
 
-pub fn init(
-    user_store: Entity<UserStore>,
-    cloud_user_store: Entity<CloudUserStore>,
-    client: Arc<Client>,
-    cx: &mut App,
-) {
+pub fn init(user_store: Entity<UserStore>, client: Arc<Client>, cx: &mut App) {
     crate::settings::init_settings(cx);
     let registry = LanguageModelRegistry::global(cx);
     registry.update(cx, |registry, cx| {
-        register_language_model_providers(
-            registry,
-            user_store,
-            cloud_user_store,
-            client.clone(),
-            cx,
-        );
+        register_language_model_providers(registry, user_store, client.clone(), cx);
     });
 
     let mut openai_compatible_providers = AllLanguageModelSettings::get_global(cx)
@@ -111,17 +100,11 @@ fn register_openai_compatible_providers(
 fn register_language_model_providers(
     registry: &mut LanguageModelRegistry,
     user_store: Entity<UserStore>,
-    cloud_user_store: Entity<CloudUserStore>,
     client: Arc<Client>,
     cx: &mut Context<LanguageModelRegistry>,
 ) {
     registry.register_provider(
-        CloudLanguageModelProvider::new(
-            user_store.clone(),
-            cloud_user_store.clone(),
-            client.clone(),
-            cx,
-        ),
+        CloudLanguageModelProvider::new(user_store.clone(), client.clone(), cx),
         cx,
     );
 

crates/language_models/src/provider/cloud.rs 🔗

@@ -2,7 +2,7 @@ use ai_onboarding::YoungAccountBanner;
 use anthropic::AnthropicModelMode;
 use anyhow::{Context as _, Result, anyhow};
 use chrono::{DateTime, Utc};
-use client::{Client, CloudUserStore, ModelRequestUsage, UserStore, zed_urls};
+use client::{Client, ModelRequestUsage, UserStore, zed_urls};
 use cloud_llm_client::{
     CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, CURRENT_PLAN_HEADER_NAME, CompletionBody,
     CompletionEvent, CompletionRequestStatus, CountTokensBody, CountTokensResponse,
@@ -117,7 +117,6 @@ pub struct State {
     client: Arc<Client>,
     llm_api_token: LlmApiToken,
     user_store: Entity<UserStore>,
-    cloud_user_store: Entity<CloudUserStore>,
     status: client::Status,
     accept_terms_of_service_task: Option<Task<Result<()>>>,
     models: Vec<Arc<cloud_llm_client::LanguageModel>>,
@@ -133,17 +132,14 @@ impl State {
     fn new(
         client: Arc<Client>,
         user_store: Entity<UserStore>,
-        cloud_user_store: Entity<CloudUserStore>,
         status: client::Status,
         cx: &mut Context<Self>,
     ) -> Self {
         let refresh_llm_token_listener = RefreshLlmTokenListener::global(cx);
-
         Self {
             client: client.clone(),
             llm_api_token: LlmApiToken::default(),
-            user_store,
-            cloud_user_store,
+            user_store: user_store.clone(),
             status,
             accept_terms_of_service_task: None,
             models: Vec::new(),
@@ -152,18 +148,12 @@ impl State {
             recommended_models: Vec::new(),
             _fetch_models_task: cx.spawn(async move |this, cx| {
                 maybe!(async move {
-                    let (client, cloud_user_store, llm_api_token) =
-                        this.read_with(cx, |this, _cx| {
-                            (
-                                client.clone(),
-                                this.cloud_user_store.clone(),
-                                this.llm_api_token.clone(),
-                            )
-                        })?;
+                    let (client, llm_api_token) = this
+                        .read_with(cx, |this, _cx| (client.clone(), this.llm_api_token.clone()))?;
 
                     loop {
-                        let is_authenticated =
-                            cloud_user_store.read_with(cx, |this, _cx| this.is_authenticated())?;
+                        let is_authenticated = user_store
+                            .read_with(cx, |user_store, _cx| user_store.current_user().is_some())?;
                         if is_authenticated {
                             break;
                         }
@@ -204,22 +194,19 @@ impl State {
     }
 
     fn is_signed_out(&self, cx: &App) -> bool {
-        !self.cloud_user_store.read(cx).is_authenticated()
+        self.user_store.read(cx).current_user().is_none()
     }
 
     fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
         let client = self.client.clone();
         cx.spawn(async move |state, cx| {
-            client
-                .authenticate_and_connect(true, &cx)
-                .await
-                .into_response()?;
+            client.sign_in_with_optional_connect(true, &cx).await?;
             state.update(cx, |_, cx| cx.notify())
         })
     }
 
     fn has_accepted_terms_of_service(&self, cx: &App) -> bool {
-        self.cloud_user_store.read(cx).has_accepted_tos()
+        self.user_store.read(cx).has_accepted_terms_of_service()
     }
 
     fn accept_terms_of_service(&mut self, cx: &mut Context<Self>) {
@@ -303,24 +290,11 @@ impl State {
 }
 
 impl CloudLanguageModelProvider {
-    pub fn new(
-        user_store: Entity<UserStore>,
-        cloud_user_store: Entity<CloudUserStore>,
-        client: Arc<Client>,
-        cx: &mut App,
-    ) -> Self {
+    pub fn new(user_store: Entity<UserStore>, client: Arc<Client>, cx: &mut App) -> Self {
         let mut status_rx = client.status();
         let status = *status_rx.borrow();
 
-        let state = cx.new(|cx| {
-            State::new(
-                client.clone(),
-                user_store.clone(),
-                cloud_user_store.clone(),
-                status,
-                cx,
-            )
-        });
+        let state = cx.new(|cx| State::new(client.clone(), user_store.clone(), status, cx));
 
         let state_ref = state.downgrade();
         let maintain_client_status = cx.spawn(async move |cx| {
@@ -632,11 +606,6 @@ impl CloudLanguageModel {
                         .and_then(|plan| plan.to_str().ok())
                         .and_then(|plan| cloud_llm_client::Plan::from_str(plan).ok())
                     {
-                        let plan = match plan {
-                            cloud_llm_client::Plan::ZedFree => proto::Plan::Free,
-                            cloud_llm_client::Plan::ZedPro => proto::Plan::ZedPro,
-                            cloud_llm_client::Plan::ZedProTrial => proto::Plan::ZedProTrial,
-                        };
                         return Err(anyhow!(ModelRequestLimitReachedError { plan }));
                     }
                 }
@@ -1281,15 +1250,15 @@ impl ConfigurationView {
 impl Render for ConfigurationView {
     fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
         let state = self.state.read(cx);
-        let cloud_user_store = state.cloud_user_store.read(cx);
+        let user_store = state.user_store.read(cx);
 
         ZedAiConfiguration {
             is_connected: !state.is_signed_out(cx),
-            plan: cloud_user_store.plan(),
-            subscription_period: cloud_user_store.subscription_period(),
-            eligible_for_trial: cloud_user_store.trial_started_at().is_none(),
+            plan: user_store.plan(),
+            subscription_period: user_store.subscription_period(),
+            eligible_for_trial: user_store.trial_started_at().is_none(),
             has_accepted_terms_of_service: state.has_accepted_terms_of_service(cx),
-            account_too_young: cloud_user_store.account_too_young(),
+            account_too_young: user_store.account_too_young(),
             accept_terms_of_service_in_progress: state.accept_terms_of_service_task.is_some(),
             accept_terms_of_service_callback: self.accept_terms_of_service_callback.clone(),
             sign_in_callback: self.sign_in_callback.clone(),

crates/onboarding/src/ai_setup_page.rs 🔗

@@ -278,7 +278,7 @@ pub(crate) fn render_ai_setup_page(
                 .child(AiUpsellCard {
                     sign_in_status: SignInStatus::SignedIn,
                     sign_in: Arc::new(|_, _| {}),
-                    user_plan: onboarding.cloud_user_store.read(cx).plan(),
+                    user_plan: onboarding.user_store.read(cx).plan(),
                 })
                 .child(render_llm_provider_section(
                     onboarding,

crates/onboarding/src/onboarding.rs 🔗

@@ -1,5 +1,5 @@
 use crate::welcome::{ShowWelcome, WelcomePage};
-use client::{Client, CloudUserStore, UserStore};
+use client::{Client, UserStore};
 use command_palette_hooks::CommandPaletteFilter;
 use db::kvp::KEY_VALUE_STORE;
 use feature_flags::{FeatureFlag, FeatureFlagViewExt as _};
@@ -220,7 +220,6 @@ struct Onboarding {
     workspace: WeakEntity<Workspace>,
     focus_handle: FocusHandle,
     selected_page: SelectedPage,
-    cloud_user_store: Entity<CloudUserStore>,
     user_store: Entity<UserStore>,
     _settings_subscription: Subscription,
 }
@@ -231,7 +230,6 @@ impl Onboarding {
             workspace: workspace.weak_handle(),
             focus_handle: cx.focus_handle(),
             selected_page: SelectedPage::Basics,
-            cloud_user_store: workspace.app_state().cloud_user_store.clone(),
             user_store: workspace.user_store().clone(),
             _settings_subscription: cx.observe_global::<SettingsStore>(move |_, cx| cx.notify()),
         })
@@ -365,9 +363,8 @@ impl Onboarding {
                             window
                                 .spawn(cx, async move |cx| {
                                     client
-                                        .authenticate_and_connect(true, &cx)
+                                        .sign_in_with_optional_connect(true, &cx)
                                         .await
-                                        .into_response()
                                         .notify_async_err(cx);
                                 })
                                 .detach();

crates/project/src/project.rs 🔗

@@ -1362,10 +1362,7 @@ impl Project {
         fs: Arc<dyn Fs>,
         cx: AsyncApp,
     ) -> Result<Entity<Self>> {
-        client
-            .authenticate_and_connect(true, &cx)
-            .await
-            .into_response()?;
+        client.connect(true, &cx).await.into_response()?;
 
         let subscriptions = [
             EntitySubscription::Project(client.subscribe_to_entity::<Self>(remote_id)?),

crates/title_bar/src/title_bar.rs 🔗

@@ -20,7 +20,7 @@ use crate::application_menu::{
 
 use auto_update::AutoUpdateStatus;
 use call::ActiveCall;
-use client::{Client, CloudUserStore, UserStore, zed_urls};
+use client::{Client, UserStore, zed_urls};
 use cloud_llm_client::Plan;
 use gpui::{
     Action, AnyElement, App, Context, Corner, Element, Entity, Focusable, InteractiveElement,
@@ -126,7 +126,6 @@ pub struct TitleBar {
     platform_titlebar: Entity<PlatformTitleBar>,
     project: Entity<Project>,
     user_store: Entity<UserStore>,
-    cloud_user_store: Entity<CloudUserStore>,
     client: Arc<Client>,
     workspace: WeakEntity<Workspace>,
     application_menu: Option<Entity<ApplicationMenu>>,
@@ -180,11 +179,9 @@ impl Render for TitleBar {
             children.push(self.banner.clone().into_any_element())
         }
 
-        let is_authenticated = self.cloud_user_store.read(cx).is_authenticated();
         let status = self.client.status();
         let status = &*status.borrow();
-
-        let show_sign_in = !is_authenticated || !matches!(status, client::Status::Connected { .. });
+        let user = self.user_store.read(cx).current_user();
 
         children.push(
             h_flex()
@@ -194,10 +191,10 @@ impl Render for TitleBar {
                 .children(self.render_call_controls(window, cx))
                 .children(self.render_connection_status(status, cx))
                 .when(
-                    show_sign_in && TitleBarSettings::get_global(cx).show_sign_in,
+                    user.is_none() && TitleBarSettings::get_global(cx).show_sign_in,
                     |el| el.child(self.render_sign_in_button(cx)),
                 )
-                .when(is_authenticated, |parent| {
+                .when(user.is_some(), |parent| {
                     parent.child(self.render_user_menu_button(cx))
                 })
                 .into_any_element(),
@@ -248,7 +245,6 @@ impl TitleBar {
     ) -> Self {
         let project = workspace.project().clone();
         let user_store = workspace.app_state().user_store.clone();
-        let cloud_user_store = workspace.app_state().cloud_user_store.clone();
         let client = workspace.app_state().client.clone();
         let active_call = ActiveCall::global(cx);
 
@@ -296,7 +292,6 @@ impl TitleBar {
             workspace: workspace.weak_handle(),
             project,
             user_store,
-            cloud_user_store,
             client,
             _subscriptions: subscriptions,
             banner,
@@ -622,9 +617,8 @@ impl TitleBar {
                 window
                     .spawn(cx, async move |cx| {
                         client
-                            .authenticate_and_connect(true, &cx)
+                            .sign_in_with_optional_connect(true, &cx)
                             .await
-                            .into_response()
                             .notify_async_err(cx);
                     })
                     .detach();
@@ -632,15 +626,15 @@ 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 = cloud_user_store.subscription_period().is_some();
-            let plan = cloud_user_store.plan().filter(|_| {
+        let user_store = self.user_store.read(cx);
+        if let Some(user) = user_store.current_user() {
+            let has_subscription_period = user_store.subscription_period().is_some();
+            let plan = 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
             });
 
-            let user_avatar = user.avatar_url.clone();
+            let user_avatar = user.avatar_uri.clone();
             let free_chip_bg = cx
                 .theme()
                 .colors()

crates/workspace/src/workspace.rs 🔗

@@ -15,7 +15,6 @@ mod toast_layer;
 mod toolbar;
 mod workspace_settings;
 
-use client::CloudUserStore;
 pub use toast_layer::{ToastAction, ToastLayer, ToastView};
 
 use anyhow::{Context as _, Result, anyhow};
@@ -840,7 +839,6 @@ pub struct AppState {
     pub languages: Arc<LanguageRegistry>,
     pub client: Arc<Client>,
     pub user_store: Entity<UserStore>,
-    pub cloud_user_store: Entity<CloudUserStore>,
     pub workspace_store: Entity<WorkspaceStore>,
     pub fs: Arc<dyn fs::Fs>,
     pub build_window_options: fn(Option<Uuid>, &mut App) -> WindowOptions,
@@ -913,8 +911,6 @@ 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(), user_store.clone(), cx));
         let workspace_store = cx.new(|cx| WorkspaceStore::new(client.clone(), cx));
 
         theme::init(theme::LoadThemes::JustBase, cx);
@@ -926,7 +922,6 @@ impl AppState {
             fs,
             languages,
             user_store,
-            cloud_user_store,
             workspace_store,
             node_runtime: NodeRuntime::unavailable(),
             build_window_options: |_, _| Default::default(),
@@ -5739,16 +5734,12 @@ 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(), 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));
         window.activate_window();
         let app_state = Arc::new(AppState {
             languages: project.read(cx).languages().clone(),
             workspace_store,
-            cloud_user_store,
             client,
             user_store,
             fs: project.read(cx).fs().clone(),
@@ -6947,10 +6938,13 @@ async fn join_channel_internal(
         match status {
             Status::Connecting
             | Status::Authenticating
+            | Status::Authenticated
             | Status::Reconnecting
             | Status::Reauthenticating => continue,
             Status::Connected { .. } => break 'outer,
-            Status::SignedOut => return Err(ErrorCode::SignedOut.into()),
+            Status::SignedOut | Status::AuthenticationError => {
+                return Err(ErrorCode::SignedOut.into());
+            }
             Status::UpgradeRequired => return Err(ErrorCode::UpgradeRequired.into()),
             Status::ConnectionError | Status::ConnectionLost | Status::ReconnectionError { .. } => {
                 return Err(ErrorCode::Disconnected.into());

crates/zed/src/main.rs 🔗

@@ -5,7 +5,7 @@ use agent_ui::AgentPanel;
 use anyhow::{Context as _, Result};
 use clap::{Parser, command};
 use cli::FORCE_CLI_MODE_ENV_VAR_NAME;
-use client::{Client, CloudUserStore, ProxySettings, UserStore, parse_zed_link};
+use client::{Client, ProxySettings, UserStore, parse_zed_link};
 use collab_ui::channel_view::ChannelView;
 use collections::HashMap;
 use db::kvp::{GLOBAL_KEY_VALUE_STORE, KEY_VALUE_STORE};
@@ -42,7 +42,7 @@ use theme::{
     ActiveTheme, IconThemeNotFoundError, SystemAppearance, ThemeNotFoundError, ThemeRegistry,
     ThemeSettings,
 };
-use util::{ConnectionResult, ResultExt, TryFutureExt, maybe};
+use util::{ResultExt, TryFutureExt, maybe};
 use uuid::Uuid;
 use welcome::{FIRST_OPEN, show_welcome_view};
 use workspace::{
@@ -457,8 +457,6 @@ 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(), user_store.clone(), cx));
         let workspace_store = cx.new(|cx| WorkspaceStore::new(client.clone(), cx));
 
         language_extension::init(
@@ -518,7 +516,6 @@ pub fn main() {
             languages: languages.clone(),
             client: client.clone(),
             user_store: user_store.clone(),
-            cloud_user_store,
             fs: fs.clone(),
             build_window_options,
             workspace_store,
@@ -556,12 +553,7 @@ pub fn main() {
         );
         supermaven::init(app_state.client.clone(), cx);
         language_model::init(app_state.client.clone(), cx);
-        language_models::init(
-            app_state.user_store.clone(),
-            app_state.cloud_user_store.clone(),
-            app_state.client.clone(),
-            cx,
-        );
+        language_models::init(app_state.user_store.clone(), app_state.client.clone(), cx);
         agent_settings::init(cx);
         agent_servers::init(cx);
         web_search::init(cx);
@@ -569,7 +561,7 @@ pub fn main() {
         snippet_provider::init(cx);
         inline_completion_registry::init(
             app_state.client.clone(),
-            app_state.cloud_user_store.clone(),
+            app_state.user_store.clone(),
             cx,
         );
         let prompt_builder = PromptBuilder::load(app_state.fs.clone(), stdout_is_a_pty(), cx);
@@ -690,17 +682,9 @@ pub fn main() {
 
         cx.spawn({
             let client = app_state.client.clone();
-            async move |cx| match authenticate(client, &cx).await {
-                ConnectionResult::Timeout => log::error!("Timeout during initial auth"),
-                ConnectionResult::ConnectionReset => {
-                    log::error!("Connection reset during initial auth")
-                }
-                ConnectionResult::Result(r) => {
-                    r.log_err();
-                }
-            }
+            async move |cx| authenticate(client, &cx).await
         })
-        .detach();
+        .detach_and_log_err(cx);
 
         let urls: Vec<_> = args
             .paths_or_urls
@@ -850,15 +834,7 @@ fn handle_open_request(request: OpenRequest, app_state: Arc<AppState>, cx: &mut
                 let client = app_state.client.clone();
                 // we continue even if authentication fails as join_channel/ open channel notes will
                 // show a visible error message.
-                match authenticate(client, &cx).await {
-                    ConnectionResult::Timeout => {
-                        log::error!("Timeout during open request handling")
-                    }
-                    ConnectionResult::ConnectionReset => {
-                        log::error!("Connection reset during open request handling")
-                    }
-                    ConnectionResult::Result(r) => r?,
-                };
+                authenticate(client, &cx).await.log_err();
 
                 if let Some(channel_id) = request.join_channel {
                     cx.update(|cx| {
@@ -908,18 +884,18 @@ fn handle_open_request(request: OpenRequest, app_state: Arc<AppState>, cx: &mut
     }
 }
 
-async fn authenticate(client: Arc<Client>, cx: &AsyncApp) -> ConnectionResult<()> {
+async fn authenticate(client: Arc<Client>, cx: &AsyncApp) -> Result<()> {
     if stdout_is_a_pty() {
         if client::IMPERSONATE_LOGIN.is_some() {
-            return client.authenticate_and_connect(false, cx).await;
+            client.sign_in_with_optional_connect(false, cx).await?;
         } else if client.has_credentials(cx).await {
-            return client.authenticate_and_connect(true, cx).await;
+            client.sign_in_with_optional_connect(true, cx).await?;
         }
     } else if client.has_credentials(cx).await {
-        return client.authenticate_and_connect(true, cx).await;
+        client.sign_in_with_optional_connect(true, cx).await?;
     }
 
-    ConnectionResult::Result(Ok(()))
+    Ok(())
 }
 
 async fn system_id() -> Result<IdType> {

crates/zed/src/zed.rs 🔗

@@ -336,7 +336,7 @@ pub fn initialize_workspace(
         let edit_prediction_button = cx.new(|cx| {
             inline_completion_button::InlineCompletionButton::new(
                 app_state.fs.clone(),
-                app_state.cloud_user_store.clone(),
+                app_state.user_store.clone(),
                 inline_completion_menu_handle.clone(),
                 cx,
             )
@@ -4488,12 +4488,7 @@ mod tests {
             );
             image_viewer::init(cx);
             language_model::init(app_state.client.clone(), cx);
-            language_models::init(
-                app_state.user_store.clone(),
-                app_state.cloud_user_store.clone(),
-                app_state.client.clone(),
-                cx,
-            );
+            language_models::init(app_state.user_store.clone(), app_state.client.clone(), cx);
             web_search::init(cx);
             web_search_providers::init(app_state.client.clone(), cx);
             let prompt_builder = PromptBuilder::load(app_state.fs.clone(), false, cx);

crates/zed/src/zed/component_preview.rs 🔗

@@ -139,8 +139,7 @@ impl ComponentPreview {
         let project_clone = project.clone();
 
         cx.spawn_in(window, async move |entity, cx| {
-            let thread_store_future =
-                load_preview_thread_store(workspace_clone.clone(), project_clone.clone(), cx);
+            let thread_store_future = load_preview_thread_store(project_clone.clone(), cx);
             let text_thread_store_future =
                 load_preview_text_thread_store(workspace_clone.clone(), project_clone.clone(), cx);
 

crates/zed/src/zed/component_preview/preview_support/active_thread.rs 🔗

@@ -12,22 +12,19 @@ use ui::{App, Window};
 use workspace::Workspace;
 
 pub fn load_preview_thread_store(
-    workspace: WeakEntity<Workspace>,
     project: Entity<Project>,
     cx: &mut AsyncApp,
 ) -> Task<Result<Entity<ThreadStore>>> {
-    workspace
-        .update(cx, |workspace, cx| {
-            ThreadStore::load(
-                project.clone(),
-                workspace.app_state().cloud_user_store.clone(),
-                cx.new(|_| ToolWorkingSet::default()),
-                None,
-                Arc::new(PromptBuilder::new(None).unwrap()),
-                cx,
-            )
-        })
-        .unwrap_or(Task::ready(Err(anyhow!("workspace dropped"))))
+    cx.update(|cx| {
+        ThreadStore::load(
+            project.clone(),
+            cx.new(|_| ToolWorkingSet::default()),
+            None,
+            Arc::new(PromptBuilder::new(None).unwrap()),
+            cx,
+        )
+    })
+    .unwrap_or(Task::ready(Err(anyhow!("workspace dropped"))))
 }
 
 pub fn load_preview_text_thread_store(

crates/zed/src/zed/inline_completion_registry.rs 🔗

@@ -1,4 +1,4 @@
-use client::{Client, CloudUserStore, DisableAiSettings};
+use client::{Client, DisableAiSettings, UserStore};
 use collections::HashMap;
 use copilot::{Copilot, CopilotCompletionProvider};
 use editor::Editor;
@@ -13,12 +13,12 @@ use util::ResultExt;
 use workspace::Workspace;
 use zeta::{ProviderDataCollection, ZetaInlineCompletionProvider};
 
-pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &mut App) {
+pub fn init(client: Arc<Client>, user_store: Entity<UserStore>, cx: &mut App) {
     let editors: Rc<RefCell<HashMap<WeakEntity<Editor>, AnyWindowHandle>>> = Rc::default();
     cx.observe_new({
         let editors = editors.clone();
         let client = client.clone();
-        let cloud_user_store = cloud_user_store.clone();
+        let user_store = user_store.clone();
         move |editor: &mut Editor, window, cx: &mut Context<Editor>| {
             if !editor.mode().is_full() {
                 return;
@@ -48,7 +48,7 @@ pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &
                 editor,
                 provider,
                 &client,
-                cloud_user_store.clone(),
+                user_store.clone(),
                 window,
                 cx,
             );
@@ -60,7 +60,7 @@ pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &
 
     let mut provider = all_language_settings(None, cx).edit_predictions.provider;
     cx.spawn({
-        let cloud_user_store = cloud_user_store.clone();
+        let user_store = user_store.clone();
         let editors = editors.clone();
         let client = client.clone();
 
@@ -72,7 +72,7 @@ pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &
                         &editors,
                         provider,
                         &client,
-                        cloud_user_store.clone(),
+                        user_store.clone(),
                         cx,
                     );
                 })
@@ -85,12 +85,12 @@ pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &
     cx.observe_global::<SettingsStore>({
         let editors = editors.clone();
         let client = client.clone();
-        let cloud_user_store = cloud_user_store.clone();
+        let user_store = user_store.clone();
         move |cx| {
             let new_provider = all_language_settings(None, cx).edit_predictions.provider;
 
             if new_provider != provider {
-                let tos_accepted = cloud_user_store.read(cx).has_accepted_tos();
+                let tos_accepted = user_store.read(cx).has_accepted_terms_of_service();
 
                 telemetry::event!(
                     "Edit Prediction Provider Changed",
@@ -104,7 +104,7 @@ pub fn init(client: Arc<Client>, cloud_user_store: Entity<CloudUserStore>, cx: &
                     &editors,
                     provider,
                     &client,
-                    cloud_user_store.clone(),
+                    user_store.clone(),
                     cx,
                 );
 
@@ -145,7 +145,7 @@ fn assign_edit_prediction_providers(
     editors: &Rc<RefCell<HashMap<WeakEntity<Editor>, AnyWindowHandle>>>,
     provider: EditPredictionProvider,
     client: &Arc<Client>,
-    cloud_user_store: Entity<CloudUserStore>,
+    user_store: Entity<UserStore>,
     cx: &mut App,
 ) {
     for (editor, window) in editors.borrow().iter() {
@@ -155,7 +155,7 @@ fn assign_edit_prediction_providers(
                     editor,
                     provider,
                     &client,
-                    cloud_user_store.clone(),
+                    user_store.clone(),
                     window,
                     cx,
                 );
@@ -210,7 +210,7 @@ fn assign_edit_prediction_provider(
     editor: &mut Editor,
     provider: EditPredictionProvider,
     client: &Arc<Client>,
-    cloud_user_store: Entity<CloudUserStore>,
+    user_store: Entity<UserStore>,
     window: &mut Window,
     cx: &mut Context<Editor>,
 ) {
@@ -241,7 +241,7 @@ fn assign_edit_prediction_provider(
             }
         }
         EditPredictionProvider::Zed => {
-            if cloud_user_store.read(cx).is_authenticated() {
+            if user_store.read(cx).current_user().is_some() {
                 let mut worktree = None;
 
                 if let Some(buffer) = &singleton_buffer {
@@ -263,7 +263,7 @@ fn assign_edit_prediction_provider(
                     .map(|workspace| workspace.downgrade());
 
                 let zeta =
-                    zeta::Zeta::register(workspace, worktree, client.clone(), cloud_user_store, cx);
+                    zeta::Zeta::register(workspace, worktree, client.clone(), user_store, cx);
 
                 if let Some(buffer) = &singleton_buffer {
                     if buffer.read(cx).file().is_some() {

crates/zeta/src/zeta.rs 🔗

@@ -16,7 +16,7 @@ pub use rate_completion_modal::*;
 
 use anyhow::{Context as _, Result, anyhow};
 use arrayvec::ArrayVec;
-use client::{Client, CloudUserStore, EditPredictionUsage};
+use client::{Client, EditPredictionUsage, UserStore};
 use cloud_llm_client::{
     AcceptEditPredictionBody, EXPIRED_LLM_TOKEN_HEADER_NAME, MINIMUM_REQUIRED_VERSION_HEADER_NAME,
     PredictEditsBody, PredictEditsResponse, ZED_VERSION_HEADER_NAME,
@@ -120,8 +120,8 @@ impl Dismissable for ZedPredictUpsell {
     }
 }
 
-pub fn should_show_upsell_modal(cloud_user_store: &Entity<CloudUserStore>, cx: &App) -> bool {
-    if cloud_user_store.read(cx).has_accepted_tos() {
+pub fn should_show_upsell_modal(user_store: &Entity<UserStore>, cx: &App) -> bool {
+    if user_store.read(cx).has_accepted_terms_of_service() {
         !ZedPredictUpsell::dismissed()
     } else {
         true
@@ -229,7 +229,7 @@ pub struct Zeta {
     _llm_token_subscription: Subscription,
     /// Whether an update to a newer version of Zed is required to continue using Zeta.
     update_required: bool,
-    cloud_user_store: Entity<CloudUserStore>,
+    user_store: Entity<UserStore>,
     license_detection_watchers: HashMap<WorktreeId, Rc<LicenseDetectionWatcher>>,
 }
 
@@ -242,11 +242,11 @@ impl Zeta {
         workspace: Option<WeakEntity<Workspace>>,
         worktree: Option<Entity<Worktree>>,
         client: Arc<Client>,
-        cloud_user_store: Entity<CloudUserStore>,
+        user_store: Entity<UserStore>,
         cx: &mut App,
     ) -> Entity<Self> {
         let this = Self::global(cx).unwrap_or_else(|| {
-            let entity = cx.new(|cx| Self::new(workspace, client, cloud_user_store, cx));
+            let entity = cx.new(|cx| Self::new(workspace, client, user_store, cx));
             cx.set_global(ZetaGlobal(entity.clone()));
             entity
         });
@@ -269,13 +269,13 @@ impl Zeta {
     }
 
     pub fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
-        self.cloud_user_store.read(cx).edit_prediction_usage()
+        self.user_store.read(cx).edit_prediction_usage()
     }
 
     fn new(
         workspace: Option<WeakEntity<Workspace>>,
         client: Arc<Client>,
-        cloud_user_store: Entity<CloudUserStore>,
+        user_store: Entity<UserStore>,
         cx: &mut Context<Self>,
     ) -> Self {
         let refresh_llm_token_listener = RefreshLlmTokenListener::global(cx);
@@ -306,7 +306,7 @@ impl Zeta {
             ),
             update_required: false,
             license_detection_watchers: HashMap::default(),
-            cloud_user_store,
+            user_store,
         }
     }
 
@@ -535,8 +535,8 @@ impl Zeta {
 
             if let Some(usage) = usage {
                 this.update(cx, |this, cx| {
-                    this.cloud_user_store.update(cx, |cloud_user_store, cx| {
-                        cloud_user_store.update_edit_prediction_usage(usage, cx);
+                    this.user_store.update(cx, |user_store, cx| {
+                        user_store.update_edit_prediction_usage(usage, cx);
                     });
                 })
                 .ok();
@@ -877,8 +877,8 @@ and then another
             if response.status().is_success() {
                 if let Some(usage) = EditPredictionUsage::from_headers(response.headers()).ok() {
                     this.update(cx, |this, cx| {
-                        this.cloud_user_store.update(cx, |cloud_user_store, cx| {
-                            cloud_user_store.update_edit_prediction_usage(usage, cx);
+                        this.user_store.update(cx, |user_store, cx| {
+                            user_store.update_edit_prediction_usage(usage, cx);
                         });
                     })?;
                 }
@@ -1559,9 +1559,9 @@ impl inline_completion::EditPredictionProvider for ZetaInlineCompletionProvider
         !self
             .zeta
             .read(cx)
-            .cloud_user_store
+            .user_store
             .read(cx)
-            .has_accepted_tos()
+            .has_accepted_terms_of_service()
     }
 
     fn is_refreshing(&self) -> bool {
@@ -1587,7 +1587,7 @@ impl inline_completion::EditPredictionProvider for ZetaInlineCompletionProvider
         if self
             .zeta
             .read(cx)
-            .cloud_user_store
+            .user_store
             .read_with(cx, |cloud_user_store, _cx| {
                 cloud_user_store.account_too_young() || cloud_user_store.has_overdue_invoices()
             })
@@ -1808,10 +1808,7 @@ mod tests {
     use client::UserStore;
     use client::test::FakeServer;
     use clock::FakeSystemClock;
-    use cloud_api_types::{
-        AuthenticatedUser, CreateLlmTokenResponse, GetAuthenticatedUserResponse, LlmToken, PlanInfo,
-    };
-    use cloud_llm_client::{CurrentUsage, Plan, UsageData, UsageLimit};
+    use cloud_api_types::{CreateLlmTokenResponse, LlmToken};
     use gpui::TestAppContext;
     use http_client::FakeHttpClient;
     use indoc::indoc;
@@ -1820,39 +1817,6 @@ mod tests {
 
     use super::*;
 
-    fn make_get_authenticated_user_response() -> GetAuthenticatedUserResponse {
-        GetAuthenticatedUserResponse {
-            user: AuthenticatedUser {
-                id: 1,
-                metrics_id: "metrics-id-1".to_string(),
-                avatar_url: "".to_string(),
-                github_login: "".to_string(),
-                name: None,
-                is_staff: false,
-                accepted_tos_at: None,
-            },
-            feature_flags: vec![],
-            plan: PlanInfo {
-                plan: Plan::ZedPro,
-                subscription_period: None,
-                usage: CurrentUsage {
-                    model_requests: UsageData {
-                        used: 0,
-                        limit: UsageLimit::Limited(500),
-                    },
-                    edit_predictions: UsageData {
-                        used: 250,
-                        limit: UsageLimit::Unlimited,
-                    },
-                },
-                trial_started_at: None,
-                is_usage_based_billing_enabled: false,
-                is_account_too_young: false,
-                has_overdue_invoices: false,
-            },
-        }
-    }
-
     #[gpui::test]
     async fn test_inline_completion_basic_interpolation(cx: &mut TestAppContext) {
         let buffer = cx.new(|cx| Buffer::local("Lorem ipsum dolor", cx));
@@ -2054,14 +2018,6 @@ mod tests {
 
         let http_client = FakeHttpClient::create(move |req| async move {
             match (req.method(), req.uri().path()) {
-                (&Method::GET, "/client/users/me") => Ok(http_client::Response::builder()
-                    .status(200)
-                    .body(
-                        serde_json::to_string(&make_get_authenticated_user_response())
-                            .unwrap()
-                            .into(),
-                    )
-                    .unwrap()),
                 (&Method::POST, "/client/llm_tokens") => Ok(http_client::Response::builder()
                     .status(200)
                     .body(
@@ -2098,9 +2054,7 @@ mod tests {
         // Construct the fake server to authenticate.
         let _server = FakeServer::for_client(42, &client, cx).await;
         let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
-        let cloud_user_store =
-            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store.clone(), cx));
-        let zeta = cx.new(|cx| Zeta::new(None, client, cloud_user_store, cx));
+        let zeta = cx.new(|cx| Zeta::new(None, client, user_store.clone(), cx));
 
         let buffer = cx.new(|cx| Buffer::local(buffer_content, cx));
         let cursor = buffer.read_with(cx, |buffer, _| buffer.anchor_before(Point::new(1, 0)));
@@ -2128,14 +2082,6 @@ mod tests {
             let completion = completion_response.clone();
             async move {
                 match (req.method(), req.uri().path()) {
-                    (&Method::GET, "/client/users/me") => Ok(http_client::Response::builder()
-                        .status(200)
-                        .body(
-                            serde_json::to_string(&make_get_authenticated_user_response())
-                                .unwrap()
-                                .into(),
-                        )
-                        .unwrap()),
                     (&Method::POST, "/client/llm_tokens") => Ok(http_client::Response::builder()
                         .status(200)
                         .body(
@@ -2172,9 +2118,7 @@ mod tests {
         // Construct the fake server to authenticate.
         let _server = FakeServer::for_client(42, &client, cx).await;
         let user_store = cx.new(|cx| UserStore::new(client.clone(), cx));
-        let cloud_user_store =
-            cx.new(|cx| CloudUserStore::new(client.cloud_client(), user_store.clone(), cx));
-        let zeta = cx.new(|cx| Zeta::new(None, client, cloud_user_store, cx));
+        let zeta = cx.new(|cx| Zeta::new(None, client, user_store.clone(), cx));
 
         let buffer = cx.new(|cx| Buffer::local(buffer_content, cx));
         let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());