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