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