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

@@ -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,

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,

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)]