Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Start separating authentication from connection to collab (#35471)

Browse source 
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>
This commit is contained in:
Antonio Scandurra 2025-08-01 19:37:38 +02:00 • committed by GitHub
parent b01d1872cc
commit f888f3fc0b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
46 changed files with 653 additions and 855 deletions
Download patch file Download diff file Expand all files Collapse all files

230
crates/client/src/user.rs
View file

@ -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<()> {
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;
let client = this
.read_with(&cx, |this, _| this.client.upgrade())?
.context("client was dropped")?;
cx.emit(Event::PlanUpdated);
cx.notify();
})?;
Ok(())
let response = client
.cloud_client()
.get_authenticated_user()
.await
.context("failed to fetch authenticated user")?;
this.update(&mut cx, |this, cx| {
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()
}
/// 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,
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()
}
/// 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 current_user_has_accepted_terms(&self) -> Option<bool> {
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>,