client: Store organization configuration

Marshall Bowers created 1 day ago

Change summary

crates/client/src/test.rs                     |  1 +
crates/client/src/user.rs                     | 12 ++++++++++++
crates/cloud_api_types/src/cloud_api_types.rs | 15 +++++++++++++++
3 files changed, 28 insertions(+)

Detailed changes

crates/client/src/test.rs

@@ -271,6 +271,7 @@ pub fn make_get_authenticated_user_response(
         organizations: vec![],
         default_organization_id: None,
         plans_by_organization: BTreeMap::new(),
+        configuration_by_organization: BTreeMap::new(),
         plan: PlanInfo {
             plan: KnownOrUnknown::Known(Plan::ZedPro),
             subscription_period: None,

crates/client/src/user.rs

@@ -5,6 +5,7 @@ use cloud_api_client::websocket_protocol::MessageToClient;
 use cloud_api_client::{
     GetAuthenticatedUserResponse, KnownOrUnknown, Organization, OrganizationId, Plan, PlanInfo,
 };
+use cloud_api_types::OrganizationConfiguration;
 use cloud_llm_client::{
     EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME, EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME, UsageLimit,
 };
@@ -117,6 +118,7 @@ pub struct UserStore {
     current_organization: Option<Arc<Organization>>,
     organizations: Vec<Arc<Organization>>,
     plans_by_organization: HashMap<OrganizationId, Plan>,
+    configuration_by_organization: HashMap<OrganizationId, OrganizationConfiguration>,
     contacts: Vec<Arc<Contact>>,
     incoming_contact_requests: Vec<Arc<User>>,
     outgoing_contact_requests: Vec<Arc<User>>,
@@ -193,6 +195,7 @@ impl UserStore {
             current_organization: None,
             organizations: Vec::new(),
             plans_by_organization: HashMap::default(),
+            configuration_by_organization: HashMap::default(),
             plan_info: None,
             edit_prediction_usage: None,
             contacts: Default::default(),
@@ -730,6 +733,13 @@ impl UserStore {
         self.plans_by_organization.get(organization_id).copied()
     }
 
+    pub fn current_organization_configuration(&self) -> Option<&OrganizationConfiguration> {
+        let current_organization = self.current_organization.as_ref()?;
+
+        self.configuration_by_organization
+            .get(&current_organization.id)
+    }
+
     pub fn plan(&self) -> Option<Plan> {
         #[cfg(debug_assertions)]
         if let Ok(plan) = std::env::var("ZED_SIMULATE_PLAN").as_ref() {
@@ -865,6 +875,8 @@ impl UserStore {
                 (organization_id, plan)
             })
             .collect();
+        self.configuration_by_organization =
+            response.configuration_by_organization.into_iter().collect();
 
         self.edit_prediction_usage = Some(EditPredictionUsage(RequestUsage {
             limit: response.plan.usage.edit_predictions.limit,

crates/cloud_api_types/src/cloud_api_types.rs

@@ -26,6 +26,8 @@ pub struct GetAuthenticatedUserResponse {
     pub default_organization_id: Option<OrganizationId>,
     #[serde(default)]
     pub plans_by_organization: BTreeMap<OrganizationId, KnownOrUnknown<Plan, String>>,
+    #[serde(default)]
+    pub configuration_by_organization: BTreeMap<OrganizationId, OrganizationConfiguration>,
     pub plan: PlanInfo,
 }
 
@@ -50,6 +52,19 @@ pub struct Organization {
     pub is_personal: bool,
 }
 
+#[derive(Debug, PartialEq, Serialize, Deserialize)]
+pub struct OrganizationConfiguration {
+    pub is_zed_model_provider_enabled: bool,
+    pub is_agent_thread_feedback_enabled: bool,
+    pub edit_prediction: OrganizationEditPredictionConfiguration,
+}
+
+#[derive(Debug, PartialEq, Serialize, Deserialize)]
+pub struct OrganizationEditPredictionConfiguration {
+    pub is_enabled: bool,
+    pub is_feedback_enabled: bool,
+}
+
 #[derive(Debug, PartialEq, Serialize, Deserialize)]
 pub struct AcceptTermsOfServiceResponse {
     pub user: AuthenticatedUser,