1mod connection_pool;
2
3use crate::api::billing::find_or_create_billing_customer;
4use crate::api::{CloudflareIpCountryHeader, SystemIdHeader};
5use crate::db::billing_subscription::SubscriptionKind;
6use crate::llm::db::LlmDatabase;
7use crate::llm::{AGENT_EXTENDED_TRIAL_FEATURE_FLAG, LlmTokenClaims};
8use crate::{
9 AppState, Error, Result, auth,
10 db::{
11 self, BufferId, Capability, Channel, ChannelId, ChannelRole, ChannelsForUser,
12 CreatedChannelMessage, Database, InviteMemberResult, MembershipUpdated, MessageId,
13 NotificationId, Project, ProjectId, RejoinedProject, RemoveChannelMemberResult, ReplicaId,
14 RespondToChannelInvite, RoomId, ServerId, UpdatedChannelMessage, User, UserId,
15 },
16 executor::Executor,
17};
18use anyhow::{Context as _, anyhow, bail};
19use async_tungstenite::tungstenite::{
20 Message as TungsteniteMessage, protocol::CloseFrame as TungsteniteCloseFrame,
21};
22use axum::{
23 Extension, Router, TypedHeader,
24 body::Body,
25 extract::{
26 ConnectInfo, WebSocketUpgrade,
27 ws::{CloseFrame as AxumCloseFrame, Message as AxumMessage},
28 },
29 headers::{Header, HeaderName},
30 http::StatusCode,
31 middleware,
32 response::IntoResponse,
33 routing::get,
34};
35use chrono::Utc;
36use collections::{HashMap, HashSet};
37pub use connection_pool::{ConnectionPool, ZedVersion};
38use core::fmt::{self, Debug, Formatter};
39use reqwest_client::ReqwestClient;
40use rpc::proto::split_repository_update;
41use supermaven_api::{CreateExternalUserRequest, SupermavenAdminApi};
42
43use futures::{
44 FutureExt, SinkExt, StreamExt, TryStreamExt, channel::oneshot, future::BoxFuture,
45 stream::FuturesUnordered,
46};
47use prometheus::{IntGauge, register_int_gauge};
48use rpc::{
49 Connection, ConnectionId, ErrorCode, ErrorCodeExt, ErrorExt, Peer, Receipt, TypedEnvelope,
50 proto::{
51 self, Ack, AnyTypedEnvelope, EntityMessage, EnvelopedMessage, LiveKitConnectionInfo,
52 RequestMessage, ShareProject, UpdateChannelBufferCollaborators,
53 },
54};
55use semantic_version::SemanticVersion;
56use serde::{Serialize, Serializer};
57use std::{
58 any::TypeId,
59 future::Future,
60 marker::PhantomData,
61 mem,
62 net::SocketAddr,
63 ops::{Deref, DerefMut},
64 rc::Rc,
65 sync::{
66 Arc, OnceLock,
67 atomic::{AtomicBool, Ordering::SeqCst},
68 },
69 time::{Duration, Instant},
70};
71use time::OffsetDateTime;
72use tokio::sync::{Semaphore, watch};
73use tower::ServiceBuilder;
74use tracing::{
75 Instrument,
76 field::{self},
77 info_span, instrument,
78};
79
80pub const RECONNECT_TIMEOUT: Duration = Duration::from_secs(30);
81
82// kubernetes gives terminated pods 10s to shutdown gracefully. After they're gone, we can clean up old resources.
83pub const CLEANUP_TIMEOUT: Duration = Duration::from_secs(15);
84
85const MESSAGE_COUNT_PER_PAGE: usize = 100;
86const MAX_MESSAGE_LEN: usize = 1024;
87const NOTIFICATION_COUNT_PER_PAGE: usize = 50;
88
89type MessageHandler =
90 Box<dyn Send + Sync + Fn(Box<dyn AnyTypedEnvelope>, Session) -> BoxFuture<'static, ()>>;
91
92struct Response<R> {
93 peer: Arc<Peer>,
94 receipt: Receipt<R>,
95 responded: Arc<AtomicBool>,
96}
97
98impl<R: RequestMessage> Response<R> {
99 fn send(self, payload: R::Response) -> Result<()> {
100 self.responded.store(true, SeqCst);
101 self.peer.respond(self.receipt, payload)?;
102 Ok(())
103 }
104}
105
106#[derive(Clone, Debug)]
107pub enum Principal {
108 User(User),
109 Impersonated { user: User, admin: User },
110}
111
112impl Principal {
113 fn user(&self) -> &User {
114 match self {
115 Principal::User(user) => user,
116 Principal::Impersonated { user, .. } => user,
117 }
118 }
119
120 fn update_span(&self, span: &tracing::Span) {
121 match &self {
122 Principal::User(user) => {
123 span.record("user_id", user.id.0);
124 span.record("login", &user.github_login);
125 }
126 Principal::Impersonated { user, admin } => {
127 span.record("user_id", user.id.0);
128 span.record("login", &user.github_login);
129 span.record("impersonator", &admin.github_login);
130 }
131 }
132 }
133}
134
135#[derive(Clone)]
136struct Session {
137 principal: Principal,
138 connection_id: ConnectionId,
139 db: Arc<tokio::sync::Mutex<DbHandle>>,
140 peer: Arc<Peer>,
141 connection_pool: Arc<parking_lot::Mutex<ConnectionPool>>,
142 app_state: Arc<AppState>,
143 supermaven_client: Option<Arc<SupermavenAdminApi>>,
144 /// The GeoIP country code for the user.
145 #[allow(unused)]
146 geoip_country_code: Option<String>,
147 system_id: Option<String>,
148 _executor: Executor,
149}
150
151impl Session {
152 async fn db(&self) -> tokio::sync::MutexGuard<DbHandle> {
153 #[cfg(test)]
154 tokio::task::yield_now().await;
155 let guard = self.db.lock().await;
156 #[cfg(test)]
157 tokio::task::yield_now().await;
158 guard
159 }
160
161 async fn connection_pool(&self) -> ConnectionPoolGuard<'_> {
162 #[cfg(test)]
163 tokio::task::yield_now().await;
164 let guard = self.connection_pool.lock();
165 ConnectionPoolGuard {
166 guard,
167 _not_send: PhantomData,
168 }
169 }
170
171 fn is_staff(&self) -> bool {
172 match &self.principal {
173 Principal::User(user) => user.admin,
174 Principal::Impersonated { .. } => true,
175 }
176 }
177
178 fn user_id(&self) -> UserId {
179 match &self.principal {
180 Principal::User(user) => user.id,
181 Principal::Impersonated { user, .. } => user.id,
182 }
183 }
184
185 pub fn email(&self) -> Option<String> {
186 match &self.principal {
187 Principal::User(user) => user.email_address.clone(),
188 Principal::Impersonated { user, .. } => user.email_address.clone(),
189 }
190 }
191}
192
193impl Debug for Session {
194 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
195 let mut result = f.debug_struct("Session");
196 match &self.principal {
197 Principal::User(user) => {
198 result.field("user", &user.github_login);
199 }
200 Principal::Impersonated { user, admin } => {
201 result.field("user", &user.github_login);
202 result.field("impersonator", &admin.github_login);
203 }
204 }
205 result.field("connection_id", &self.connection_id).finish()
206 }
207}
208
209struct DbHandle(Arc<Database>);
210
211impl Deref for DbHandle {
212 type Target = Database;
213
214 fn deref(&self) -> &Self::Target {
215 self.0.as_ref()
216 }
217}
218
219pub struct Server {
220 id: parking_lot::Mutex<ServerId>,
221 peer: Arc<Peer>,
222 pub(crate) connection_pool: Arc<parking_lot::Mutex<ConnectionPool>>,
223 app_state: Arc<AppState>,
224 handlers: HashMap<TypeId, MessageHandler>,
225 teardown: watch::Sender<bool>,
226}
227
228pub(crate) struct ConnectionPoolGuard<'a> {
229 guard: parking_lot::MutexGuard<'a, ConnectionPool>,
230 _not_send: PhantomData<Rc<()>>,
231}
232
233#[derive(Serialize)]
234pub struct ServerSnapshot<'a> {
235 peer: &'a Peer,
236 #[serde(serialize_with = "serialize_deref")]
237 connection_pool: ConnectionPoolGuard<'a>,
238}
239
240pub fn serialize_deref<S, T, U>(value: &T, serializer: S) -> Result<S::Ok, S::Error>
241where
242 S: Serializer,
243 T: Deref<Target = U>,
244 U: Serialize,
245{
246 Serialize::serialize(value.deref(), serializer)
247}
248
249impl Server {
250 pub fn new(id: ServerId, app_state: Arc<AppState>) -> Arc<Self> {
251 let mut server = Self {
252 id: parking_lot::Mutex::new(id),
253 peer: Peer::new(id.0 as u32),
254 app_state: app_state.clone(),
255 connection_pool: Default::default(),
256 handlers: Default::default(),
257 teardown: watch::channel(false).0,
258 };
259
260 server
261 .add_request_handler(ping)
262 .add_request_handler(create_room)
263 .add_request_handler(join_room)
264 .add_request_handler(rejoin_room)
265 .add_request_handler(leave_room)
266 .add_request_handler(set_room_participant_role)
267 .add_request_handler(call)
268 .add_request_handler(cancel_call)
269 .add_message_handler(decline_call)
270 .add_request_handler(update_participant_location)
271 .add_request_handler(share_project)
272 .add_message_handler(unshare_project)
273 .add_request_handler(join_project)
274 .add_message_handler(leave_project)
275 .add_request_handler(update_project)
276 .add_request_handler(update_worktree)
277 .add_request_handler(update_repository)
278 .add_request_handler(remove_repository)
279 .add_message_handler(start_language_server)
280 .add_message_handler(update_language_server)
281 .add_message_handler(update_diagnostic_summary)
282 .add_message_handler(update_worktree_settings)
283 .add_request_handler(forward_read_only_project_request::<proto::GetHover>)
284 .add_request_handler(forward_read_only_project_request::<proto::GetDefinition>)
285 .add_request_handler(forward_read_only_project_request::<proto::GetTypeDefinition>)
286 .add_request_handler(forward_read_only_project_request::<proto::GetReferences>)
287 .add_request_handler(forward_find_search_candidates_request)
288 .add_request_handler(forward_read_only_project_request::<proto::GetDocumentHighlights>)
289 .add_request_handler(forward_read_only_project_request::<proto::GetDocumentSymbols>)
290 .add_request_handler(forward_read_only_project_request::<proto::GetProjectSymbols>)
291 .add_request_handler(forward_read_only_project_request::<proto::OpenBufferForSymbol>)
292 .add_request_handler(forward_read_only_project_request::<proto::OpenBufferById>)
293 .add_request_handler(forward_read_only_project_request::<proto::SynchronizeBuffers>)
294 .add_request_handler(forward_read_only_project_request::<proto::InlayHints>)
295 .add_request_handler(forward_read_only_project_request::<proto::ResolveInlayHint>)
296 .add_request_handler(forward_mutating_project_request::<proto::GetCodeLens>)
297 .add_request_handler(forward_read_only_project_request::<proto::OpenBufferByPath>)
298 .add_request_handler(forward_read_only_project_request::<proto::GitGetBranches>)
299 .add_request_handler(forward_read_only_project_request::<proto::OpenUnstagedDiff>)
300 .add_request_handler(forward_read_only_project_request::<proto::OpenUncommittedDiff>)
301 .add_request_handler(forward_read_only_project_request::<proto::LspExtExpandMacro>)
302 .add_request_handler(forward_read_only_project_request::<proto::LspExtOpenDocs>)
303 .add_request_handler(forward_mutating_project_request::<proto::LspExtRunnables>)
304 .add_request_handler(
305 forward_read_only_project_request::<proto::LspExtSwitchSourceHeader>,
306 )
307 .add_request_handler(forward_read_only_project_request::<proto::LspExtGoToParentModule>)
308 .add_request_handler(forward_read_only_project_request::<proto::LspExtCancelFlycheck>)
309 .add_request_handler(forward_read_only_project_request::<proto::LspExtRunFlycheck>)
310 .add_request_handler(forward_read_only_project_request::<proto::LspExtClearFlycheck>)
311 .add_request_handler(
312 forward_read_only_project_request::<proto::LanguageServerIdForName>,
313 )
314 .add_request_handler(
315 forward_mutating_project_request::<proto::RegisterBufferWithLanguageServers>,
316 )
317 .add_request_handler(forward_mutating_project_request::<proto::UpdateGitBranch>)
318 .add_request_handler(forward_mutating_project_request::<proto::GetCompletions>)
319 .add_request_handler(
320 forward_mutating_project_request::<proto::ApplyCompletionAdditionalEdits>,
321 )
322 .add_request_handler(forward_mutating_project_request::<proto::OpenNewBuffer>)
323 .add_request_handler(
324 forward_mutating_project_request::<proto::ResolveCompletionDocumentation>,
325 )
326 .add_request_handler(forward_mutating_project_request::<proto::GetCodeActions>)
327 .add_request_handler(forward_mutating_project_request::<proto::ApplyCodeAction>)
328 .add_request_handler(forward_mutating_project_request::<proto::PrepareRename>)
329 .add_request_handler(forward_mutating_project_request::<proto::PerformRename>)
330 .add_request_handler(forward_mutating_project_request::<proto::ReloadBuffers>)
331 .add_request_handler(forward_mutating_project_request::<proto::ApplyCodeActionKind>)
332 .add_request_handler(forward_mutating_project_request::<proto::FormatBuffers>)
333 .add_request_handler(forward_mutating_project_request::<proto::CreateProjectEntry>)
334 .add_request_handler(forward_mutating_project_request::<proto::RenameProjectEntry>)
335 .add_request_handler(forward_mutating_project_request::<proto::CopyProjectEntry>)
336 .add_request_handler(forward_mutating_project_request::<proto::DeleteProjectEntry>)
337 .add_request_handler(forward_mutating_project_request::<proto::ExpandProjectEntry>)
338 .add_request_handler(
339 forward_mutating_project_request::<proto::ExpandAllForProjectEntry>,
340 )
341 .add_request_handler(forward_mutating_project_request::<proto::OnTypeFormatting>)
342 .add_request_handler(forward_mutating_project_request::<proto::SaveBuffer>)
343 .add_request_handler(forward_mutating_project_request::<proto::BlameBuffer>)
344 .add_request_handler(forward_mutating_project_request::<proto::MultiLspQuery>)
345 .add_request_handler(forward_mutating_project_request::<proto::RestartLanguageServers>)
346 .add_request_handler(forward_mutating_project_request::<proto::StopLanguageServers>)
347 .add_request_handler(forward_mutating_project_request::<proto::LinkedEditingRange>)
348 .add_message_handler(create_buffer_for_peer)
349 .add_request_handler(update_buffer)
350 .add_message_handler(broadcast_project_message_from_host::<proto::RefreshInlayHints>)
351 .add_message_handler(broadcast_project_message_from_host::<proto::RefreshCodeLens>)
352 .add_message_handler(broadcast_project_message_from_host::<proto::UpdateBufferFile>)
353 .add_message_handler(broadcast_project_message_from_host::<proto::BufferReloaded>)
354 .add_message_handler(broadcast_project_message_from_host::<proto::BufferSaved>)
355 .add_message_handler(broadcast_project_message_from_host::<proto::UpdateDiffBases>)
356 .add_request_handler(get_users)
357 .add_request_handler(fuzzy_search_users)
358 .add_request_handler(request_contact)
359 .add_request_handler(remove_contact)
360 .add_request_handler(respond_to_contact_request)
361 .add_message_handler(subscribe_to_channels)
362 .add_request_handler(create_channel)
363 .add_request_handler(delete_channel)
364 .add_request_handler(invite_channel_member)
365 .add_request_handler(remove_channel_member)
366 .add_request_handler(set_channel_member_role)
367 .add_request_handler(set_channel_visibility)
368 .add_request_handler(rename_channel)
369 .add_request_handler(join_channel_buffer)
370 .add_request_handler(leave_channel_buffer)
371 .add_message_handler(update_channel_buffer)
372 .add_request_handler(rejoin_channel_buffers)
373 .add_request_handler(get_channel_members)
374 .add_request_handler(respond_to_channel_invite)
375 .add_request_handler(join_channel)
376 .add_request_handler(join_channel_chat)
377 .add_message_handler(leave_channel_chat)
378 .add_request_handler(send_channel_message)
379 .add_request_handler(remove_channel_message)
380 .add_request_handler(update_channel_message)
381 .add_request_handler(get_channel_messages)
382 .add_request_handler(get_channel_messages_by_id)
383 .add_request_handler(get_notifications)
384 .add_request_handler(mark_notification_as_read)
385 .add_request_handler(move_channel)
386 .add_request_handler(follow)
387 .add_message_handler(unfollow)
388 .add_message_handler(update_followers)
389 .add_request_handler(get_private_user_info)
390 .add_request_handler(get_llm_api_token)
391 .add_request_handler(accept_terms_of_service)
392 .add_message_handler(acknowledge_channel_message)
393 .add_message_handler(acknowledge_buffer_version)
394 .add_request_handler(get_supermaven_api_key)
395 .add_request_handler(forward_mutating_project_request::<proto::OpenContext>)
396 .add_request_handler(forward_mutating_project_request::<proto::CreateContext>)
397 .add_request_handler(forward_mutating_project_request::<proto::SynchronizeContexts>)
398 .add_request_handler(forward_mutating_project_request::<proto::Stage>)
399 .add_request_handler(forward_mutating_project_request::<proto::Unstage>)
400 .add_request_handler(forward_mutating_project_request::<proto::Commit>)
401 .add_request_handler(forward_mutating_project_request::<proto::GitInit>)
402 .add_request_handler(forward_read_only_project_request::<proto::GetRemotes>)
403 .add_request_handler(forward_read_only_project_request::<proto::GitShow>)
404 .add_request_handler(forward_read_only_project_request::<proto::LoadCommitDiff>)
405 .add_request_handler(forward_read_only_project_request::<proto::GitReset>)
406 .add_request_handler(forward_read_only_project_request::<proto::GitCheckoutFiles>)
407 .add_request_handler(forward_mutating_project_request::<proto::SetIndexText>)
408 .add_request_handler(forward_mutating_project_request::<proto::ToggleBreakpoint>)
409 .add_message_handler(broadcast_project_message_from_host::<proto::BreakpointsForFile>)
410 .add_request_handler(forward_mutating_project_request::<proto::OpenCommitMessageBuffer>)
411 .add_request_handler(forward_mutating_project_request::<proto::GitDiff>)
412 .add_request_handler(forward_mutating_project_request::<proto::GitCreateBranch>)
413 .add_request_handler(forward_mutating_project_request::<proto::GitChangeBranch>)
414 .add_request_handler(forward_mutating_project_request::<proto::CheckForPushedCommits>)
415 .add_message_handler(broadcast_project_message_from_host::<proto::AdvertiseContexts>)
416 .add_message_handler(update_context);
417
418 Arc::new(server)
419 }
420
421 pub async fn start(&self) -> Result<()> {
422 let server_id = *self.id.lock();
423 let app_state = self.app_state.clone();
424 let peer = self.peer.clone();
425 let timeout = self.app_state.executor.sleep(CLEANUP_TIMEOUT);
426 let pool = self.connection_pool.clone();
427 let livekit_client = self.app_state.livekit_client.clone();
428
429 let span = info_span!("start server");
430 self.app_state.executor.spawn_detached(
431 async move {
432 tracing::info!("waiting for cleanup timeout");
433 timeout.await;
434 tracing::info!("cleanup timeout expired, retrieving stale rooms");
435 if let Some((room_ids, channel_ids)) = app_state
436 .db
437 .stale_server_resource_ids(&app_state.config.zed_environment, server_id)
438 .await
439 .trace_err()
440 {
441 tracing::info!(stale_room_count = room_ids.len(), "retrieved stale rooms");
442 tracing::info!(
443 stale_channel_buffer_count = channel_ids.len(),
444 "retrieved stale channel buffers"
445 );
446
447 for channel_id in channel_ids {
448 if let Some(refreshed_channel_buffer) = app_state
449 .db
450 .clear_stale_channel_buffer_collaborators(channel_id, server_id)
451 .await
452 .trace_err()
453 {
454 for connection_id in refreshed_channel_buffer.connection_ids {
455 peer.send(
456 connection_id,
457 proto::UpdateChannelBufferCollaborators {
458 channel_id: channel_id.to_proto(),
459 collaborators: refreshed_channel_buffer
460 .collaborators
461 .clone(),
462 },
463 )
464 .trace_err();
465 }
466 }
467 }
468
469 for room_id in room_ids {
470 let mut contacts_to_update = HashSet::default();
471 let mut canceled_calls_to_user_ids = Vec::new();
472 let mut livekit_room = String::new();
473 let mut delete_livekit_room = false;
474
475 if let Some(mut refreshed_room) = app_state
476 .db
477 .clear_stale_room_participants(room_id, server_id)
478 .await
479 .trace_err()
480 {
481 tracing::info!(
482 room_id = room_id.0,
483 new_participant_count = refreshed_room.room.participants.len(),
484 "refreshed room"
485 );
486 room_updated(&refreshed_room.room, &peer);
487 if let Some(channel) = refreshed_room.channel.as_ref() {
488 channel_updated(channel, &refreshed_room.room, &peer, &pool.lock());
489 }
490 contacts_to_update
491 .extend(refreshed_room.stale_participant_user_ids.iter().copied());
492 contacts_to_update
493 .extend(refreshed_room.canceled_calls_to_user_ids.iter().copied());
494 canceled_calls_to_user_ids =
495 mem::take(&mut refreshed_room.canceled_calls_to_user_ids);
496 livekit_room = mem::take(&mut refreshed_room.room.livekit_room);
497 delete_livekit_room = refreshed_room.room.participants.is_empty();
498 }
499
500 {
501 let pool = pool.lock();
502 for canceled_user_id in canceled_calls_to_user_ids {
503 for connection_id in pool.user_connection_ids(canceled_user_id) {
504 peer.send(
505 connection_id,
506 proto::CallCanceled {
507 room_id: room_id.to_proto(),
508 },
509 )
510 .trace_err();
511 }
512 }
513 }
514
515 for user_id in contacts_to_update {
516 let busy = app_state.db.is_user_busy(user_id).await.trace_err();
517 let contacts = app_state.db.get_contacts(user_id).await.trace_err();
518 if let Some((busy, contacts)) = busy.zip(contacts) {
519 let pool = pool.lock();
520 let updated_contact = contact_for_user(user_id, busy, &pool);
521 for contact in contacts {
522 if let db::Contact::Accepted {
523 user_id: contact_user_id,
524 ..
525 } = contact
526 {
527 for contact_conn_id in
528 pool.user_connection_ids(contact_user_id)
529 {
530 peer.send(
531 contact_conn_id,
532 proto::UpdateContacts {
533 contacts: vec![updated_contact.clone()],
534 remove_contacts: Default::default(),
535 incoming_requests: Default::default(),
536 remove_incoming_requests: Default::default(),
537 outgoing_requests: Default::default(),
538 remove_outgoing_requests: Default::default(),
539 },
540 )
541 .trace_err();
542 }
543 }
544 }
545 }
546 }
547
548 if let Some(live_kit) = livekit_client.as_ref() {
549 if delete_livekit_room {
550 live_kit.delete_room(livekit_room).await.trace_err();
551 }
552 }
553 }
554 }
555
556 app_state
557 .db
558 .delete_stale_servers(&app_state.config.zed_environment, server_id)
559 .await
560 .trace_err();
561 }
562 .instrument(span),
563 );
564 Ok(())
565 }
566
567 pub fn teardown(&self) {
568 self.peer.teardown();
569 self.connection_pool.lock().reset();
570 let _ = self.teardown.send(true);
571 }
572
573 #[cfg(test)]
574 pub fn reset(&self, id: ServerId) {
575 self.teardown();
576 *self.id.lock() = id;
577 self.peer.reset(id.0 as u32);
578 let _ = self.teardown.send(false);
579 }
580
581 #[cfg(test)]
582 pub fn id(&self) -> ServerId {
583 *self.id.lock()
584 }
585
586 fn add_handler<F, Fut, M>(&mut self, handler: F) -> &mut Self
587 where
588 F: 'static + Send + Sync + Fn(TypedEnvelope<M>, Session) -> Fut,
589 Fut: 'static + Send + Future<Output = Result<()>>,
590 M: EnvelopedMessage,
591 {
592 let prev_handler = self.handlers.insert(
593 TypeId::of::<M>(),
594 Box::new(move |envelope, session| {
595 let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();
596 let received_at = envelope.received_at;
597 tracing::info!("message received");
598 let start_time = Instant::now();
599 let future = (handler)(*envelope, session);
600 async move {
601 let result = future.await;
602 let total_duration_ms = received_at.elapsed().as_micros() as f64 / 1000.0;
603 let processing_duration_ms = start_time.elapsed().as_micros() as f64 / 1000.0;
604 let queue_duration_ms = total_duration_ms - processing_duration_ms;
605 let payload_type = M::NAME;
606
607 match result {
608 Err(error) => {
609 tracing::error!(
610 ?error,
611 total_duration_ms,
612 processing_duration_ms,
613 queue_duration_ms,
614 payload_type,
615 "error handling message"
616 )
617 }
618 Ok(()) => tracing::info!(
619 total_duration_ms,
620 processing_duration_ms,
621 queue_duration_ms,
622 "finished handling message"
623 ),
624 }
625 }
626 .boxed()
627 }),
628 );
629 if prev_handler.is_some() {
630 panic!("registered a handler for the same message twice");
631 }
632 self
633 }
634
635 fn add_message_handler<F, Fut, M>(&mut self, handler: F) -> &mut Self
636 where
637 F: 'static + Send + Sync + Fn(M, Session) -> Fut,
638 Fut: 'static + Send + Future<Output = Result<()>>,
639 M: EnvelopedMessage,
640 {
641 self.add_handler(move |envelope, session| handler(envelope.payload, session));
642 self
643 }
644
645 fn add_request_handler<F, Fut, M>(&mut self, handler: F) -> &mut Self
646 where
647 F: 'static + Send + Sync + Fn(M, Response<M>, Session) -> Fut,
648 Fut: Send + Future<Output = Result<()>>,
649 M: RequestMessage,
650 {
651 let handler = Arc::new(handler);
652 self.add_handler(move |envelope, session| {
653 let receipt = envelope.receipt();
654 let handler = handler.clone();
655 async move {
656 let peer = session.peer.clone();
657 let responded = Arc::new(AtomicBool::default());
658 let response = Response {
659 peer: peer.clone(),
660 responded: responded.clone(),
661 receipt,
662 };
663 match (handler)(envelope.payload, response, session).await {
664 Ok(()) => {
665 if responded.load(std::sync::atomic::Ordering::SeqCst) {
666 Ok(())
667 } else {
668 Err(anyhow!("handler did not send a response"))?
669 }
670 }
671 Err(error) => {
672 let proto_err = match &error {
673 Error::Internal(err) => err.to_proto(),
674 _ => ErrorCode::Internal.message(format!("{error}")).to_proto(),
675 };
676 peer.respond_with_error(receipt, proto_err)?;
677 Err(error)
678 }
679 }
680 }
681 })
682 }
683
684 pub fn handle_connection(
685 self: &Arc<Self>,
686 connection: Connection,
687 address: String,
688 principal: Principal,
689 zed_version: ZedVersion,
690 geoip_country_code: Option<String>,
691 system_id: Option<String>,
692 send_connection_id: Option<oneshot::Sender<ConnectionId>>,
693 executor: Executor,
694 ) -> impl Future<Output = ()> + use<> {
695 let this = self.clone();
696 let span = info_span!("handle connection", %address,
697 connection_id=field::Empty,
698 user_id=field::Empty,
699 login=field::Empty,
700 impersonator=field::Empty,
701 geoip_country_code=field::Empty
702 );
703 principal.update_span(&span);
704 if let Some(country_code) = geoip_country_code.as_ref() {
705 span.record("geoip_country_code", country_code);
706 }
707
708 let mut teardown = self.teardown.subscribe();
709 async move {
710 if *teardown.borrow() {
711 tracing::error!("server is tearing down");
712 return
713 }
714 let (connection_id, handle_io, mut incoming_rx) = this
715 .peer
716 .add_connection(connection, {
717 let executor = executor.clone();
718 move |duration| executor.sleep(duration)
719 });
720 tracing::Span::current().record("connection_id", format!("{}", connection_id));
721
722 tracing::info!("connection opened");
723
724 let user_agent = format!("Zed Server/{}", env!("CARGO_PKG_VERSION"));
725 let http_client = match ReqwestClient::user_agent(&user_agent) {
726 Ok(http_client) => Arc::new(http_client),
727 Err(error) => {
728 tracing::error!(?error, "failed to create HTTP client");
729 return;
730 }
731 };
732
733 let supermaven_client = this.app_state.config.supermaven_admin_api_key.clone().map(|supermaven_admin_api_key| Arc::new(SupermavenAdminApi::new(
734 supermaven_admin_api_key.to_string(),
735 http_client.clone(),
736 )));
737
738 let session = Session {
739 principal: principal.clone(),
740 connection_id,
741 db: Arc::new(tokio::sync::Mutex::new(DbHandle(this.app_state.db.clone()))),
742 peer: this.peer.clone(),
743 connection_pool: this.connection_pool.clone(),
744 app_state: this.app_state.clone(),
745 geoip_country_code,
746 system_id,
747 _executor: executor.clone(),
748 supermaven_client,
749 };
750
751 if let Err(error) = this.send_initial_client_update(connection_id, zed_version, send_connection_id, &session).await {
752 tracing::error!(?error, "failed to send initial client update");
753 return;
754 }
755
756 let handle_io = handle_io.fuse();
757 futures::pin_mut!(handle_io);
758
759 // Handlers for foreground messages are pushed into the following `FuturesUnordered`.
760 // This prevents deadlocks when e.g., client A performs a request to client B and
761 // client B performs a request to client A. If both clients stop processing further
762 // messages until their respective request completes, they won't have a chance to
763 // respond to the other client's request and cause a deadlock.
764 //
765 // This arrangement ensures we will attempt to process earlier messages first, but fall
766 // back to processing messages arrived later in the spirit of making progress.
767 let mut foreground_message_handlers = FuturesUnordered::new();
768 let concurrent_handlers = Arc::new(Semaphore::new(256));
769 loop {
770 let next_message = async {
771 let permit = concurrent_handlers.clone().acquire_owned().await.unwrap();
772 let message = incoming_rx.next().await;
773 (permit, message)
774 }.fuse();
775 futures::pin_mut!(next_message);
776 futures::select_biased! {
777 _ = teardown.changed().fuse() => return,
778 result = handle_io => {
779 if let Err(error) = result {
780 tracing::error!(?error, "error handling I/O");
781 }
782 break;
783 }
784 _ = foreground_message_handlers.next() => {}
785 next_message = next_message => {
786 let (permit, message) = next_message;
787 if let Some(message) = message {
788 let type_name = message.payload_type_name();
789 // note: we copy all the fields from the parent span so we can query them in the logs.
790 // (https://github.com/tokio-rs/tracing/issues/2670).
791 let span = tracing::info_span!("receive message", %connection_id, %address, type_name,
792 user_id=field::Empty,
793 login=field::Empty,
794 impersonator=field::Empty,
795 );
796 principal.update_span(&span);
797 let span_enter = span.enter();
798 if let Some(handler) = this.handlers.get(&message.payload_type_id()) {
799 let is_background = message.is_background();
800 let handle_message = (handler)(message, session.clone());
801 drop(span_enter);
802
803 let handle_message = async move {
804 handle_message.await;
805 drop(permit);
806 }.instrument(span);
807 if is_background {
808 executor.spawn_detached(handle_message);
809 } else {
810 foreground_message_handlers.push(handle_message);
811 }
812 } else {
813 tracing::error!("no message handler");
814 }
815 } else {
816 tracing::info!("connection closed");
817 break;
818 }
819 }
820 }
821 }
822
823 drop(foreground_message_handlers);
824 tracing::info!("signing out");
825 if let Err(error) = connection_lost(session, teardown, executor).await {
826 tracing::error!(?error, "error signing out");
827 }
828
829 }.instrument(span)
830 }
831
832 async fn send_initial_client_update(
833 &self,
834 connection_id: ConnectionId,
835 zed_version: ZedVersion,
836 mut send_connection_id: Option<oneshot::Sender<ConnectionId>>,
837 session: &Session,
838 ) -> Result<()> {
839 self.peer.send(
840 connection_id,
841 proto::Hello {
842 peer_id: Some(connection_id.into()),
843 },
844 )?;
845 tracing::info!("sent hello message");
846 if let Some(send_connection_id) = send_connection_id.take() {
847 let _ = send_connection_id.send(connection_id);
848 }
849
850 match &session.principal {
851 Principal::User(user) | Principal::Impersonated { user, admin: _ } => {
852 if !user.connected_once {
853 self.peer.send(connection_id, proto::ShowContacts {})?;
854 self.app_state
855 .db
856 .set_user_connected_once(user.id, true)
857 .await?;
858 }
859
860 update_user_plan(session).await?;
861
862 let contacts = self.app_state.db.get_contacts(user.id).await?;
863
864 {
865 let mut pool = self.connection_pool.lock();
866 pool.add_connection(connection_id, user.id, user.admin, zed_version);
867 self.peer.send(
868 connection_id,
869 build_initial_contacts_update(contacts, &pool),
870 )?;
871 }
872
873 if should_auto_subscribe_to_channels(zed_version) {
874 subscribe_user_to_channels(user.id, session).await?;
875 }
876
877 if let Some(incoming_call) =
878 self.app_state.db.incoming_call_for_user(user.id).await?
879 {
880 self.peer.send(connection_id, incoming_call)?;
881 }
882
883 update_user_contacts(user.id, session).await?;
884 }
885 }
886
887 Ok(())
888 }
889
890 pub async fn invite_code_redeemed(
891 self: &Arc<Self>,
892 inviter_id: UserId,
893 invitee_id: UserId,
894 ) -> Result<()> {
895 if let Some(user) = self.app_state.db.get_user_by_id(inviter_id).await? {
896 if let Some(code) = &user.invite_code {
897 let pool = self.connection_pool.lock();
898 let invitee_contact = contact_for_user(invitee_id, false, &pool);
899 for connection_id in pool.user_connection_ids(inviter_id) {
900 self.peer.send(
901 connection_id,
902 proto::UpdateContacts {
903 contacts: vec![invitee_contact.clone()],
904 ..Default::default()
905 },
906 )?;
907 self.peer.send(
908 connection_id,
909 proto::UpdateInviteInfo {
910 url: format!("{}{}", self.app_state.config.invite_link_prefix, &code),
911 count: user.invite_count as u32,
912 },
913 )?;
914 }
915 }
916 }
917 Ok(())
918 }
919
920 pub async fn invite_count_updated(self: &Arc<Self>, user_id: UserId) -> Result<()> {
921 if let Some(user) = self.app_state.db.get_user_by_id(user_id).await? {
922 if let Some(invite_code) = &user.invite_code {
923 let pool = self.connection_pool.lock();
924 for connection_id in pool.user_connection_ids(user_id) {
925 self.peer.send(
926 connection_id,
927 proto::UpdateInviteInfo {
928 url: format!(
929 "{}{}",
930 self.app_state.config.invite_link_prefix, invite_code
931 ),
932 count: user.invite_count as u32,
933 },
934 )?;
935 }
936 }
937 }
938 Ok(())
939 }
940
941 pub async fn update_plan_for_user(self: &Arc<Self>, user_id: UserId) -> Result<()> {
942 let user = self
943 .app_state
944 .db
945 .get_user_by_id(user_id)
946 .await?
947 .context("user not found")?;
948
949 let update_user_plan = make_update_user_plan_message(
950 &user,
951 user.admin,
952 &self.app_state.db,
953 self.app_state.llm_db.clone(),
954 )
955 .await?;
956
957 let pool = self.connection_pool.lock();
958 for connection_id in pool.user_connection_ids(user_id) {
959 self.peer
960 .send(connection_id, update_user_plan.clone())
961 .trace_err();
962 }
963
964 Ok(())
965 }
966
967 pub async fn refresh_llm_tokens_for_user(self: &Arc<Self>, user_id: UserId) {
968 let pool = self.connection_pool.lock();
969 for connection_id in pool.user_connection_ids(user_id) {
970 self.peer
971 .send(connection_id, proto::RefreshLlmToken {})
972 .trace_err();
973 }
974 }
975
976 pub async fn snapshot(self: &Arc<Self>) -> ServerSnapshot {
977 ServerSnapshot {
978 connection_pool: ConnectionPoolGuard {
979 guard: self.connection_pool.lock(),
980 _not_send: PhantomData,
981 },
982 peer: &self.peer,
983 }
984 }
985}
986
987impl Deref for ConnectionPoolGuard<'_> {
988 type Target = ConnectionPool;
989
990 fn deref(&self) -> &Self::Target {
991 &self.guard
992 }
993}
994
995impl DerefMut for ConnectionPoolGuard<'_> {
996 fn deref_mut(&mut self) -> &mut Self::Target {
997 &mut self.guard
998 }
999}
1000
1001impl Drop for ConnectionPoolGuard<'_> {
1002 fn drop(&mut self) {
1003 #[cfg(test)]
1004 self.check_invariants();
1005 }
1006}
1007
1008fn broadcast<F>(
1009 sender_id: Option<ConnectionId>,
1010 receiver_ids: impl IntoIterator<Item = ConnectionId>,
1011 mut f: F,
1012) where
1013 F: FnMut(ConnectionId) -> anyhow::Result<()>,
1014{
1015 for receiver_id in receiver_ids {
1016 if Some(receiver_id) != sender_id {
1017 if let Err(error) = f(receiver_id) {
1018 tracing::error!("failed to send to {:?} {}", receiver_id, error);
1019 }
1020 }
1021 }
1022}
1023
1024pub struct ProtocolVersion(u32);
1025
1026impl Header for ProtocolVersion {
1027 fn name() -> &'static HeaderName {
1028 static ZED_PROTOCOL_VERSION: OnceLock<HeaderName> = OnceLock::new();
1029 ZED_PROTOCOL_VERSION.get_or_init(|| HeaderName::from_static("x-zed-protocol-version"))
1030 }
1031
1032 fn decode<'i, I>(values: &mut I) -> Result<Self, axum::headers::Error>
1033 where
1034 Self: Sized,
1035 I: Iterator<Item = &'i axum::http::HeaderValue>,
1036 {
1037 let version = values
1038 .next()
1039 .ok_or_else(axum::headers::Error::invalid)?
1040 .to_str()
1041 .map_err(|_| axum::headers::Error::invalid())?
1042 .parse()
1043 .map_err(|_| axum::headers::Error::invalid())?;
1044 Ok(Self(version))
1045 }
1046
1047 fn encode<E: Extend<axum::http::HeaderValue>>(&self, values: &mut E) {
1048 values.extend([self.0.to_string().parse().unwrap()]);
1049 }
1050}
1051
1052pub struct AppVersionHeader(SemanticVersion);
1053impl Header for AppVersionHeader {
1054 fn name() -> &'static HeaderName {
1055 static ZED_APP_VERSION: OnceLock<HeaderName> = OnceLock::new();
1056 ZED_APP_VERSION.get_or_init(|| HeaderName::from_static("x-zed-app-version"))
1057 }
1058
1059 fn decode<'i, I>(values: &mut I) -> Result<Self, axum::headers::Error>
1060 where
1061 Self: Sized,
1062 I: Iterator<Item = &'i axum::http::HeaderValue>,
1063 {
1064 let version = values
1065 .next()
1066 .ok_or_else(axum::headers::Error::invalid)?
1067 .to_str()
1068 .map_err(|_| axum::headers::Error::invalid())?
1069 .parse()
1070 .map_err(|_| axum::headers::Error::invalid())?;
1071 Ok(Self(version))
1072 }
1073
1074 fn encode<E: Extend<axum::http::HeaderValue>>(&self, values: &mut E) {
1075 values.extend([self.0.to_string().parse().unwrap()]);
1076 }
1077}
1078
1079pub fn routes(server: Arc<Server>) -> Router<(), Body> {
1080 Router::new()
1081 .route("/rpc", get(handle_websocket_request))
1082 .layer(
1083 ServiceBuilder::new()
1084 .layer(Extension(server.app_state.clone()))
1085 .layer(middleware::from_fn(auth::validate_header)),
1086 )
1087 .route("/metrics", get(handle_metrics))
1088 .layer(Extension(server))
1089}
1090
1091pub async fn handle_websocket_request(
1092 TypedHeader(ProtocolVersion(protocol_version)): TypedHeader<ProtocolVersion>,
1093 app_version_header: Option<TypedHeader<AppVersionHeader>>,
1094 ConnectInfo(socket_address): ConnectInfo<SocketAddr>,
1095 Extension(server): Extension<Arc<Server>>,
1096 Extension(principal): Extension<Principal>,
1097 country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,
1098 system_id_header: Option<TypedHeader<SystemIdHeader>>,
1099 ws: WebSocketUpgrade,
1100) -> axum::response::Response {
1101 if protocol_version != rpc::PROTOCOL_VERSION {
1102 return (
1103 StatusCode::UPGRADE_REQUIRED,
1104 "client must be upgraded".to_string(),
1105 )
1106 .into_response();
1107 }
1108
1109 let Some(version) = app_version_header.map(|header| ZedVersion(header.0.0)) else {
1110 return (
1111 StatusCode::UPGRADE_REQUIRED,
1112 "no version header found".to_string(),
1113 )
1114 .into_response();
1115 };
1116
1117 if !version.can_collaborate() {
1118 return (
1119 StatusCode::UPGRADE_REQUIRED,
1120 "client must be upgraded".to_string(),
1121 )
1122 .into_response();
1123 }
1124
1125 let socket_address = socket_address.to_string();
1126 ws.on_upgrade(move |socket| {
1127 let socket = socket
1128 .map_ok(to_tungstenite_message)
1129 .err_into()
1130 .with(|message| async move { to_axum_message(message) });
1131 let connection = Connection::new(Box::pin(socket));
1132 async move {
1133 server
1134 .handle_connection(
1135 connection,
1136 socket_address,
1137 principal,
1138 version,
1139 country_code_header.map(|header| header.to_string()),
1140 system_id_header.map(|header| header.to_string()),
1141 None,
1142 Executor::Production,
1143 )
1144 .await;
1145 }
1146 })
1147}
1148
1149pub async fn handle_metrics(Extension(server): Extension<Arc<Server>>) -> Result<String> {
1150 static CONNECTIONS_METRIC: OnceLock<IntGauge> = OnceLock::new();
1151 let connections_metric = CONNECTIONS_METRIC
1152 .get_or_init(|| register_int_gauge!("connections", "number of connections").unwrap());
1153
1154 let connections = server
1155 .connection_pool
1156 .lock()
1157 .connections()
1158 .filter(|connection| !connection.admin)
1159 .count();
1160 connections_metric.set(connections as _);
1161
1162 static SHARED_PROJECTS_METRIC: OnceLock<IntGauge> = OnceLock::new();
1163 let shared_projects_metric = SHARED_PROJECTS_METRIC.get_or_init(|| {
1164 register_int_gauge!(
1165 "shared_projects",
1166 "number of open projects with one or more guests"
1167 )
1168 .unwrap()
1169 });
1170
1171 let shared_projects = server.app_state.db.project_count_excluding_admins().await?;
1172 shared_projects_metric.set(shared_projects as _);
1173
1174 let encoder = prometheus::TextEncoder::new();
1175 let metric_families = prometheus::gather();
1176 let encoded_metrics = encoder
1177 .encode_to_string(&metric_families)
1178 .map_err(|err| anyhow!("{err}"))?;
1179 Ok(encoded_metrics)
1180}
1181
1182#[instrument(err, skip(executor))]
1183async fn connection_lost(
1184 session: Session,
1185 mut teardown: watch::Receiver<bool>,
1186 executor: Executor,
1187) -> Result<()> {
1188 session.peer.disconnect(session.connection_id);
1189 session
1190 .connection_pool()
1191 .await
1192 .remove_connection(session.connection_id)?;
1193
1194 session
1195 .db()
1196 .await
1197 .connection_lost(session.connection_id)
1198 .await
1199 .trace_err();
1200
1201 futures::select_biased! {
1202 _ = executor.sleep(RECONNECT_TIMEOUT).fuse() => {
1203
1204 log::info!("connection lost, removing all resources for user:{}, connection:{:?}", session.user_id(), session.connection_id);
1205 leave_room_for_session(&session, session.connection_id).await.trace_err();
1206 leave_channel_buffers_for_session(&session)
1207 .await
1208 .trace_err();
1209
1210 if !session
1211 .connection_pool()
1212 .await
1213 .is_user_online(session.user_id())
1214 {
1215 let db = session.db().await;
1216 if let Some(room) = db.decline_call(None, session.user_id()).await.trace_err().flatten() {
1217 room_updated(&room, &session.peer);
1218 }
1219 }
1220
1221 update_user_contacts(session.user_id(), &session).await?;
1222 },
1223 _ = teardown.changed().fuse() => {}
1224 }
1225
1226 Ok(())
1227}
1228
1229/// Acknowledges a ping from a client, used to keep the connection alive.
1230async fn ping(_: proto::Ping, response: Response<proto::Ping>, _session: Session) -> Result<()> {
1231 response.send(proto::Ack {})?;
1232 Ok(())
1233}
1234
1235/// Creates a new room for calling (outside of channels)
1236async fn create_room(
1237 _request: proto::CreateRoom,
1238 response: Response<proto::CreateRoom>,
1239 session: Session,
1240) -> Result<()> {
1241 let livekit_room = nanoid::nanoid!(30);
1242
1243 let live_kit_connection_info = util::maybe!(async {
1244 let live_kit = session.app_state.livekit_client.as_ref();
1245 let live_kit = live_kit?;
1246 let user_id = session.user_id().to_string();
1247
1248 let token = live_kit
1249 .room_token(&livekit_room, &user_id.to_string())
1250 .trace_err()?;
1251
1252 Some(proto::LiveKitConnectionInfo {
1253 server_url: live_kit.url().into(),
1254 token,
1255 can_publish: true,
1256 })
1257 })
1258 .await;
1259
1260 let room = session
1261 .db()
1262 .await
1263 .create_room(session.user_id(), session.connection_id, &livekit_room)
1264 .await?;
1265
1266 response.send(proto::CreateRoomResponse {
1267 room: Some(room.clone()),
1268 live_kit_connection_info,
1269 })?;
1270
1271 update_user_contacts(session.user_id(), &session).await?;
1272 Ok(())
1273}
1274
1275/// Join a room from an invitation. Equivalent to joining a channel if there is one.
1276async fn join_room(
1277 request: proto::JoinRoom,
1278 response: Response<proto::JoinRoom>,
1279 session: Session,
1280) -> Result<()> {
1281 let room_id = RoomId::from_proto(request.id);
1282
1283 let channel_id = session.db().await.channel_id_for_room(room_id).await?;
1284
1285 if let Some(channel_id) = channel_id {
1286 return join_channel_internal(channel_id, Box::new(response), session).await;
1287 }
1288
1289 let joined_room = {
1290 let room = session
1291 .db()
1292 .await
1293 .join_room(room_id, session.user_id(), session.connection_id)
1294 .await?;
1295 room_updated(&room.room, &session.peer);
1296 room.into_inner()
1297 };
1298
1299 for connection_id in session
1300 .connection_pool()
1301 .await
1302 .user_connection_ids(session.user_id())
1303 {
1304 session
1305 .peer
1306 .send(
1307 connection_id,
1308 proto::CallCanceled {
1309 room_id: room_id.to_proto(),
1310 },
1311 )
1312 .trace_err();
1313 }
1314
1315 let live_kit_connection_info = if let Some(live_kit) = session.app_state.livekit_client.as_ref()
1316 {
1317 live_kit
1318 .room_token(
1319 &joined_room.room.livekit_room,
1320 &session.user_id().to_string(),
1321 )
1322 .trace_err()
1323 .map(|token| proto::LiveKitConnectionInfo {
1324 server_url: live_kit.url().into(),
1325 token,
1326 can_publish: true,
1327 })
1328 } else {
1329 None
1330 };
1331
1332 response.send(proto::JoinRoomResponse {
1333 room: Some(joined_room.room),
1334 channel_id: None,
1335 live_kit_connection_info,
1336 })?;
1337
1338 update_user_contacts(session.user_id(), &session).await?;
1339 Ok(())
1340}
1341
1342/// Rejoin room is used to reconnect to a room after connection errors.
1343async fn rejoin_room(
1344 request: proto::RejoinRoom,
1345 response: Response<proto::RejoinRoom>,
1346 session: Session,
1347) -> Result<()> {
1348 let room;
1349 let channel;
1350 {
1351 let mut rejoined_room = session
1352 .db()
1353 .await
1354 .rejoin_room(request, session.user_id(), session.connection_id)
1355 .await?;
1356
1357 response.send(proto::RejoinRoomResponse {
1358 room: Some(rejoined_room.room.clone()),
1359 reshared_projects: rejoined_room
1360 .reshared_projects
1361 .iter()
1362 .map(|project| proto::ResharedProject {
1363 id: project.id.to_proto(),
1364 collaborators: project
1365 .collaborators
1366 .iter()
1367 .map(|collaborator| collaborator.to_proto())
1368 .collect(),
1369 })
1370 .collect(),
1371 rejoined_projects: rejoined_room
1372 .rejoined_projects
1373 .iter()
1374 .map(|rejoined_project| rejoined_project.to_proto())
1375 .collect(),
1376 })?;
1377 room_updated(&rejoined_room.room, &session.peer);
1378
1379 for project in &rejoined_room.reshared_projects {
1380 for collaborator in &project.collaborators {
1381 session
1382 .peer
1383 .send(
1384 collaborator.connection_id,
1385 proto::UpdateProjectCollaborator {
1386 project_id: project.id.to_proto(),
1387 old_peer_id: Some(project.old_connection_id.into()),
1388 new_peer_id: Some(session.connection_id.into()),
1389 },
1390 )
1391 .trace_err();
1392 }
1393
1394 broadcast(
1395 Some(session.connection_id),
1396 project
1397 .collaborators
1398 .iter()
1399 .map(|collaborator| collaborator.connection_id),
1400 |connection_id| {
1401 session.peer.forward_send(
1402 session.connection_id,
1403 connection_id,
1404 proto::UpdateProject {
1405 project_id: project.id.to_proto(),
1406 worktrees: project.worktrees.clone(),
1407 },
1408 )
1409 },
1410 );
1411 }
1412
1413 notify_rejoined_projects(&mut rejoined_room.rejoined_projects, &session)?;
1414
1415 let rejoined_room = rejoined_room.into_inner();
1416
1417 room = rejoined_room.room;
1418 channel = rejoined_room.channel;
1419 }
1420
1421 if let Some(channel) = channel {
1422 channel_updated(
1423 &channel,
1424 &room,
1425 &session.peer,
1426 &*session.connection_pool().await,
1427 );
1428 }
1429
1430 update_user_contacts(session.user_id(), &session).await?;
1431 Ok(())
1432}
1433
1434fn notify_rejoined_projects(
1435 rejoined_projects: &mut Vec<RejoinedProject>,
1436 session: &Session,
1437) -> Result<()> {
1438 for project in rejoined_projects.iter() {
1439 for collaborator in &project.collaborators {
1440 session
1441 .peer
1442 .send(
1443 collaborator.connection_id,
1444 proto::UpdateProjectCollaborator {
1445 project_id: project.id.to_proto(),
1446 old_peer_id: Some(project.old_connection_id.into()),
1447 new_peer_id: Some(session.connection_id.into()),
1448 },
1449 )
1450 .trace_err();
1451 }
1452 }
1453
1454 for project in rejoined_projects {
1455 for worktree in mem::take(&mut project.worktrees) {
1456 // Stream this worktree's entries.
1457 let message = proto::UpdateWorktree {
1458 project_id: project.id.to_proto(),
1459 worktree_id: worktree.id,
1460 abs_path: worktree.abs_path.clone(),
1461 root_name: worktree.root_name,
1462 updated_entries: worktree.updated_entries,
1463 removed_entries: worktree.removed_entries,
1464 scan_id: worktree.scan_id,
1465 is_last_update: worktree.completed_scan_id == worktree.scan_id,
1466 updated_repositories: worktree.updated_repositories,
1467 removed_repositories: worktree.removed_repositories,
1468 };
1469 for update in proto::split_worktree_update(message) {
1470 session.peer.send(session.connection_id, update)?;
1471 }
1472
1473 // Stream this worktree's diagnostics.
1474 for summary in worktree.diagnostic_summaries {
1475 session.peer.send(
1476 session.connection_id,
1477 proto::UpdateDiagnosticSummary {
1478 project_id: project.id.to_proto(),
1479 worktree_id: worktree.id,
1480 summary: Some(summary),
1481 },
1482 )?;
1483 }
1484
1485 for settings_file in worktree.settings_files {
1486 session.peer.send(
1487 session.connection_id,
1488 proto::UpdateWorktreeSettings {
1489 project_id: project.id.to_proto(),
1490 worktree_id: worktree.id,
1491 path: settings_file.path,
1492 content: Some(settings_file.content),
1493 kind: Some(settings_file.kind.to_proto().into()),
1494 },
1495 )?;
1496 }
1497 }
1498
1499 for repository in mem::take(&mut project.updated_repositories) {
1500 for update in split_repository_update(repository) {
1501 session.peer.send(session.connection_id, update)?;
1502 }
1503 }
1504
1505 for id in mem::take(&mut project.removed_repositories) {
1506 session.peer.send(
1507 session.connection_id,
1508 proto::RemoveRepository {
1509 project_id: project.id.to_proto(),
1510 id,
1511 },
1512 )?;
1513 }
1514 }
1515
1516 Ok(())
1517}
1518
1519/// leave room disconnects from the room.
1520async fn leave_room(
1521 _: proto::LeaveRoom,
1522 response: Response<proto::LeaveRoom>,
1523 session: Session,
1524) -> Result<()> {
1525 leave_room_for_session(&session, session.connection_id).await?;
1526 response.send(proto::Ack {})?;
1527 Ok(())
1528}
1529
1530/// Updates the permissions of someone else in the room.
1531async fn set_room_participant_role(
1532 request: proto::SetRoomParticipantRole,
1533 response: Response<proto::SetRoomParticipantRole>,
1534 session: Session,
1535) -> Result<()> {
1536 let user_id = UserId::from_proto(request.user_id);
1537 let role = ChannelRole::from(request.role());
1538
1539 let (livekit_room, can_publish) = {
1540 let room = session
1541 .db()
1542 .await
1543 .set_room_participant_role(
1544 session.user_id(),
1545 RoomId::from_proto(request.room_id),
1546 user_id,
1547 role,
1548 )
1549 .await?;
1550
1551 let livekit_room = room.livekit_room.clone();
1552 let can_publish = ChannelRole::from(request.role()).can_use_microphone();
1553 room_updated(&room, &session.peer);
1554 (livekit_room, can_publish)
1555 };
1556
1557 if let Some(live_kit) = session.app_state.livekit_client.as_ref() {
1558 live_kit
1559 .update_participant(
1560 livekit_room.clone(),
1561 request.user_id.to_string(),
1562 livekit_api::proto::ParticipantPermission {
1563 can_subscribe: true,
1564 can_publish,
1565 can_publish_data: can_publish,
1566 hidden: false,
1567 recorder: false,
1568 },
1569 )
1570 .await
1571 .trace_err();
1572 }
1573
1574 response.send(proto::Ack {})?;
1575 Ok(())
1576}
1577
1578/// Call someone else into the current room
1579async fn call(
1580 request: proto::Call,
1581 response: Response<proto::Call>,
1582 session: Session,
1583) -> Result<()> {
1584 let room_id = RoomId::from_proto(request.room_id);
1585 let calling_user_id = session.user_id();
1586 let calling_connection_id = session.connection_id;
1587 let called_user_id = UserId::from_proto(request.called_user_id);
1588 let initial_project_id = request.initial_project_id.map(ProjectId::from_proto);
1589 if !session
1590 .db()
1591 .await
1592 .has_contact(calling_user_id, called_user_id)
1593 .await?
1594 {
1595 return Err(anyhow!("cannot call a user who isn't a contact"))?;
1596 }
1597
1598 let incoming_call = {
1599 let (room, incoming_call) = &mut *session
1600 .db()
1601 .await
1602 .call(
1603 room_id,
1604 calling_user_id,
1605 calling_connection_id,
1606 called_user_id,
1607 initial_project_id,
1608 )
1609 .await?;
1610 room_updated(room, &session.peer);
1611 mem::take(incoming_call)
1612 };
1613 update_user_contacts(called_user_id, &session).await?;
1614
1615 let mut calls = session
1616 .connection_pool()
1617 .await
1618 .user_connection_ids(called_user_id)
1619 .map(|connection_id| session.peer.request(connection_id, incoming_call.clone()))
1620 .collect::<FuturesUnordered<_>>();
1621
1622 while let Some(call_response) = calls.next().await {
1623 match call_response.as_ref() {
1624 Ok(_) => {
1625 response.send(proto::Ack {})?;
1626 return Ok(());
1627 }
1628 Err(_) => {
1629 call_response.trace_err();
1630 }
1631 }
1632 }
1633
1634 {
1635 let room = session
1636 .db()
1637 .await
1638 .call_failed(room_id, called_user_id)
1639 .await?;
1640 room_updated(&room, &session.peer);
1641 }
1642 update_user_contacts(called_user_id, &session).await?;
1643
1644 Err(anyhow!("failed to ring user"))?
1645}
1646
1647/// Cancel an outgoing call.
1648async fn cancel_call(
1649 request: proto::CancelCall,
1650 response: Response<proto::CancelCall>,
1651 session: Session,
1652) -> Result<()> {
1653 let called_user_id = UserId::from_proto(request.called_user_id);
1654 let room_id = RoomId::from_proto(request.room_id);
1655 {
1656 let room = session
1657 .db()
1658 .await
1659 .cancel_call(room_id, session.connection_id, called_user_id)
1660 .await?;
1661 room_updated(&room, &session.peer);
1662 }
1663
1664 for connection_id in session
1665 .connection_pool()
1666 .await
1667 .user_connection_ids(called_user_id)
1668 {
1669 session
1670 .peer
1671 .send(
1672 connection_id,
1673 proto::CallCanceled {
1674 room_id: room_id.to_proto(),
1675 },
1676 )
1677 .trace_err();
1678 }
1679 response.send(proto::Ack {})?;
1680
1681 update_user_contacts(called_user_id, &session).await?;
1682 Ok(())
1683}
1684
1685/// Decline an incoming call.
1686async fn decline_call(message: proto::DeclineCall, session: Session) -> Result<()> {
1687 let room_id = RoomId::from_proto(message.room_id);
1688 {
1689 let room = session
1690 .db()
1691 .await
1692 .decline_call(Some(room_id), session.user_id())
1693 .await?
1694 .context("declining call")?;
1695 room_updated(&room, &session.peer);
1696 }
1697
1698 for connection_id in session
1699 .connection_pool()
1700 .await
1701 .user_connection_ids(session.user_id())
1702 {
1703 session
1704 .peer
1705 .send(
1706 connection_id,
1707 proto::CallCanceled {
1708 room_id: room_id.to_proto(),
1709 },
1710 )
1711 .trace_err();
1712 }
1713 update_user_contacts(session.user_id(), &session).await?;
1714 Ok(())
1715}
1716
1717/// Updates other participants in the room with your current location.
1718async fn update_participant_location(
1719 request: proto::UpdateParticipantLocation,
1720 response: Response<proto::UpdateParticipantLocation>,
1721 session: Session,
1722) -> Result<()> {
1723 let room_id = RoomId::from_proto(request.room_id);
1724 let location = request.location.context("invalid location")?;
1725
1726 let db = session.db().await;
1727 let room = db
1728 .update_room_participant_location(room_id, session.connection_id, location)
1729 .await?;
1730
1731 room_updated(&room, &session.peer);
1732 response.send(proto::Ack {})?;
1733 Ok(())
1734}
1735
1736/// Share a project into the room.
1737async fn share_project(
1738 request: proto::ShareProject,
1739 response: Response<proto::ShareProject>,
1740 session: Session,
1741) -> Result<()> {
1742 let (project_id, room) = &*session
1743 .db()
1744 .await
1745 .share_project(
1746 RoomId::from_proto(request.room_id),
1747 session.connection_id,
1748 &request.worktrees,
1749 request.is_ssh_project,
1750 )
1751 .await?;
1752 response.send(proto::ShareProjectResponse {
1753 project_id: project_id.to_proto(),
1754 })?;
1755 room_updated(room, &session.peer);
1756
1757 Ok(())
1758}
1759
1760/// Unshare a project from the room.
1761async fn unshare_project(message: proto::UnshareProject, session: Session) -> Result<()> {
1762 let project_id = ProjectId::from_proto(message.project_id);
1763 unshare_project_internal(project_id, session.connection_id, &session).await
1764}
1765
1766async fn unshare_project_internal(
1767 project_id: ProjectId,
1768 connection_id: ConnectionId,
1769 session: &Session,
1770) -> Result<()> {
1771 let delete = {
1772 let room_guard = session
1773 .db()
1774 .await
1775 .unshare_project(project_id, connection_id)
1776 .await?;
1777
1778 let (delete, room, guest_connection_ids) = &*room_guard;
1779
1780 let message = proto::UnshareProject {
1781 project_id: project_id.to_proto(),
1782 };
1783
1784 broadcast(
1785 Some(connection_id),
1786 guest_connection_ids.iter().copied(),
1787 |conn_id| session.peer.send(conn_id, message.clone()),
1788 );
1789 if let Some(room) = room {
1790 room_updated(room, &session.peer);
1791 }
1792
1793 *delete
1794 };
1795
1796 if delete {
1797 let db = session.db().await;
1798 db.delete_project(project_id).await?;
1799 }
1800
1801 Ok(())
1802}
1803
1804/// Join someone elses shared project.
1805async fn join_project(
1806 request: proto::JoinProject,
1807 response: Response<proto::JoinProject>,
1808 session: Session,
1809) -> Result<()> {
1810 let project_id = ProjectId::from_proto(request.project_id);
1811
1812 tracing::info!(%project_id, "join project");
1813
1814 let db = session.db().await;
1815 let (project, replica_id) = &mut *db
1816 .join_project(project_id, session.connection_id, session.user_id())
1817 .await?;
1818 drop(db);
1819 tracing::info!(%project_id, "join remote project");
1820 join_project_internal(response, session, project, replica_id)
1821}
1822
1823trait JoinProjectInternalResponse {
1824 fn send(self, result: proto::JoinProjectResponse) -> Result<()>;
1825}
1826impl JoinProjectInternalResponse for Response<proto::JoinProject> {
1827 fn send(self, result: proto::JoinProjectResponse) -> Result<()> {
1828 Response::<proto::JoinProject>::send(self, result)
1829 }
1830}
1831
1832fn join_project_internal(
1833 response: impl JoinProjectInternalResponse,
1834 session: Session,
1835 project: &mut Project,
1836 replica_id: &ReplicaId,
1837) -> Result<()> {
1838 let collaborators = project
1839 .collaborators
1840 .iter()
1841 .filter(|collaborator| collaborator.connection_id != session.connection_id)
1842 .map(|collaborator| collaborator.to_proto())
1843 .collect::<Vec<_>>();
1844 let project_id = project.id;
1845 let guest_user_id = session.user_id();
1846
1847 let worktrees = project
1848 .worktrees
1849 .iter()
1850 .map(|(id, worktree)| proto::WorktreeMetadata {
1851 id: *id,
1852 root_name: worktree.root_name.clone(),
1853 visible: worktree.visible,
1854 abs_path: worktree.abs_path.clone(),
1855 })
1856 .collect::<Vec<_>>();
1857
1858 let add_project_collaborator = proto::AddProjectCollaborator {
1859 project_id: project_id.to_proto(),
1860 collaborator: Some(proto::Collaborator {
1861 peer_id: Some(session.connection_id.into()),
1862 replica_id: replica_id.0 as u32,
1863 user_id: guest_user_id.to_proto(),
1864 is_host: false,
1865 }),
1866 };
1867
1868 for collaborator in &collaborators {
1869 session
1870 .peer
1871 .send(
1872 collaborator.peer_id.unwrap().into(),
1873 add_project_collaborator.clone(),
1874 )
1875 .trace_err();
1876 }
1877
1878 // First, we send the metadata associated with each worktree.
1879 response.send(proto::JoinProjectResponse {
1880 project_id: project.id.0 as u64,
1881 worktrees: worktrees.clone(),
1882 replica_id: replica_id.0 as u32,
1883 collaborators: collaborators.clone(),
1884 language_servers: project.language_servers.clone(),
1885 role: project.role.into(),
1886 })?;
1887
1888 for (worktree_id, worktree) in mem::take(&mut project.worktrees) {
1889 // Stream this worktree's entries.
1890 let message = proto::UpdateWorktree {
1891 project_id: project_id.to_proto(),
1892 worktree_id,
1893 abs_path: worktree.abs_path.clone(),
1894 root_name: worktree.root_name,
1895 updated_entries: worktree.entries,
1896 removed_entries: Default::default(),
1897 scan_id: worktree.scan_id,
1898 is_last_update: worktree.scan_id == worktree.completed_scan_id,
1899 updated_repositories: worktree.legacy_repository_entries.into_values().collect(),
1900 removed_repositories: Default::default(),
1901 };
1902 for update in proto::split_worktree_update(message) {
1903 session.peer.send(session.connection_id, update.clone())?;
1904 }
1905
1906 // Stream this worktree's diagnostics.
1907 for summary in worktree.diagnostic_summaries {
1908 session.peer.send(
1909 session.connection_id,
1910 proto::UpdateDiagnosticSummary {
1911 project_id: project_id.to_proto(),
1912 worktree_id: worktree.id,
1913 summary: Some(summary),
1914 },
1915 )?;
1916 }
1917
1918 for settings_file in worktree.settings_files {
1919 session.peer.send(
1920 session.connection_id,
1921 proto::UpdateWorktreeSettings {
1922 project_id: project_id.to_proto(),
1923 worktree_id: worktree.id,
1924 path: settings_file.path,
1925 content: Some(settings_file.content),
1926 kind: Some(settings_file.kind.to_proto() as i32),
1927 },
1928 )?;
1929 }
1930 }
1931
1932 for repository in mem::take(&mut project.repositories) {
1933 for update in split_repository_update(repository) {
1934 session.peer.send(session.connection_id, update)?;
1935 }
1936 }
1937
1938 for language_server in &project.language_servers {
1939 session.peer.send(
1940 session.connection_id,
1941 proto::UpdateLanguageServer {
1942 project_id: project_id.to_proto(),
1943 language_server_id: language_server.id,
1944 variant: Some(
1945 proto::update_language_server::Variant::DiskBasedDiagnosticsUpdated(
1946 proto::LspDiskBasedDiagnosticsUpdated {},
1947 ),
1948 ),
1949 },
1950 )?;
1951 }
1952
1953 Ok(())
1954}
1955
1956/// Leave someone elses shared project.
1957async fn leave_project(request: proto::LeaveProject, session: Session) -> Result<()> {
1958 let sender_id = session.connection_id;
1959 let project_id = ProjectId::from_proto(request.project_id);
1960 let db = session.db().await;
1961
1962 let (room, project) = &*db.leave_project(project_id, sender_id).await?;
1963 tracing::info!(
1964 %project_id,
1965 "leave project"
1966 );
1967
1968 project_left(project, &session);
1969 if let Some(room) = room {
1970 room_updated(room, &session.peer);
1971 }
1972
1973 Ok(())
1974}
1975
1976/// Updates other participants with changes to the project
1977async fn update_project(
1978 request: proto::UpdateProject,
1979 response: Response<proto::UpdateProject>,
1980 session: Session,
1981) -> Result<()> {
1982 let project_id = ProjectId::from_proto(request.project_id);
1983 let (room, guest_connection_ids) = &*session
1984 .db()
1985 .await
1986 .update_project(project_id, session.connection_id, &request.worktrees)
1987 .await?;
1988 broadcast(
1989 Some(session.connection_id),
1990 guest_connection_ids.iter().copied(),
1991 |connection_id| {
1992 session
1993 .peer
1994 .forward_send(session.connection_id, connection_id, request.clone())
1995 },
1996 );
1997 if let Some(room) = room {
1998 room_updated(room, &session.peer);
1999 }
2000 response.send(proto::Ack {})?;
2001
2002 Ok(())
2003}
2004
2005/// Updates other participants with changes to the worktree
2006async fn update_worktree(
2007 request: proto::UpdateWorktree,
2008 response: Response<proto::UpdateWorktree>,
2009 session: Session,
2010) -> Result<()> {
2011 let guest_connection_ids = session
2012 .db()
2013 .await
2014 .update_worktree(&request, session.connection_id)
2015 .await?;
2016
2017 broadcast(
2018 Some(session.connection_id),
2019 guest_connection_ids.iter().copied(),
2020 |connection_id| {
2021 session
2022 .peer
2023 .forward_send(session.connection_id, connection_id, request.clone())
2024 },
2025 );
2026 response.send(proto::Ack {})?;
2027 Ok(())
2028}
2029
2030async fn update_repository(
2031 request: proto::UpdateRepository,
2032 response: Response<proto::UpdateRepository>,
2033 session: Session,
2034) -> Result<()> {
2035 let guest_connection_ids = session
2036 .db()
2037 .await
2038 .update_repository(&request, session.connection_id)
2039 .await?;
2040
2041 broadcast(
2042 Some(session.connection_id),
2043 guest_connection_ids.iter().copied(),
2044 |connection_id| {
2045 session
2046 .peer
2047 .forward_send(session.connection_id, connection_id, request.clone())
2048 },
2049 );
2050 response.send(proto::Ack {})?;
2051 Ok(())
2052}
2053
2054async fn remove_repository(
2055 request: proto::RemoveRepository,
2056 response: Response<proto::RemoveRepository>,
2057 session: Session,
2058) -> Result<()> {
2059 let guest_connection_ids = session
2060 .db()
2061 .await
2062 .remove_repository(&request, session.connection_id)
2063 .await?;
2064
2065 broadcast(
2066 Some(session.connection_id),
2067 guest_connection_ids.iter().copied(),
2068 |connection_id| {
2069 session
2070 .peer
2071 .forward_send(session.connection_id, connection_id, request.clone())
2072 },
2073 );
2074 response.send(proto::Ack {})?;
2075 Ok(())
2076}
2077
2078/// Updates other participants with changes to the diagnostics
2079async fn update_diagnostic_summary(
2080 message: proto::UpdateDiagnosticSummary,
2081 session: Session,
2082) -> Result<()> {
2083 let guest_connection_ids = session
2084 .db()
2085 .await
2086 .update_diagnostic_summary(&message, session.connection_id)
2087 .await?;
2088
2089 broadcast(
2090 Some(session.connection_id),
2091 guest_connection_ids.iter().copied(),
2092 |connection_id| {
2093 session
2094 .peer
2095 .forward_send(session.connection_id, connection_id, message.clone())
2096 },
2097 );
2098
2099 Ok(())
2100}
2101
2102/// Updates other participants with changes to the worktree settings
2103async fn update_worktree_settings(
2104 message: proto::UpdateWorktreeSettings,
2105 session: Session,
2106) -> Result<()> {
2107 let guest_connection_ids = session
2108 .db()
2109 .await
2110 .update_worktree_settings(&message, session.connection_id)
2111 .await?;
2112
2113 broadcast(
2114 Some(session.connection_id),
2115 guest_connection_ids.iter().copied(),
2116 |connection_id| {
2117 session
2118 .peer
2119 .forward_send(session.connection_id, connection_id, message.clone())
2120 },
2121 );
2122
2123 Ok(())
2124}
2125
2126/// Notify other participants that a language server has started.
2127async fn start_language_server(
2128 request: proto::StartLanguageServer,
2129 session: Session,
2130) -> Result<()> {
2131 let guest_connection_ids = session
2132 .db()
2133 .await
2134 .start_language_server(&request, session.connection_id)
2135 .await?;
2136
2137 broadcast(
2138 Some(session.connection_id),
2139 guest_connection_ids.iter().copied(),
2140 |connection_id| {
2141 session
2142 .peer
2143 .forward_send(session.connection_id, connection_id, request.clone())
2144 },
2145 );
2146 Ok(())
2147}
2148
2149/// Notify other participants that a language server has changed.
2150async fn update_language_server(
2151 request: proto::UpdateLanguageServer,
2152 session: Session,
2153) -> Result<()> {
2154 let project_id = ProjectId::from_proto(request.project_id);
2155 let project_connection_ids = session
2156 .db()
2157 .await
2158 .project_connection_ids(project_id, session.connection_id, true)
2159 .await?;
2160 broadcast(
2161 Some(session.connection_id),
2162 project_connection_ids.iter().copied(),
2163 |connection_id| {
2164 session
2165 .peer
2166 .forward_send(session.connection_id, connection_id, request.clone())
2167 },
2168 );
2169 Ok(())
2170}
2171
2172/// forward a project request to the host. These requests should be read only
2173/// as guests are allowed to send them.
2174async fn forward_read_only_project_request<T>(
2175 request: T,
2176 response: Response<T>,
2177 session: Session,
2178) -> Result<()>
2179where
2180 T: EntityMessage + RequestMessage,
2181{
2182 let project_id = ProjectId::from_proto(request.remote_entity_id());
2183 let host_connection_id = session
2184 .db()
2185 .await
2186 .host_for_read_only_project_request(project_id, session.connection_id)
2187 .await?;
2188 let payload = session
2189 .peer
2190 .forward_request(session.connection_id, host_connection_id, request)
2191 .await?;
2192 response.send(payload)?;
2193 Ok(())
2194}
2195
2196async fn forward_find_search_candidates_request(
2197 request: proto::FindSearchCandidates,
2198 response: Response<proto::FindSearchCandidates>,
2199 session: Session,
2200) -> Result<()> {
2201 let project_id = ProjectId::from_proto(request.remote_entity_id());
2202 let host_connection_id = session
2203 .db()
2204 .await
2205 .host_for_read_only_project_request(project_id, session.connection_id)
2206 .await?;
2207 let payload = session
2208 .peer
2209 .forward_request(session.connection_id, host_connection_id, request)
2210 .await?;
2211 response.send(payload)?;
2212 Ok(())
2213}
2214
2215/// forward a project request to the host. These requests are disallowed
2216/// for guests.
2217async fn forward_mutating_project_request<T>(
2218 request: T,
2219 response: Response<T>,
2220 session: Session,
2221) -> Result<()>
2222where
2223 T: EntityMessage + RequestMessage,
2224{
2225 let project_id = ProjectId::from_proto(request.remote_entity_id());
2226
2227 let host_connection_id = session
2228 .db()
2229 .await
2230 .host_for_mutating_project_request(project_id, session.connection_id)
2231 .await?;
2232 let payload = session
2233 .peer
2234 .forward_request(session.connection_id, host_connection_id, request)
2235 .await?;
2236 response.send(payload)?;
2237 Ok(())
2238}
2239
2240/// Notify other participants that a new buffer has been created
2241async fn create_buffer_for_peer(
2242 request: proto::CreateBufferForPeer,
2243 session: Session,
2244) -> Result<()> {
2245 session
2246 .db()
2247 .await
2248 .check_user_is_project_host(
2249 ProjectId::from_proto(request.project_id),
2250 session.connection_id,
2251 )
2252 .await?;
2253 let peer_id = request.peer_id.context("invalid peer id")?;
2254 session
2255 .peer
2256 .forward_send(session.connection_id, peer_id.into(), request)?;
2257 Ok(())
2258}
2259
2260/// Notify other participants that a buffer has been updated. This is
2261/// allowed for guests as long as the update is limited to selections.
2262async fn update_buffer(
2263 request: proto::UpdateBuffer,
2264 response: Response<proto::UpdateBuffer>,
2265 session: Session,
2266) -> Result<()> {
2267 let project_id = ProjectId::from_proto(request.project_id);
2268 let mut capability = Capability::ReadOnly;
2269
2270 for op in request.operations.iter() {
2271 match op.variant {
2272 None | Some(proto::operation::Variant::UpdateSelections(_)) => {}
2273 Some(_) => capability = Capability::ReadWrite,
2274 }
2275 }
2276
2277 let host = {
2278 let guard = session
2279 .db()
2280 .await
2281 .connections_for_buffer_update(project_id, session.connection_id, capability)
2282 .await?;
2283
2284 let (host, guests) = &*guard;
2285
2286 broadcast(
2287 Some(session.connection_id),
2288 guests.clone(),
2289 |connection_id| {
2290 session
2291 .peer
2292 .forward_send(session.connection_id, connection_id, request.clone())
2293 },
2294 );
2295
2296 *host
2297 };
2298
2299 if host != session.connection_id {
2300 session
2301 .peer
2302 .forward_request(session.connection_id, host, request.clone())
2303 .await?;
2304 }
2305
2306 response.send(proto::Ack {})?;
2307 Ok(())
2308}
2309
2310async fn update_context(message: proto::UpdateContext, session: Session) -> Result<()> {
2311 let project_id = ProjectId::from_proto(message.project_id);
2312
2313 let operation = message.operation.as_ref().context("invalid operation")?;
2314 let capability = match operation.variant.as_ref() {
2315 Some(proto::context_operation::Variant::BufferOperation(buffer_op)) => {
2316 if let Some(buffer_op) = buffer_op.operation.as_ref() {
2317 match buffer_op.variant {
2318 None | Some(proto::operation::Variant::UpdateSelections(_)) => {
2319 Capability::ReadOnly
2320 }
2321 _ => Capability::ReadWrite,
2322 }
2323 } else {
2324 Capability::ReadWrite
2325 }
2326 }
2327 Some(_) => Capability::ReadWrite,
2328 None => Capability::ReadOnly,
2329 };
2330
2331 let guard = session
2332 .db()
2333 .await
2334 .connections_for_buffer_update(project_id, session.connection_id, capability)
2335 .await?;
2336
2337 let (host, guests) = &*guard;
2338
2339 broadcast(
2340 Some(session.connection_id),
2341 guests.iter().chain([host]).copied(),
2342 |connection_id| {
2343 session
2344 .peer
2345 .forward_send(session.connection_id, connection_id, message.clone())
2346 },
2347 );
2348
2349 Ok(())
2350}
2351
2352/// Notify other participants that a project has been updated.
2353async fn broadcast_project_message_from_host<T: EntityMessage<Entity = ShareProject>>(
2354 request: T,
2355 session: Session,
2356) -> Result<()> {
2357 let project_id = ProjectId::from_proto(request.remote_entity_id());
2358 let project_connection_ids = session
2359 .db()
2360 .await
2361 .project_connection_ids(project_id, session.connection_id, false)
2362 .await?;
2363
2364 broadcast(
2365 Some(session.connection_id),
2366 project_connection_ids.iter().copied(),
2367 |connection_id| {
2368 session
2369 .peer
2370 .forward_send(session.connection_id, connection_id, request.clone())
2371 },
2372 );
2373 Ok(())
2374}
2375
2376/// Start following another user in a call.
2377async fn follow(
2378 request: proto::Follow,
2379 response: Response<proto::Follow>,
2380 session: Session,
2381) -> Result<()> {
2382 let room_id = RoomId::from_proto(request.room_id);
2383 let project_id = request.project_id.map(ProjectId::from_proto);
2384 let leader_id = request.leader_id.context("invalid leader id")?.into();
2385 let follower_id = session.connection_id;
2386
2387 session
2388 .db()
2389 .await
2390 .check_room_participants(room_id, leader_id, session.connection_id)
2391 .await?;
2392
2393 let response_payload = session
2394 .peer
2395 .forward_request(session.connection_id, leader_id, request)
2396 .await?;
2397 response.send(response_payload)?;
2398
2399 if let Some(project_id) = project_id {
2400 let room = session
2401 .db()
2402 .await
2403 .follow(room_id, project_id, leader_id, follower_id)
2404 .await?;
2405 room_updated(&room, &session.peer);
2406 }
2407
2408 Ok(())
2409}
2410
2411/// Stop following another user in a call.
2412async fn unfollow(request: proto::Unfollow, session: Session) -> Result<()> {
2413 let room_id = RoomId::from_proto(request.room_id);
2414 let project_id = request.project_id.map(ProjectId::from_proto);
2415 let leader_id = request.leader_id.context("invalid leader id")?.into();
2416 let follower_id = session.connection_id;
2417
2418 session
2419 .db()
2420 .await
2421 .check_room_participants(room_id, leader_id, session.connection_id)
2422 .await?;
2423
2424 session
2425 .peer
2426 .forward_send(session.connection_id, leader_id, request)?;
2427
2428 if let Some(project_id) = project_id {
2429 let room = session
2430 .db()
2431 .await
2432 .unfollow(room_id, project_id, leader_id, follower_id)
2433 .await?;
2434 room_updated(&room, &session.peer);
2435 }
2436
2437 Ok(())
2438}
2439
2440/// Notify everyone following you of your current location.
2441async fn update_followers(request: proto::UpdateFollowers, session: Session) -> Result<()> {
2442 let room_id = RoomId::from_proto(request.room_id);
2443 let database = session.db.lock().await;
2444
2445 let connection_ids = if let Some(project_id) = request.project_id {
2446 let project_id = ProjectId::from_proto(project_id);
2447 database
2448 .project_connection_ids(project_id, session.connection_id, true)
2449 .await?
2450 } else {
2451 database
2452 .room_connection_ids(room_id, session.connection_id)
2453 .await?
2454 };
2455
2456 // For now, don't send view update messages back to that view's current leader.
2457 let peer_id_to_omit = request.variant.as_ref().and_then(|variant| match variant {
2458 proto::update_followers::Variant::UpdateView(payload) => payload.leader_id,
2459 _ => None,
2460 });
2461
2462 for connection_id in connection_ids.iter().cloned() {
2463 if Some(connection_id.into()) != peer_id_to_omit && connection_id != session.connection_id {
2464 session
2465 .peer
2466 .forward_send(session.connection_id, connection_id, request.clone())?;
2467 }
2468 }
2469 Ok(())
2470}
2471
2472/// Get public data about users.
2473async fn get_users(
2474 request: proto::GetUsers,
2475 response: Response<proto::GetUsers>,
2476 session: Session,
2477) -> Result<()> {
2478 let user_ids = request
2479 .user_ids
2480 .into_iter()
2481 .map(UserId::from_proto)
2482 .collect();
2483 let users = session
2484 .db()
2485 .await
2486 .get_users_by_ids(user_ids)
2487 .await?
2488 .into_iter()
2489 .map(|user| proto::User {
2490 id: user.id.to_proto(),
2491 avatar_url: format!("https://github.com/{}.png?size=128", user.github_login),
2492 github_login: user.github_login,
2493 email: user.email_address,
2494 name: user.name,
2495 })
2496 .collect();
2497 response.send(proto::UsersResponse { users })?;
2498 Ok(())
2499}
2500
2501/// Search for users (to invite) buy Github login
2502async fn fuzzy_search_users(
2503 request: proto::FuzzySearchUsers,
2504 response: Response<proto::FuzzySearchUsers>,
2505 session: Session,
2506) -> Result<()> {
2507 let query = request.query;
2508 let users = match query.len() {
2509 0 => vec![],
2510 1 | 2 => session
2511 .db()
2512 .await
2513 .get_user_by_github_login(&query)
2514 .await?
2515 .into_iter()
2516 .collect(),
2517 _ => session.db().await.fuzzy_search_users(&query, 10).await?,
2518 };
2519 let users = users
2520 .into_iter()
2521 .filter(|user| user.id != session.user_id())
2522 .map(|user| proto::User {
2523 id: user.id.to_proto(),
2524 avatar_url: format!("https://github.com/{}.png?size=128", user.github_login),
2525 github_login: user.github_login,
2526 name: user.name,
2527 email: user.email_address,
2528 })
2529 .collect();
2530 response.send(proto::UsersResponse { users })?;
2531 Ok(())
2532}
2533
2534/// Send a contact request to another user.
2535async fn request_contact(
2536 request: proto::RequestContact,
2537 response: Response<proto::RequestContact>,
2538 session: Session,
2539) -> Result<()> {
2540 let requester_id = session.user_id();
2541 let responder_id = UserId::from_proto(request.responder_id);
2542 if requester_id == responder_id {
2543 return Err(anyhow!("cannot add yourself as a contact"))?;
2544 }
2545
2546 let notifications = session
2547 .db()
2548 .await
2549 .send_contact_request(requester_id, responder_id)
2550 .await?;
2551
2552 // Update outgoing contact requests of requester
2553 let mut update = proto::UpdateContacts::default();
2554 update.outgoing_requests.push(responder_id.to_proto());
2555 for connection_id in session
2556 .connection_pool()
2557 .await
2558 .user_connection_ids(requester_id)
2559 {
2560 session.peer.send(connection_id, update.clone())?;
2561 }
2562
2563 // Update incoming contact requests of responder
2564 let mut update = proto::UpdateContacts::default();
2565 update
2566 .incoming_requests
2567 .push(proto::IncomingContactRequest {
2568 requester_id: requester_id.to_proto(),
2569 });
2570 let connection_pool = session.connection_pool().await;
2571 for connection_id in connection_pool.user_connection_ids(responder_id) {
2572 session.peer.send(connection_id, update.clone())?;
2573 }
2574
2575 send_notifications(&connection_pool, &session.peer, notifications);
2576
2577 response.send(proto::Ack {})?;
2578 Ok(())
2579}
2580
2581/// Accept or decline a contact request
2582async fn respond_to_contact_request(
2583 request: proto::RespondToContactRequest,
2584 response: Response<proto::RespondToContactRequest>,
2585 session: Session,
2586) -> Result<()> {
2587 let responder_id = session.user_id();
2588 let requester_id = UserId::from_proto(request.requester_id);
2589 let db = session.db().await;
2590 if request.response == proto::ContactRequestResponse::Dismiss as i32 {
2591 db.dismiss_contact_notification(responder_id, requester_id)
2592 .await?;
2593 } else {
2594 let accept = request.response == proto::ContactRequestResponse::Accept as i32;
2595
2596 let notifications = db
2597 .respond_to_contact_request(responder_id, requester_id, accept)
2598 .await?;
2599 let requester_busy = db.is_user_busy(requester_id).await?;
2600 let responder_busy = db.is_user_busy(responder_id).await?;
2601
2602 let pool = session.connection_pool().await;
2603 // Update responder with new contact
2604 let mut update = proto::UpdateContacts::default();
2605 if accept {
2606 update
2607 .contacts
2608 .push(contact_for_user(requester_id, requester_busy, &pool));
2609 }
2610 update
2611 .remove_incoming_requests
2612 .push(requester_id.to_proto());
2613 for connection_id in pool.user_connection_ids(responder_id) {
2614 session.peer.send(connection_id, update.clone())?;
2615 }
2616
2617 // Update requester with new contact
2618 let mut update = proto::UpdateContacts::default();
2619 if accept {
2620 update
2621 .contacts
2622 .push(contact_for_user(responder_id, responder_busy, &pool));
2623 }
2624 update
2625 .remove_outgoing_requests
2626 .push(responder_id.to_proto());
2627
2628 for connection_id in pool.user_connection_ids(requester_id) {
2629 session.peer.send(connection_id, update.clone())?;
2630 }
2631
2632 send_notifications(&pool, &session.peer, notifications);
2633 }
2634
2635 response.send(proto::Ack {})?;
2636 Ok(())
2637}
2638
2639/// Remove a contact.
2640async fn remove_contact(
2641 request: proto::RemoveContact,
2642 response: Response<proto::RemoveContact>,
2643 session: Session,
2644) -> Result<()> {
2645 let requester_id = session.user_id();
2646 let responder_id = UserId::from_proto(request.user_id);
2647 let db = session.db().await;
2648 let (contact_accepted, deleted_notification_id) =
2649 db.remove_contact(requester_id, responder_id).await?;
2650
2651 let pool = session.connection_pool().await;
2652 // Update outgoing contact requests of requester
2653 let mut update = proto::UpdateContacts::default();
2654 if contact_accepted {
2655 update.remove_contacts.push(responder_id.to_proto());
2656 } else {
2657 update
2658 .remove_outgoing_requests
2659 .push(responder_id.to_proto());
2660 }
2661 for connection_id in pool.user_connection_ids(requester_id) {
2662 session.peer.send(connection_id, update.clone())?;
2663 }
2664
2665 // Update incoming contact requests of responder
2666 let mut update = proto::UpdateContacts::default();
2667 if contact_accepted {
2668 update.remove_contacts.push(requester_id.to_proto());
2669 } else {
2670 update
2671 .remove_incoming_requests
2672 .push(requester_id.to_proto());
2673 }
2674 for connection_id in pool.user_connection_ids(responder_id) {
2675 session.peer.send(connection_id, update.clone())?;
2676 if let Some(notification_id) = deleted_notification_id {
2677 session.peer.send(
2678 connection_id,
2679 proto::DeleteNotification {
2680 notification_id: notification_id.to_proto(),
2681 },
2682 )?;
2683 }
2684 }
2685
2686 response.send(proto::Ack {})?;
2687 Ok(())
2688}
2689
2690fn should_auto_subscribe_to_channels(version: ZedVersion) -> bool {
2691 version.0.minor() < 139
2692}
2693
2694async fn current_plan(db: &Arc<Database>, user_id: UserId, is_staff: bool) -> Result<proto::Plan> {
2695 if is_staff {
2696 return Ok(proto::Plan::ZedPro);
2697 }
2698
2699 let subscription = db.get_active_billing_subscription(user_id).await?;
2700 let subscription_kind = subscription.and_then(|subscription| subscription.kind);
2701
2702 let plan = if let Some(subscription_kind) = subscription_kind {
2703 match subscription_kind {
2704 SubscriptionKind::ZedPro => proto::Plan::ZedPro,
2705 SubscriptionKind::ZedProTrial => proto::Plan::ZedProTrial,
2706 SubscriptionKind::ZedFree => proto::Plan::Free,
2707 }
2708 } else {
2709 proto::Plan::Free
2710 };
2711
2712 Ok(plan)
2713}
2714
2715async fn make_update_user_plan_message(
2716 user: &User,
2717 is_staff: bool,
2718 db: &Arc<Database>,
2719 llm_db: Option<Arc<LlmDatabase>>,
2720) -> Result<proto::UpdateUserPlan> {
2721 let feature_flags = db.get_user_flags(user.id).await?;
2722 let plan = current_plan(db, user.id, is_staff).await?;
2723 let billing_customer = db.get_billing_customer_by_user_id(user.id).await?;
2724 let billing_preferences = db.get_billing_preferences(user.id).await?;
2725
2726 let (subscription_period, usage) = if let Some(llm_db) = llm_db {
2727 let subscription = db.get_active_billing_subscription(user.id).await?;
2728
2729 let subscription_period =
2730 crate::db::billing_subscription::Model::current_period(subscription, is_staff);
2731
2732 let usage = if let Some((period_start_at, period_end_at)) = subscription_period {
2733 llm_db
2734 .get_subscription_usage_for_period(user.id, period_start_at, period_end_at)
2735 .await?
2736 } else {
2737 None
2738 };
2739
2740 (subscription_period, usage)
2741 } else {
2742 (None, None)
2743 };
2744
2745 let account_too_young =
2746 !matches!(plan, proto::Plan::ZedPro) && user.account_age() < MIN_ACCOUNT_AGE_FOR_LLM_USE;
2747
2748 Ok(proto::UpdateUserPlan {
2749 plan: plan.into(),
2750 trial_started_at: billing_customer
2751 .and_then(|billing_customer| billing_customer.trial_started_at)
2752 .map(|trial_started_at| trial_started_at.and_utc().timestamp() as u64),
2753 is_usage_based_billing_enabled: if is_staff {
2754 Some(true)
2755 } else {
2756 billing_preferences.map(|preferences| preferences.model_request_overages_enabled)
2757 },
2758 subscription_period: subscription_period.map(|(started_at, ended_at)| {
2759 proto::SubscriptionPeriod {
2760 started_at: started_at.timestamp() as u64,
2761 ended_at: ended_at.timestamp() as u64,
2762 }
2763 }),
2764 account_too_young: Some(account_too_young),
2765 usage: usage.map(|usage| {
2766 let plan = match plan {
2767 proto::Plan::Free => zed_llm_client::Plan::ZedFree,
2768 proto::Plan::ZedPro => zed_llm_client::Plan::ZedPro,
2769 proto::Plan::ZedProTrial => zed_llm_client::Plan::ZedProTrial,
2770 };
2771
2772 let model_requests_limit = match plan.model_requests_limit() {
2773 zed_llm_client::UsageLimit::Limited(limit) => {
2774 let limit = if plan == zed_llm_client::Plan::ZedProTrial
2775 && feature_flags
2776 .iter()
2777 .any(|flag| flag == AGENT_EXTENDED_TRIAL_FEATURE_FLAG)
2778 {
2779 1_000
2780 } else {
2781 limit
2782 };
2783
2784 zed_llm_client::UsageLimit::Limited(limit)
2785 }
2786 zed_llm_client::UsageLimit::Unlimited => zed_llm_client::UsageLimit::Unlimited,
2787 };
2788
2789 proto::SubscriptionUsage {
2790 model_requests_usage_amount: usage.model_requests as u32,
2791 model_requests_usage_limit: Some(proto::UsageLimit {
2792 variant: Some(match model_requests_limit {
2793 zed_llm_client::UsageLimit::Limited(limit) => {
2794 proto::usage_limit::Variant::Limited(proto::usage_limit::Limited {
2795 limit: limit as u32,
2796 })
2797 }
2798 zed_llm_client::UsageLimit::Unlimited => {
2799 proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {})
2800 }
2801 }),
2802 }),
2803 edit_predictions_usage_amount: usage.edit_predictions as u32,
2804 edit_predictions_usage_limit: Some(proto::UsageLimit {
2805 variant: Some(match plan.edit_predictions_limit() {
2806 zed_llm_client::UsageLimit::Limited(limit) => {
2807 proto::usage_limit::Variant::Limited(proto::usage_limit::Limited {
2808 limit: limit as u32,
2809 })
2810 }
2811 zed_llm_client::UsageLimit::Unlimited => {
2812 proto::usage_limit::Variant::Unlimited(proto::usage_limit::Unlimited {})
2813 }
2814 }),
2815 }),
2816 }
2817 }),
2818 })
2819}
2820
2821async fn update_user_plan(session: &Session) -> Result<()> {
2822 let db = session.db().await;
2823
2824 let update_user_plan = make_update_user_plan_message(
2825 session.principal.user(),
2826 session.is_staff(),
2827 &db.0,
2828 session.app_state.llm_db.clone(),
2829 )
2830 .await?;
2831
2832 session
2833 .peer
2834 .send(session.connection_id, update_user_plan)
2835 .trace_err();
2836
2837 Ok(())
2838}
2839
2840async fn subscribe_to_channels(_: proto::SubscribeToChannels, session: Session) -> Result<()> {
2841 subscribe_user_to_channels(session.user_id(), &session).await?;
2842 Ok(())
2843}
2844
2845async fn subscribe_user_to_channels(user_id: UserId, session: &Session) -> Result<(), Error> {
2846 let channels_for_user = session.db().await.get_channels_for_user(user_id).await?;
2847 let mut pool = session.connection_pool().await;
2848 for membership in &channels_for_user.channel_memberships {
2849 pool.subscribe_to_channel(user_id, membership.channel_id, membership.role)
2850 }
2851 session.peer.send(
2852 session.connection_id,
2853 build_update_user_channels(&channels_for_user),
2854 )?;
2855 session.peer.send(
2856 session.connection_id,
2857 build_channels_update(channels_for_user),
2858 )?;
2859 Ok(())
2860}
2861
2862/// Creates a new channel.
2863async fn create_channel(
2864 request: proto::CreateChannel,
2865 response: Response<proto::CreateChannel>,
2866 session: Session,
2867) -> Result<()> {
2868 let db = session.db().await;
2869
2870 let parent_id = request.parent_id.map(ChannelId::from_proto);
2871 let (channel, membership) = db
2872 .create_channel(&request.name, parent_id, session.user_id())
2873 .await?;
2874
2875 let root_id = channel.root_id();
2876 let channel = Channel::from_model(channel);
2877
2878 response.send(proto::CreateChannelResponse {
2879 channel: Some(channel.to_proto()),
2880 parent_id: request.parent_id,
2881 })?;
2882
2883 let mut connection_pool = session.connection_pool().await;
2884 if let Some(membership) = membership {
2885 connection_pool.subscribe_to_channel(
2886 membership.user_id,
2887 membership.channel_id,
2888 membership.role,
2889 );
2890 let update = proto::UpdateUserChannels {
2891 channel_memberships: vec![proto::ChannelMembership {
2892 channel_id: membership.channel_id.to_proto(),
2893 role: membership.role.into(),
2894 }],
2895 ..Default::default()
2896 };
2897 for connection_id in connection_pool.user_connection_ids(membership.user_id) {
2898 session.peer.send(connection_id, update.clone())?;
2899 }
2900 }
2901
2902 for (connection_id, role) in connection_pool.channel_connection_ids(root_id) {
2903 if !role.can_see_channel(channel.visibility) {
2904 continue;
2905 }
2906
2907 let update = proto::UpdateChannels {
2908 channels: vec![channel.to_proto()],
2909 ..Default::default()
2910 };
2911 session.peer.send(connection_id, update.clone())?;
2912 }
2913
2914 Ok(())
2915}
2916
2917/// Delete a channel
2918async fn delete_channel(
2919 request: proto::DeleteChannel,
2920 response: Response<proto::DeleteChannel>,
2921 session: Session,
2922) -> Result<()> {
2923 let db = session.db().await;
2924
2925 let channel_id = request.channel_id;
2926 let (root_channel, removed_channels) = db
2927 .delete_channel(ChannelId::from_proto(channel_id), session.user_id())
2928 .await?;
2929 response.send(proto::Ack {})?;
2930
2931 // Notify members of removed channels
2932 let mut update = proto::UpdateChannels::default();
2933 update
2934 .delete_channels
2935 .extend(removed_channels.into_iter().map(|id| id.to_proto()));
2936
2937 let connection_pool = session.connection_pool().await;
2938 for (connection_id, _) in connection_pool.channel_connection_ids(root_channel) {
2939 session.peer.send(connection_id, update.clone())?;
2940 }
2941
2942 Ok(())
2943}
2944
2945/// Invite someone to join a channel.
2946async fn invite_channel_member(
2947 request: proto::InviteChannelMember,
2948 response: Response<proto::InviteChannelMember>,
2949 session: Session,
2950) -> Result<()> {
2951 let db = session.db().await;
2952 let channel_id = ChannelId::from_proto(request.channel_id);
2953 let invitee_id = UserId::from_proto(request.user_id);
2954 let InviteMemberResult {
2955 channel,
2956 notifications,
2957 } = db
2958 .invite_channel_member(
2959 channel_id,
2960 invitee_id,
2961 session.user_id(),
2962 request.role().into(),
2963 )
2964 .await?;
2965
2966 let update = proto::UpdateChannels {
2967 channel_invitations: vec![channel.to_proto()],
2968 ..Default::default()
2969 };
2970
2971 let connection_pool = session.connection_pool().await;
2972 for connection_id in connection_pool.user_connection_ids(invitee_id) {
2973 session.peer.send(connection_id, update.clone())?;
2974 }
2975
2976 send_notifications(&connection_pool, &session.peer, notifications);
2977
2978 response.send(proto::Ack {})?;
2979 Ok(())
2980}
2981
2982/// remove someone from a channel
2983async fn remove_channel_member(
2984 request: proto::RemoveChannelMember,
2985 response: Response<proto::RemoveChannelMember>,
2986 session: Session,
2987) -> Result<()> {
2988 let db = session.db().await;
2989 let channel_id = ChannelId::from_proto(request.channel_id);
2990 let member_id = UserId::from_proto(request.user_id);
2991
2992 let RemoveChannelMemberResult {
2993 membership_update,
2994 notification_id,
2995 } = db
2996 .remove_channel_member(channel_id, member_id, session.user_id())
2997 .await?;
2998
2999 let mut connection_pool = session.connection_pool().await;
3000 notify_membership_updated(
3001 &mut connection_pool,
3002 membership_update,
3003 member_id,
3004 &session.peer,
3005 );
3006 for connection_id in connection_pool.user_connection_ids(member_id) {
3007 if let Some(notification_id) = notification_id {
3008 session
3009 .peer
3010 .send(
3011 connection_id,
3012 proto::DeleteNotification {
3013 notification_id: notification_id.to_proto(),
3014 },
3015 )
3016 .trace_err();
3017 }
3018 }
3019
3020 response.send(proto::Ack {})?;
3021 Ok(())
3022}
3023
3024/// Toggle the channel between public and private.
3025/// Care is taken to maintain the invariant that public channels only descend from public channels,
3026/// (though members-only channels can appear at any point in the hierarchy).
3027async fn set_channel_visibility(
3028 request: proto::SetChannelVisibility,
3029 response: Response<proto::SetChannelVisibility>,
3030 session: Session,
3031) -> Result<()> {
3032 let db = session.db().await;
3033 let channel_id = ChannelId::from_proto(request.channel_id);
3034 let visibility = request.visibility().into();
3035
3036 let channel_model = db
3037 .set_channel_visibility(channel_id, visibility, session.user_id())
3038 .await?;
3039 let root_id = channel_model.root_id();
3040 let channel = Channel::from_model(channel_model);
3041
3042 let mut connection_pool = session.connection_pool().await;
3043 for (user_id, role) in connection_pool
3044 .channel_user_ids(root_id)
3045 .collect::<Vec<_>>()
3046 .into_iter()
3047 {
3048 let update = if role.can_see_channel(channel.visibility) {
3049 connection_pool.subscribe_to_channel(user_id, channel_id, role);
3050 proto::UpdateChannels {
3051 channels: vec![channel.to_proto()],
3052 ..Default::default()
3053 }
3054 } else {
3055 connection_pool.unsubscribe_from_channel(&user_id, &channel_id);
3056 proto::UpdateChannels {
3057 delete_channels: vec![channel.id.to_proto()],
3058 ..Default::default()
3059 }
3060 };
3061
3062 for connection_id in connection_pool.user_connection_ids(user_id) {
3063 session.peer.send(connection_id, update.clone())?;
3064 }
3065 }
3066
3067 response.send(proto::Ack {})?;
3068 Ok(())
3069}
3070
3071/// Alter the role for a user in the channel.
3072async fn set_channel_member_role(
3073 request: proto::SetChannelMemberRole,
3074 response: Response<proto::SetChannelMemberRole>,
3075 session: Session,
3076) -> Result<()> {
3077 let db = session.db().await;
3078 let channel_id = ChannelId::from_proto(request.channel_id);
3079 let member_id = UserId::from_proto(request.user_id);
3080 let result = db
3081 .set_channel_member_role(
3082 channel_id,
3083 session.user_id(),
3084 member_id,
3085 request.role().into(),
3086 )
3087 .await?;
3088
3089 match result {
3090 db::SetMemberRoleResult::MembershipUpdated(membership_update) => {
3091 let mut connection_pool = session.connection_pool().await;
3092 notify_membership_updated(
3093 &mut connection_pool,
3094 membership_update,
3095 member_id,
3096 &session.peer,
3097 )
3098 }
3099 db::SetMemberRoleResult::InviteUpdated(channel) => {
3100 let update = proto::UpdateChannels {
3101 channel_invitations: vec![channel.to_proto()],
3102 ..Default::default()
3103 };
3104
3105 for connection_id in session
3106 .connection_pool()
3107 .await
3108 .user_connection_ids(member_id)
3109 {
3110 session.peer.send(connection_id, update.clone())?;
3111 }
3112 }
3113 }
3114
3115 response.send(proto::Ack {})?;
3116 Ok(())
3117}
3118
3119/// Change the name of a channel
3120async fn rename_channel(
3121 request: proto::RenameChannel,
3122 response: Response<proto::RenameChannel>,
3123 session: Session,
3124) -> Result<()> {
3125 let db = session.db().await;
3126 let channel_id = ChannelId::from_proto(request.channel_id);
3127 let channel_model = db
3128 .rename_channel(channel_id, session.user_id(), &request.name)
3129 .await?;
3130 let root_id = channel_model.root_id();
3131 let channel = Channel::from_model(channel_model);
3132
3133 response.send(proto::RenameChannelResponse {
3134 channel: Some(channel.to_proto()),
3135 })?;
3136
3137 let connection_pool = session.connection_pool().await;
3138 let update = proto::UpdateChannels {
3139 channels: vec![channel.to_proto()],
3140 ..Default::default()
3141 };
3142 for (connection_id, role) in connection_pool.channel_connection_ids(root_id) {
3143 if role.can_see_channel(channel.visibility) {
3144 session.peer.send(connection_id, update.clone())?;
3145 }
3146 }
3147
3148 Ok(())
3149}
3150
3151/// Move a channel to a new parent.
3152async fn move_channel(
3153 request: proto::MoveChannel,
3154 response: Response<proto::MoveChannel>,
3155 session: Session,
3156) -> Result<()> {
3157 let channel_id = ChannelId::from_proto(request.channel_id);
3158 let to = ChannelId::from_proto(request.to);
3159
3160 let (root_id, channels) = session
3161 .db()
3162 .await
3163 .move_channel(channel_id, to, session.user_id())
3164 .await?;
3165
3166 let connection_pool = session.connection_pool().await;
3167 for (connection_id, role) in connection_pool.channel_connection_ids(root_id) {
3168 let channels = channels
3169 .iter()
3170 .filter_map(|channel| {
3171 if role.can_see_channel(channel.visibility) {
3172 Some(channel.to_proto())
3173 } else {
3174 None
3175 }
3176 })
3177 .collect::<Vec<_>>();
3178 if channels.is_empty() {
3179 continue;
3180 }
3181
3182 let update = proto::UpdateChannels {
3183 channels,
3184 ..Default::default()
3185 };
3186
3187 session.peer.send(connection_id, update.clone())?;
3188 }
3189
3190 response.send(Ack {})?;
3191 Ok(())
3192}
3193
3194/// Get the list of channel members
3195async fn get_channel_members(
3196 request: proto::GetChannelMembers,
3197 response: Response<proto::GetChannelMembers>,
3198 session: Session,
3199) -> Result<()> {
3200 let db = session.db().await;
3201 let channel_id = ChannelId::from_proto(request.channel_id);
3202 let limit = if request.limit == 0 {
3203 u16::MAX as u64
3204 } else {
3205 request.limit
3206 };
3207 let (members, users) = db
3208 .get_channel_participant_details(channel_id, &request.query, limit, session.user_id())
3209 .await?;
3210 response.send(proto::GetChannelMembersResponse { members, users })?;
3211 Ok(())
3212}
3213
3214/// Accept or decline a channel invitation.
3215async fn respond_to_channel_invite(
3216 request: proto::RespondToChannelInvite,
3217 response: Response<proto::RespondToChannelInvite>,
3218 session: Session,
3219) -> Result<()> {
3220 let db = session.db().await;
3221 let channel_id = ChannelId::from_proto(request.channel_id);
3222 let RespondToChannelInvite {
3223 membership_update,
3224 notifications,
3225 } = db
3226 .respond_to_channel_invite(channel_id, session.user_id(), request.accept)
3227 .await?;
3228
3229 let mut connection_pool = session.connection_pool().await;
3230 if let Some(membership_update) = membership_update {
3231 notify_membership_updated(
3232 &mut connection_pool,
3233 membership_update,
3234 session.user_id(),
3235 &session.peer,
3236 );
3237 } else {
3238 let update = proto::UpdateChannels {
3239 remove_channel_invitations: vec![channel_id.to_proto()],
3240 ..Default::default()
3241 };
3242
3243 for connection_id in connection_pool.user_connection_ids(session.user_id()) {
3244 session.peer.send(connection_id, update.clone())?;
3245 }
3246 };
3247
3248 send_notifications(&connection_pool, &session.peer, notifications);
3249
3250 response.send(proto::Ack {})?;
3251
3252 Ok(())
3253}
3254
3255/// Join the channels' room
3256async fn join_channel(
3257 request: proto::JoinChannel,
3258 response: Response<proto::JoinChannel>,
3259 session: Session,
3260) -> Result<()> {
3261 let channel_id = ChannelId::from_proto(request.channel_id);
3262 join_channel_internal(channel_id, Box::new(response), session).await
3263}
3264
3265trait JoinChannelInternalResponse {
3266 fn send(self, result: proto::JoinRoomResponse) -> Result<()>;
3267}
3268impl JoinChannelInternalResponse for Response<proto::JoinChannel> {
3269 fn send(self, result: proto::JoinRoomResponse) -> Result<()> {
3270 Response::<proto::JoinChannel>::send(self, result)
3271 }
3272}
3273impl JoinChannelInternalResponse for Response<proto::JoinRoom> {
3274 fn send(self, result: proto::JoinRoomResponse) -> Result<()> {
3275 Response::<proto::JoinRoom>::send(self, result)
3276 }
3277}
3278
3279async fn join_channel_internal(
3280 channel_id: ChannelId,
3281 response: Box<impl JoinChannelInternalResponse>,
3282 session: Session,
3283) -> Result<()> {
3284 let joined_room = {
3285 let mut db = session.db().await;
3286 // If zed quits without leaving the room, and the user re-opens zed before the
3287 // RECONNECT_TIMEOUT, we need to make sure that we kick the user out of the previous
3288 // room they were in.
3289 if let Some(connection) = db.stale_room_connection(session.user_id()).await? {
3290 tracing::info!(
3291 stale_connection_id = %connection,
3292 "cleaning up stale connection",
3293 );
3294 drop(db);
3295 leave_room_for_session(&session, connection).await?;
3296 db = session.db().await;
3297 }
3298
3299 let (joined_room, membership_updated, role) = db
3300 .join_channel(channel_id, session.user_id(), session.connection_id)
3301 .await?;
3302
3303 let live_kit_connection_info =
3304 session
3305 .app_state
3306 .livekit_client
3307 .as_ref()
3308 .and_then(|live_kit| {
3309 let (can_publish, token) = if role == ChannelRole::Guest {
3310 (
3311 false,
3312 live_kit
3313 .guest_token(
3314 &joined_room.room.livekit_room,
3315 &session.user_id().to_string(),
3316 )
3317 .trace_err()?,
3318 )
3319 } else {
3320 (
3321 true,
3322 live_kit
3323 .room_token(
3324 &joined_room.room.livekit_room,
3325 &session.user_id().to_string(),
3326 )
3327 .trace_err()?,
3328 )
3329 };
3330
3331 Some(LiveKitConnectionInfo {
3332 server_url: live_kit.url().into(),
3333 token,
3334 can_publish,
3335 })
3336 });
3337
3338 response.send(proto::JoinRoomResponse {
3339 room: Some(joined_room.room.clone()),
3340 channel_id: joined_room
3341 .channel
3342 .as_ref()
3343 .map(|channel| channel.id.to_proto()),
3344 live_kit_connection_info,
3345 })?;
3346
3347 let mut connection_pool = session.connection_pool().await;
3348 if let Some(membership_updated) = membership_updated {
3349 notify_membership_updated(
3350 &mut connection_pool,
3351 membership_updated,
3352 session.user_id(),
3353 &session.peer,
3354 );
3355 }
3356
3357 room_updated(&joined_room.room, &session.peer);
3358
3359 joined_room
3360 };
3361
3362 channel_updated(
3363 &joined_room.channel.context("channel not returned")?,
3364 &joined_room.room,
3365 &session.peer,
3366 &*session.connection_pool().await,
3367 );
3368
3369 update_user_contacts(session.user_id(), &session).await?;
3370 Ok(())
3371}
3372
3373/// Start editing the channel notes
3374async fn join_channel_buffer(
3375 request: proto::JoinChannelBuffer,
3376 response: Response<proto::JoinChannelBuffer>,
3377 session: Session,
3378) -> Result<()> {
3379 let db = session.db().await;
3380 let channel_id = ChannelId::from_proto(request.channel_id);
3381
3382 let open_response = db
3383 .join_channel_buffer(channel_id, session.user_id(), session.connection_id)
3384 .await?;
3385
3386 let collaborators = open_response.collaborators.clone();
3387 response.send(open_response)?;
3388
3389 let update = UpdateChannelBufferCollaborators {
3390 channel_id: channel_id.to_proto(),
3391 collaborators: collaborators.clone(),
3392 };
3393 channel_buffer_updated(
3394 session.connection_id,
3395 collaborators
3396 .iter()
3397 .filter_map(|collaborator| Some(collaborator.peer_id?.into())),
3398 &update,
3399 &session.peer,
3400 );
3401
3402 Ok(())
3403}
3404
3405/// Edit the channel notes
3406async fn update_channel_buffer(
3407 request: proto::UpdateChannelBuffer,
3408 session: Session,
3409) -> Result<()> {
3410 let db = session.db().await;
3411 let channel_id = ChannelId::from_proto(request.channel_id);
3412
3413 let (collaborators, epoch, version) = db
3414 .update_channel_buffer(channel_id, session.user_id(), &request.operations)
3415 .await?;
3416
3417 channel_buffer_updated(
3418 session.connection_id,
3419 collaborators.clone(),
3420 &proto::UpdateChannelBuffer {
3421 channel_id: channel_id.to_proto(),
3422 operations: request.operations,
3423 },
3424 &session.peer,
3425 );
3426
3427 let pool = &*session.connection_pool().await;
3428
3429 let non_collaborators =
3430 pool.channel_connection_ids(channel_id)
3431 .filter_map(|(connection_id, _)| {
3432 if collaborators.contains(&connection_id) {
3433 None
3434 } else {
3435 Some(connection_id)
3436 }
3437 });
3438
3439 broadcast(None, non_collaborators, |peer_id| {
3440 session.peer.send(
3441 peer_id,
3442 proto::UpdateChannels {
3443 latest_channel_buffer_versions: vec![proto::ChannelBufferVersion {
3444 channel_id: channel_id.to_proto(),
3445 epoch: epoch as u64,
3446 version: version.clone(),
3447 }],
3448 ..Default::default()
3449 },
3450 )
3451 });
3452
3453 Ok(())
3454}
3455
3456/// Rejoin the channel notes after a connection blip
3457async fn rejoin_channel_buffers(
3458 request: proto::RejoinChannelBuffers,
3459 response: Response<proto::RejoinChannelBuffers>,
3460 session: Session,
3461) -> Result<()> {
3462 let db = session.db().await;
3463 let buffers = db
3464 .rejoin_channel_buffers(&request.buffers, session.user_id(), session.connection_id)
3465 .await?;
3466
3467 for rejoined_buffer in &buffers {
3468 let collaborators_to_notify = rejoined_buffer
3469 .buffer
3470 .collaborators
3471 .iter()
3472 .filter_map(|c| Some(c.peer_id?.into()));
3473 channel_buffer_updated(
3474 session.connection_id,
3475 collaborators_to_notify,
3476 &proto::UpdateChannelBufferCollaborators {
3477 channel_id: rejoined_buffer.buffer.channel_id,
3478 collaborators: rejoined_buffer.buffer.collaborators.clone(),
3479 },
3480 &session.peer,
3481 );
3482 }
3483
3484 response.send(proto::RejoinChannelBuffersResponse {
3485 buffers: buffers.into_iter().map(|b| b.buffer).collect(),
3486 })?;
3487
3488 Ok(())
3489}
3490
3491/// Stop editing the channel notes
3492async fn leave_channel_buffer(
3493 request: proto::LeaveChannelBuffer,
3494 response: Response<proto::LeaveChannelBuffer>,
3495 session: Session,
3496) -> Result<()> {
3497 let db = session.db().await;
3498 let channel_id = ChannelId::from_proto(request.channel_id);
3499
3500 let left_buffer = db
3501 .leave_channel_buffer(channel_id, session.connection_id)
3502 .await?;
3503
3504 response.send(Ack {})?;
3505
3506 channel_buffer_updated(
3507 session.connection_id,
3508 left_buffer.connections,
3509 &proto::UpdateChannelBufferCollaborators {
3510 channel_id: channel_id.to_proto(),
3511 collaborators: left_buffer.collaborators,
3512 },
3513 &session.peer,
3514 );
3515
3516 Ok(())
3517}
3518
3519fn channel_buffer_updated<T: EnvelopedMessage>(
3520 sender_id: ConnectionId,
3521 collaborators: impl IntoIterator<Item = ConnectionId>,
3522 message: &T,
3523 peer: &Peer,
3524) {
3525 broadcast(Some(sender_id), collaborators, |peer_id| {
3526 peer.send(peer_id, message.clone())
3527 });
3528}
3529
3530fn send_notifications(
3531 connection_pool: &ConnectionPool,
3532 peer: &Peer,
3533 notifications: db::NotificationBatch,
3534) {
3535 for (user_id, notification) in notifications {
3536 for connection_id in connection_pool.user_connection_ids(user_id) {
3537 if let Err(error) = peer.send(
3538 connection_id,
3539 proto::AddNotification {
3540 notification: Some(notification.clone()),
3541 },
3542 ) {
3543 tracing::error!(
3544 "failed to send notification to {:?} {}",
3545 connection_id,
3546 error
3547 );
3548 }
3549 }
3550 }
3551}
3552
3553/// Send a message to the channel
3554async fn send_channel_message(
3555 request: proto::SendChannelMessage,
3556 response: Response<proto::SendChannelMessage>,
3557 session: Session,
3558) -> Result<()> {
3559 // Validate the message body.
3560 let body = request.body.trim().to_string();
3561 if body.len() > MAX_MESSAGE_LEN {
3562 return Err(anyhow!("message is too long"))?;
3563 }
3564 if body.is_empty() {
3565 return Err(anyhow!("message can't be blank"))?;
3566 }
3567
3568 // TODO: adjust mentions if body is trimmed
3569
3570 let timestamp = OffsetDateTime::now_utc();
3571 let nonce = request.nonce.context("nonce can't be blank")?;
3572
3573 let channel_id = ChannelId::from_proto(request.channel_id);
3574 let CreatedChannelMessage {
3575 message_id,
3576 participant_connection_ids,
3577 notifications,
3578 } = session
3579 .db()
3580 .await
3581 .create_channel_message(
3582 channel_id,
3583 session.user_id(),
3584 &body,
3585 &request.mentions,
3586 timestamp,
3587 nonce.clone().into(),
3588 request.reply_to_message_id.map(MessageId::from_proto),
3589 )
3590 .await?;
3591
3592 let message = proto::ChannelMessage {
3593 sender_id: session.user_id().to_proto(),
3594 id: message_id.to_proto(),
3595 body,
3596 mentions: request.mentions,
3597 timestamp: timestamp.unix_timestamp() as u64,
3598 nonce: Some(nonce),
3599 reply_to_message_id: request.reply_to_message_id,
3600 edited_at: None,
3601 };
3602 broadcast(
3603 Some(session.connection_id),
3604 participant_connection_ids.clone(),
3605 |connection| {
3606 session.peer.send(
3607 connection,
3608 proto::ChannelMessageSent {
3609 channel_id: channel_id.to_proto(),
3610 message: Some(message.clone()),
3611 },
3612 )
3613 },
3614 );
3615 response.send(proto::SendChannelMessageResponse {
3616 message: Some(message),
3617 })?;
3618
3619 let pool = &*session.connection_pool().await;
3620 let non_participants =
3621 pool.channel_connection_ids(channel_id)
3622 .filter_map(|(connection_id, _)| {
3623 if participant_connection_ids.contains(&connection_id) {
3624 None
3625 } else {
3626 Some(connection_id)
3627 }
3628 });
3629 broadcast(None, non_participants, |peer_id| {
3630 session.peer.send(
3631 peer_id,
3632 proto::UpdateChannels {
3633 latest_channel_message_ids: vec![proto::ChannelMessageId {
3634 channel_id: channel_id.to_proto(),
3635 message_id: message_id.to_proto(),
3636 }],
3637 ..Default::default()
3638 },
3639 )
3640 });
3641 send_notifications(pool, &session.peer, notifications);
3642
3643 Ok(())
3644}
3645
3646/// Delete a channel message
3647async fn remove_channel_message(
3648 request: proto::RemoveChannelMessage,
3649 response: Response<proto::RemoveChannelMessage>,
3650 session: Session,
3651) -> Result<()> {
3652 let channel_id = ChannelId::from_proto(request.channel_id);
3653 let message_id = MessageId::from_proto(request.message_id);
3654 let (connection_ids, existing_notification_ids) = session
3655 .db()
3656 .await
3657 .remove_channel_message(channel_id, message_id, session.user_id())
3658 .await?;
3659
3660 broadcast(
3661 Some(session.connection_id),
3662 connection_ids,
3663 move |connection| {
3664 session.peer.send(connection, request.clone())?;
3665
3666 for notification_id in &existing_notification_ids {
3667 session.peer.send(
3668 connection,
3669 proto::DeleteNotification {
3670 notification_id: (*notification_id).to_proto(),
3671 },
3672 )?;
3673 }
3674
3675 Ok(())
3676 },
3677 );
3678 response.send(proto::Ack {})?;
3679 Ok(())
3680}
3681
3682async fn update_channel_message(
3683 request: proto::UpdateChannelMessage,
3684 response: Response<proto::UpdateChannelMessage>,
3685 session: Session,
3686) -> Result<()> {
3687 let channel_id = ChannelId::from_proto(request.channel_id);
3688 let message_id = MessageId::from_proto(request.message_id);
3689 let updated_at = OffsetDateTime::now_utc();
3690 let UpdatedChannelMessage {
3691 message_id,
3692 participant_connection_ids,
3693 notifications,
3694 reply_to_message_id,
3695 timestamp,
3696 deleted_mention_notification_ids,
3697 updated_mention_notifications,
3698 } = session
3699 .db()
3700 .await
3701 .update_channel_message(
3702 channel_id,
3703 message_id,
3704 session.user_id(),
3705 request.body.as_str(),
3706 &request.mentions,
3707 updated_at,
3708 )
3709 .await?;
3710
3711 let nonce = request.nonce.clone().context("nonce can't be blank")?;
3712
3713 let message = proto::ChannelMessage {
3714 sender_id: session.user_id().to_proto(),
3715 id: message_id.to_proto(),
3716 body: request.body.clone(),
3717 mentions: request.mentions.clone(),
3718 timestamp: timestamp.assume_utc().unix_timestamp() as u64,
3719 nonce: Some(nonce),
3720 reply_to_message_id: reply_to_message_id.map(|id| id.to_proto()),
3721 edited_at: Some(updated_at.unix_timestamp() as u64),
3722 };
3723
3724 response.send(proto::Ack {})?;
3725
3726 let pool = &*session.connection_pool().await;
3727 broadcast(
3728 Some(session.connection_id),
3729 participant_connection_ids,
3730 |connection| {
3731 session.peer.send(
3732 connection,
3733 proto::ChannelMessageUpdate {
3734 channel_id: channel_id.to_proto(),
3735 message: Some(message.clone()),
3736 },
3737 )?;
3738
3739 for notification_id in &deleted_mention_notification_ids {
3740 session.peer.send(
3741 connection,
3742 proto::DeleteNotification {
3743 notification_id: (*notification_id).to_proto(),
3744 },
3745 )?;
3746 }
3747
3748 for notification in &updated_mention_notifications {
3749 session.peer.send(
3750 connection,
3751 proto::UpdateNotification {
3752 notification: Some(notification.clone()),
3753 },
3754 )?;
3755 }
3756
3757 Ok(())
3758 },
3759 );
3760
3761 send_notifications(pool, &session.peer, notifications);
3762
3763 Ok(())
3764}
3765
3766/// Mark a channel message as read
3767async fn acknowledge_channel_message(
3768 request: proto::AckChannelMessage,
3769 session: Session,
3770) -> Result<()> {
3771 let channel_id = ChannelId::from_proto(request.channel_id);
3772 let message_id = MessageId::from_proto(request.message_id);
3773 let notifications = session
3774 .db()
3775 .await
3776 .observe_channel_message(channel_id, session.user_id(), message_id)
3777 .await?;
3778 send_notifications(
3779 &*session.connection_pool().await,
3780 &session.peer,
3781 notifications,
3782 );
3783 Ok(())
3784}
3785
3786/// Mark a buffer version as synced
3787async fn acknowledge_buffer_version(
3788 request: proto::AckBufferOperation,
3789 session: Session,
3790) -> Result<()> {
3791 let buffer_id = BufferId::from_proto(request.buffer_id);
3792 session
3793 .db()
3794 .await
3795 .observe_buffer_version(
3796 buffer_id,
3797 session.user_id(),
3798 request.epoch as i32,
3799 &request.version,
3800 )
3801 .await?;
3802 Ok(())
3803}
3804
3805/// Get a Supermaven API key for the user
3806async fn get_supermaven_api_key(
3807 _request: proto::GetSupermavenApiKey,
3808 response: Response<proto::GetSupermavenApiKey>,
3809 session: Session,
3810) -> Result<()> {
3811 let user_id: String = session.user_id().to_string();
3812 if !session.is_staff() {
3813 return Err(anyhow!("supermaven not enabled for this account"))?;
3814 }
3815
3816 let email = session.email().context("user must have an email")?;
3817
3818 let supermaven_admin_api = session
3819 .supermaven_client
3820 .as_ref()
3821 .context("supermaven not configured")?;
3822
3823 let result = supermaven_admin_api
3824 .try_get_or_create_user(CreateExternalUserRequest { id: user_id, email })
3825 .await?;
3826
3827 response.send(proto::GetSupermavenApiKeyResponse {
3828 api_key: result.api_key,
3829 })?;
3830
3831 Ok(())
3832}
3833
3834/// Start receiving chat updates for a channel
3835async fn join_channel_chat(
3836 request: proto::JoinChannelChat,
3837 response: Response<proto::JoinChannelChat>,
3838 session: Session,
3839) -> Result<()> {
3840 let channel_id = ChannelId::from_proto(request.channel_id);
3841
3842 let db = session.db().await;
3843 db.join_channel_chat(channel_id, session.connection_id, session.user_id())
3844 .await?;
3845 let messages = db
3846 .get_channel_messages(channel_id, session.user_id(), MESSAGE_COUNT_PER_PAGE, None)
3847 .await?;
3848 response.send(proto::JoinChannelChatResponse {
3849 done: messages.len() < MESSAGE_COUNT_PER_PAGE,
3850 messages,
3851 })?;
3852 Ok(())
3853}
3854
3855/// Stop receiving chat updates for a channel
3856async fn leave_channel_chat(request: proto::LeaveChannelChat, session: Session) -> Result<()> {
3857 let channel_id = ChannelId::from_proto(request.channel_id);
3858 session
3859 .db()
3860 .await
3861 .leave_channel_chat(channel_id, session.connection_id, session.user_id())
3862 .await?;
3863 Ok(())
3864}
3865
3866/// Retrieve the chat history for a channel
3867async fn get_channel_messages(
3868 request: proto::GetChannelMessages,
3869 response: Response<proto::GetChannelMessages>,
3870 session: Session,
3871) -> Result<()> {
3872 let channel_id = ChannelId::from_proto(request.channel_id);
3873 let messages = session
3874 .db()
3875 .await
3876 .get_channel_messages(
3877 channel_id,
3878 session.user_id(),
3879 MESSAGE_COUNT_PER_PAGE,
3880 Some(MessageId::from_proto(request.before_message_id)),
3881 )
3882 .await?;
3883 response.send(proto::GetChannelMessagesResponse {
3884 done: messages.len() < MESSAGE_COUNT_PER_PAGE,
3885 messages,
3886 })?;
3887 Ok(())
3888}
3889
3890/// Retrieve specific chat messages
3891async fn get_channel_messages_by_id(
3892 request: proto::GetChannelMessagesById,
3893 response: Response<proto::GetChannelMessagesById>,
3894 session: Session,
3895) -> Result<()> {
3896 let message_ids = request
3897 .message_ids
3898 .iter()
3899 .map(|id| MessageId::from_proto(*id))
3900 .collect::<Vec<_>>();
3901 let messages = session
3902 .db()
3903 .await
3904 .get_channel_messages_by_id(session.user_id(), &message_ids)
3905 .await?;
3906 response.send(proto::GetChannelMessagesResponse {
3907 done: messages.len() < MESSAGE_COUNT_PER_PAGE,
3908 messages,
3909 })?;
3910 Ok(())
3911}
3912
3913/// Retrieve the current users notifications
3914async fn get_notifications(
3915 request: proto::GetNotifications,
3916 response: Response<proto::GetNotifications>,
3917 session: Session,
3918) -> Result<()> {
3919 let notifications = session
3920 .db()
3921 .await
3922 .get_notifications(
3923 session.user_id(),
3924 NOTIFICATION_COUNT_PER_PAGE,
3925 request.before_id.map(db::NotificationId::from_proto),
3926 )
3927 .await?;
3928 response.send(proto::GetNotificationsResponse {
3929 done: notifications.len() < NOTIFICATION_COUNT_PER_PAGE,
3930 notifications,
3931 })?;
3932 Ok(())
3933}
3934
3935/// Mark notifications as read
3936async fn mark_notification_as_read(
3937 request: proto::MarkNotificationRead,
3938 response: Response<proto::MarkNotificationRead>,
3939 session: Session,
3940) -> Result<()> {
3941 let database = &session.db().await;
3942 let notifications = database
3943 .mark_notification_as_read_by_id(
3944 session.user_id(),
3945 NotificationId::from_proto(request.notification_id),
3946 )
3947 .await?;
3948 send_notifications(
3949 &*session.connection_pool().await,
3950 &session.peer,
3951 notifications,
3952 );
3953 response.send(proto::Ack {})?;
3954 Ok(())
3955}
3956
3957/// Get the current users information
3958async fn get_private_user_info(
3959 _request: proto::GetPrivateUserInfo,
3960 response: Response<proto::GetPrivateUserInfo>,
3961 session: Session,
3962) -> Result<()> {
3963 let db = session.db().await;
3964
3965 let metrics_id = db.get_user_metrics_id(session.user_id()).await?;
3966 let user = db
3967 .get_user_by_id(session.user_id())
3968 .await?
3969 .context("user not found")?;
3970 let flags = db.get_user_flags(session.user_id()).await?;
3971
3972 response.send(proto::GetPrivateUserInfoResponse {
3973 metrics_id,
3974 staff: user.admin,
3975 flags,
3976 accepted_tos_at: user.accepted_tos_at.map(|t| t.and_utc().timestamp() as u64),
3977 })?;
3978 Ok(())
3979}
3980
3981/// Accept the terms of service (tos) on behalf of the current user
3982async fn accept_terms_of_service(
3983 _request: proto::AcceptTermsOfService,
3984 response: Response<proto::AcceptTermsOfService>,
3985 session: Session,
3986) -> Result<()> {
3987 let db = session.db().await;
3988
3989 let accepted_tos_at = Utc::now();
3990 db.set_user_accepted_tos_at(session.user_id(), Some(accepted_tos_at.naive_utc()))
3991 .await?;
3992
3993 response.send(proto::AcceptTermsOfServiceResponse {
3994 accepted_tos_at: accepted_tos_at.timestamp() as u64,
3995 })?;
3996 Ok(())
3997}
3998
3999/// The minimum account age an account must have in order to use the LLM service.
4000pub const MIN_ACCOUNT_AGE_FOR_LLM_USE: chrono::Duration = chrono::Duration::days(30);
4001
4002async fn get_llm_api_token(
4003 _request: proto::GetLlmToken,
4004 response: Response<proto::GetLlmToken>,
4005 session: Session,
4006) -> Result<()> {
4007 let db = session.db().await;
4008
4009 let flags = db.get_user_flags(session.user_id()).await?;
4010
4011 let user_id = session.user_id();
4012 let user = db
4013 .get_user_by_id(user_id)
4014 .await?
4015 .with_context(|| format!("user {user_id} not found"))?;
4016
4017 if user.accepted_tos_at.is_none() {
4018 Err(anyhow!("terms of service not accepted"))?
4019 }
4020
4021 let stripe_client = session
4022 .app_state
4023 .stripe_client
4024 .as_ref()
4025 .context("failed to retrieve Stripe client")?;
4026
4027 let stripe_billing = session
4028 .app_state
4029 .stripe_billing
4030 .as_ref()
4031 .context("failed to retrieve Stripe billing object")?;
4032
4033 let billing_customer =
4034 if let Some(billing_customer) = db.get_billing_customer_by_user_id(user.id).await? {
4035 billing_customer
4036 } else {
4037 let customer_id = stripe_billing
4038 .find_or_create_customer_by_email(user.email_address.as_deref())
4039 .await?;
4040
4041 find_or_create_billing_customer(
4042 &session.app_state,
4043 &stripe_client,
4044 stripe::Expandable::Id(customer_id),
4045 )
4046 .await?
4047 .context("billing customer not found")?
4048 };
4049
4050 let billing_subscription =
4051 if let Some(billing_subscription) = db.get_active_billing_subscription(user.id).await? {
4052 billing_subscription
4053 } else {
4054 let stripe_customer_id = billing_customer
4055 .stripe_customer_id
4056 .parse::<stripe::CustomerId>()
4057 .context("failed to parse Stripe customer ID from database")?;
4058
4059 let stripe_subscription = stripe_billing
4060 .subscribe_to_zed_free(stripe_customer_id)
4061 .await?;
4062
4063 db.create_billing_subscription(&db::CreateBillingSubscriptionParams {
4064 billing_customer_id: billing_customer.id,
4065 kind: Some(SubscriptionKind::ZedFree),
4066 stripe_subscription_id: stripe_subscription.id.to_string(),
4067 stripe_subscription_status: stripe_subscription.status.into(),
4068 stripe_cancellation_reason: None,
4069 stripe_current_period_start: Some(stripe_subscription.current_period_start),
4070 stripe_current_period_end: Some(stripe_subscription.current_period_end),
4071 })
4072 .await?
4073 };
4074
4075 let billing_preferences = db.get_billing_preferences(user.id).await?;
4076
4077 let token = LlmTokenClaims::create(
4078 &user,
4079 session.is_staff(),
4080 billing_preferences,
4081 &flags,
4082 billing_subscription,
4083 session.system_id.clone(),
4084 &session.app_state.config,
4085 )?;
4086 response.send(proto::GetLlmTokenResponse { token })?;
4087 Ok(())
4088}
4089
4090fn to_axum_message(message: TungsteniteMessage) -> anyhow::Result<AxumMessage> {
4091 let message = match message {
4092 TungsteniteMessage::Text(payload) => AxumMessage::Text(payload.as_str().to_string()),
4093 TungsteniteMessage::Binary(payload) => AxumMessage::Binary(payload.into()),
4094 TungsteniteMessage::Ping(payload) => AxumMessage::Ping(payload.into()),
4095 TungsteniteMessage::Pong(payload) => AxumMessage::Pong(payload.into()),
4096 TungsteniteMessage::Close(frame) => AxumMessage::Close(frame.map(|frame| AxumCloseFrame {
4097 code: frame.code.into(),
4098 reason: frame.reason.as_str().to_owned().into(),
4099 })),
4100 // We should never receive a frame while reading the message, according
4101 // to the `tungstenite` maintainers:
4102 //
4103 // > It cannot occur when you read messages from the WebSocket, but it
4104 // > can be used when you want to send the raw frames (e.g. you want to
4105 // > send the frames to the WebSocket without composing the full message first).
4106 // >
4107 // > — https://github.com/snapview/tungstenite-rs/issues/268
4108 TungsteniteMessage::Frame(_) => {
4109 bail!("received an unexpected frame while reading the message")
4110 }
4111 };
4112
4113 Ok(message)
4114}
4115
4116fn to_tungstenite_message(message: AxumMessage) -> TungsteniteMessage {
4117 match message {
4118 AxumMessage::Text(payload) => TungsteniteMessage::Text(payload.into()),
4119 AxumMessage::Binary(payload) => TungsteniteMessage::Binary(payload.into()),
4120 AxumMessage::Ping(payload) => TungsteniteMessage::Ping(payload.into()),
4121 AxumMessage::Pong(payload) => TungsteniteMessage::Pong(payload.into()),
4122 AxumMessage::Close(frame) => {
4123 TungsteniteMessage::Close(frame.map(|frame| TungsteniteCloseFrame {
4124 code: frame.code.into(),
4125 reason: frame.reason.as_ref().into(),
4126 }))
4127 }
4128 }
4129}
4130
4131fn notify_membership_updated(
4132 connection_pool: &mut ConnectionPool,
4133 result: MembershipUpdated,
4134 user_id: UserId,
4135 peer: &Peer,
4136) {
4137 for membership in &result.new_channels.channel_memberships {
4138 connection_pool.subscribe_to_channel(user_id, membership.channel_id, membership.role)
4139 }
4140 for channel_id in &result.removed_channels {
4141 connection_pool.unsubscribe_from_channel(&user_id, channel_id)
4142 }
4143
4144 let user_channels_update = proto::UpdateUserChannels {
4145 channel_memberships: result
4146 .new_channels
4147 .channel_memberships
4148 .iter()
4149 .map(|cm| proto::ChannelMembership {
4150 channel_id: cm.channel_id.to_proto(),
4151 role: cm.role.into(),
4152 })
4153 .collect(),
4154 ..Default::default()
4155 };
4156
4157 let mut update = build_channels_update(result.new_channels);
4158 update.delete_channels = result
4159 .removed_channels
4160 .into_iter()
4161 .map(|id| id.to_proto())
4162 .collect();
4163 update.remove_channel_invitations = vec![result.channel_id.to_proto()];
4164
4165 for connection_id in connection_pool.user_connection_ids(user_id) {
4166 peer.send(connection_id, user_channels_update.clone())
4167 .trace_err();
4168 peer.send(connection_id, update.clone()).trace_err();
4169 }
4170}
4171
4172fn build_update_user_channels(channels: &ChannelsForUser) -> proto::UpdateUserChannels {
4173 proto::UpdateUserChannels {
4174 channel_memberships: channels
4175 .channel_memberships
4176 .iter()
4177 .map(|m| proto::ChannelMembership {
4178 channel_id: m.channel_id.to_proto(),
4179 role: m.role.into(),
4180 })
4181 .collect(),
4182 observed_channel_buffer_version: channels.observed_buffer_versions.clone(),
4183 observed_channel_message_id: channels.observed_channel_messages.clone(),
4184 }
4185}
4186
4187fn build_channels_update(channels: ChannelsForUser) -> proto::UpdateChannels {
4188 let mut update = proto::UpdateChannels::default();
4189
4190 for channel in channels.channels {
4191 update.channels.push(channel.to_proto());
4192 }
4193
4194 update.latest_channel_buffer_versions = channels.latest_buffer_versions;
4195 update.latest_channel_message_ids = channels.latest_channel_messages;
4196
4197 for (channel_id, participants) in channels.channel_participants {
4198 update
4199 .channel_participants
4200 .push(proto::ChannelParticipants {
4201 channel_id: channel_id.to_proto(),
4202 participant_user_ids: participants.into_iter().map(|id| id.to_proto()).collect(),
4203 });
4204 }
4205
4206 for channel in channels.invited_channels {
4207 update.channel_invitations.push(channel.to_proto());
4208 }
4209
4210 update
4211}
4212
4213fn build_initial_contacts_update(
4214 contacts: Vec<db::Contact>,
4215 pool: &ConnectionPool,
4216) -> proto::UpdateContacts {
4217 let mut update = proto::UpdateContacts::default();
4218
4219 for contact in contacts {
4220 match contact {
4221 db::Contact::Accepted { user_id, busy } => {
4222 update.contacts.push(contact_for_user(user_id, busy, pool));
4223 }
4224 db::Contact::Outgoing { user_id } => update.outgoing_requests.push(user_id.to_proto()),
4225 db::Contact::Incoming { user_id } => {
4226 update
4227 .incoming_requests
4228 .push(proto::IncomingContactRequest {
4229 requester_id: user_id.to_proto(),
4230 })
4231 }
4232 }
4233 }
4234
4235 update
4236}
4237
4238fn contact_for_user(user_id: UserId, busy: bool, pool: &ConnectionPool) -> proto::Contact {
4239 proto::Contact {
4240 user_id: user_id.to_proto(),
4241 online: pool.is_user_online(user_id),
4242 busy,
4243 }
4244}
4245
4246fn room_updated(room: &proto::Room, peer: &Peer) {
4247 broadcast(
4248 None,
4249 room.participants
4250 .iter()
4251 .filter_map(|participant| Some(participant.peer_id?.into())),
4252 |peer_id| {
4253 peer.send(
4254 peer_id,
4255 proto::RoomUpdated {
4256 room: Some(room.clone()),
4257 },
4258 )
4259 },
4260 );
4261}
4262
4263fn channel_updated(
4264 channel: &db::channel::Model,
4265 room: &proto::Room,
4266 peer: &Peer,
4267 pool: &ConnectionPool,
4268) {
4269 let participants = room
4270 .participants
4271 .iter()
4272 .map(|p| p.user_id)
4273 .collect::<Vec<_>>();
4274
4275 broadcast(
4276 None,
4277 pool.channel_connection_ids(channel.root_id())
4278 .filter_map(|(channel_id, role)| {
4279 role.can_see_channel(channel.visibility)
4280 .then_some(channel_id)
4281 }),
4282 |peer_id| {
4283 peer.send(
4284 peer_id,
4285 proto::UpdateChannels {
4286 channel_participants: vec![proto::ChannelParticipants {
4287 channel_id: channel.id.to_proto(),
4288 participant_user_ids: participants.clone(),
4289 }],
4290 ..Default::default()
4291 },
4292 )
4293 },
4294 );
4295}
4296
4297async fn update_user_contacts(user_id: UserId, session: &Session) -> Result<()> {
4298 let db = session.db().await;
4299
4300 let contacts = db.get_contacts(user_id).await?;
4301 let busy = db.is_user_busy(user_id).await?;
4302
4303 let pool = session.connection_pool().await;
4304 let updated_contact = contact_for_user(user_id, busy, &pool);
4305 for contact in contacts {
4306 if let db::Contact::Accepted {
4307 user_id: contact_user_id,
4308 ..
4309 } = contact
4310 {
4311 for contact_conn_id in pool.user_connection_ids(contact_user_id) {
4312 session
4313 .peer
4314 .send(
4315 contact_conn_id,
4316 proto::UpdateContacts {
4317 contacts: vec![updated_contact.clone()],
4318 remove_contacts: Default::default(),
4319 incoming_requests: Default::default(),
4320 remove_incoming_requests: Default::default(),
4321 outgoing_requests: Default::default(),
4322 remove_outgoing_requests: Default::default(),
4323 },
4324 )
4325 .trace_err();
4326 }
4327 }
4328 }
4329 Ok(())
4330}
4331
4332async fn leave_room_for_session(session: &Session, connection_id: ConnectionId) -> Result<()> {
4333 let mut contacts_to_update = HashSet::default();
4334
4335 let room_id;
4336 let canceled_calls_to_user_ids;
4337 let livekit_room;
4338 let delete_livekit_room;
4339 let room;
4340 let channel;
4341
4342 if let Some(mut left_room) = session.db().await.leave_room(connection_id).await? {
4343 contacts_to_update.insert(session.user_id());
4344
4345 for project in left_room.left_projects.values() {
4346 project_left(project, session);
4347 }
4348
4349 room_id = RoomId::from_proto(left_room.room.id);
4350 canceled_calls_to_user_ids = mem::take(&mut left_room.canceled_calls_to_user_ids);
4351 livekit_room = mem::take(&mut left_room.room.livekit_room);
4352 delete_livekit_room = left_room.deleted;
4353 room = mem::take(&mut left_room.room);
4354 channel = mem::take(&mut left_room.channel);
4355
4356 room_updated(&room, &session.peer);
4357 } else {
4358 return Ok(());
4359 }
4360
4361 if let Some(channel) = channel {
4362 channel_updated(
4363 &channel,
4364 &room,
4365 &session.peer,
4366 &*session.connection_pool().await,
4367 );
4368 }
4369
4370 {
4371 let pool = session.connection_pool().await;
4372 for canceled_user_id in canceled_calls_to_user_ids {
4373 for connection_id in pool.user_connection_ids(canceled_user_id) {
4374 session
4375 .peer
4376 .send(
4377 connection_id,
4378 proto::CallCanceled {
4379 room_id: room_id.to_proto(),
4380 },
4381 )
4382 .trace_err();
4383 }
4384 contacts_to_update.insert(canceled_user_id);
4385 }
4386 }
4387
4388 for contact_user_id in contacts_to_update {
4389 update_user_contacts(contact_user_id, session).await?;
4390 }
4391
4392 if let Some(live_kit) = session.app_state.livekit_client.as_ref() {
4393 live_kit
4394 .remove_participant(livekit_room.clone(), session.user_id().to_string())
4395 .await
4396 .trace_err();
4397
4398 if delete_livekit_room {
4399 live_kit.delete_room(livekit_room).await.trace_err();
4400 }
4401 }
4402
4403 Ok(())
4404}
4405
4406async fn leave_channel_buffers_for_session(session: &Session) -> Result<()> {
4407 let left_channel_buffers = session
4408 .db()
4409 .await
4410 .leave_channel_buffers(session.connection_id)
4411 .await?;
4412
4413 for left_buffer in left_channel_buffers {
4414 channel_buffer_updated(
4415 session.connection_id,
4416 left_buffer.connections,
4417 &proto::UpdateChannelBufferCollaborators {
4418 channel_id: left_buffer.channel_id.to_proto(),
4419 collaborators: left_buffer.collaborators,
4420 },
4421 &session.peer,
4422 );
4423 }
4424
4425 Ok(())
4426}
4427
4428fn project_left(project: &db::LeftProject, session: &Session) {
4429 for connection_id in &project.connection_ids {
4430 if project.should_unshare {
4431 session
4432 .peer
4433 .send(
4434 *connection_id,
4435 proto::UnshareProject {
4436 project_id: project.id.to_proto(),
4437 },
4438 )
4439 .trace_err();
4440 } else {
4441 session
4442 .peer
4443 .send(
4444 *connection_id,
4445 proto::RemoveProjectCollaborator {
4446 project_id: project.id.to_proto(),
4447 peer_id: Some(session.connection_id.into()),
4448 },
4449 )
4450 .trace_err();
4451 }
4452 }
4453}
4454
4455pub trait ResultExt {
4456 type Ok;
4457
4458 fn trace_err(self) -> Option<Self::Ok>;
4459}
4460
4461impl<T, E> ResultExt for Result<T, E>
4462where
4463 E: std::fmt::Debug,
4464{
4465 type Ok = T;
4466
4467 #[track_caller]
4468 fn trace_err(self) -> Option<T> {
4469 match self {
4470 Ok(value) => Some(value),
4471 Err(error) => {
4472 tracing::error!("{:?}", error);
4473 None
4474 }
4475 }
4476 }
4477}