From bb4f771f0e28a07d980edf8ca8fa6a6f596d1512 Mon Sep 17 00:00:00 2001
From: Marshall Bowers <git@maxdeviant.com>
Date: Wed, 11 Mar 2026 18:22:17 -0400
Subject: [PATCH] client: Populate plans for organizations (#51334)

This PR makes it so we populate the `plans_by_organization` collection
with the plans returned from the server.

Release Notes:

- N/A
---
 crates/client/src/test.rs                     |  8 ++++++--
 crates/client/src/user.rs                     | 17 ++++++++++++++++-
 crates/cloud_api_types/src/cloud_api_types.rs |  5 ++++-
 3 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/crates/client/src/test.rs b/crates/client/src/test.rs
index 5102664a8c08ba336f3ae506aadb68eb2a537935..b506cee822ff9c2e4e31f262886a26ac1acbd134 100644
--- a/crates/client/src/test.rs
+++ b/crates/client/src/test.rs
@@ -1,4 +1,6 @@
-use crate::{Client, Connection, Credentials, EstablishConnectionError, UserStore};
+use std::collections::BTreeMap;
+use std::sync::Arc;
+
 use anyhow::{Context as _, Result, anyhow};
 use cloud_api_client::{
     AuthenticatedUser, GetAuthenticatedUserResponse, KnownOrUnknown, Plan, PlanInfo,
@@ -9,7 +11,8 @@ use gpui::{AppContext as _, Entity, TestAppContext};
 use http_client::{AsyncBody, Method, Request, http};
 use parking_lot::Mutex;
 use rpc::{ConnectionId, Peer, Receipt, TypedEnvelope, proto};
-use std::sync::Arc;
+
+use crate::{Client, Connection, Credentials, EstablishConnectionError, UserStore};
 
 pub struct FakeServer {
     peer: Arc<Peer>,
@@ -266,6 +269,7 @@ pub fn make_get_authenticated_user_response(
         },
         feature_flags: vec![],
         organizations: vec![],
+        plans_by_organization: BTreeMap::new(),
         plan: PlanInfo {
             plan: KnownOrUnknown::Known(Plan::ZedPro),
             subscription_period: None,
diff --git a/crates/client/src/user.rs b/crates/client/src/user.rs
index 5d38569cfd86c38e5b4780621db40d1f2a3b745c..71b05dc58f54379f8dfb2ec46d4c280926a56bea 100644
--- a/crates/client/src/user.rs
+++ b/crates/client/src/user.rs
@@ -3,7 +3,7 @@ use anyhow::{Context as _, Result};
 use chrono::{DateTime, Utc};
 use cloud_api_client::websocket_protocol::MessageToClient;
 use cloud_api_client::{
-    GetAuthenticatedUserResponse, Organization, OrganizationId, Plan, PlanInfo,
+    GetAuthenticatedUserResponse, KnownOrUnknown, Organization, OrganizationId, Plan, PlanInfo,
 };
 use cloud_llm_client::{
     EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME, EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME, UsageLimit,
@@ -817,6 +817,21 @@ impl UserStore {
 
         self.organizations = response.organizations.into_iter().map(Arc::new).collect();
         self.current_organization = self.organizations.first().cloned();
+        self.plans_by_organization = response
+            .plans_by_organization
+            .into_iter()
+            .map(|(organization_id, plan)| {
+                let plan = match plan {
+                    KnownOrUnknown::Known(plan) => plan,
+                    KnownOrUnknown::Unknown(_) => {
+                        // If we get a plan that we don't recognize, fall back to the Free plan.
+                        Plan::ZedFree
+                    }
+                };
+
+                (organization_id, plan)
+            })
+            .collect();
 
         self.edit_prediction_usage = Some(EditPredictionUsage(RequestUsage {
             limit: response.plan.usage.edit_predictions.limit,
diff --git a/crates/cloud_api_types/src/cloud_api_types.rs b/crates/cloud_api_types/src/cloud_api_types.rs
index 42d3442bfc016f5cb1a39ba421ccdfe386bcbc65..e2c517edcc78e37bc2eab7055c5ac8d79c9db5b2 100644
--- a/crates/cloud_api_types/src/cloud_api_types.rs
+++ b/crates/cloud_api_types/src/cloud_api_types.rs
@@ -4,6 +4,7 @@ mod plan;
 mod timestamp;
 pub mod websocket_protocol;
 
+use std::collections::BTreeMap;
 use std::sync::Arc;
 
 use serde::{Deserialize, Serialize};
@@ -21,6 +22,8 @@ pub struct GetAuthenticatedUserResponse {
     pub feature_flags: Vec<String>,
     #[serde(default)]
     pub organizations: Vec<Organization>,
+    #[serde(default)]
+    pub plans_by_organization: BTreeMap<OrganizationId, KnownOrUnknown<Plan, String>>,
     pub plan: PlanInfo,
 }
 
@@ -35,7 +38,7 @@ pub struct AuthenticatedUser {
     pub accepted_tos_at: Option<Timestamp>,
 }
 
-#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Serialize, Deserialize)]
 pub struct OrganizationId(pub Arc<str>);
 
 #[derive(Debug, PartialEq, Serialize, Deserialize)]