mod connection_pool; use crate::api::{CloudflareIpCountryHeader, SystemIdHeader};use crate::llm::LlmTokenClaims;use crate::{ auth, db::{ self, BufferId, Capability, Channel, ChannelId, ChannelRole, ChannelsForUser, CreatedChannelMessage, Database, InviteMemberResult, MembershipUpdated, MessageId, NotificationId, Project, ProjectId, RejoinedProject, RemoveChannelMemberResult, ReplicaId, RespondToChannelInvite, RoomId, ServerId, UpdatedChannelMessage, User, UserId, }, executor::Executor, AppState, Config, Error, RateLimit, Result,};use anyhow::{anyhow, bail, Context as _};use async_tungstenite::tungstenite::{ protocol::CloseFrame as TungsteniteCloseFrame, Message as TungsteniteMessage,};use axum::{ body::Body, extract::{ ws::{CloseFrame as AxumCloseFrame, Message as AxumMessage}, ConnectInfo, WebSocketUpgrade, }, headers::{Header, HeaderName}, http::StatusCode, middleware, response::IntoResponse, routing::get, Extension, Router, TypedHeader,};use chrono::Utc;use collections::{HashMap, HashSet};pub use connection_pool::{ConnectionPool, ZedVersion};use core::fmt::{self, Debug, Formatter};use http_client::HttpClient;use open_ai::{OpenAiEmbeddingModel, OPEN_AI_API_URL};use reqwest_client::ReqwestClient;use sha2::Digest;use supermaven_api::{CreateExternalUserRequest, SupermavenAdminApi}; use futures::{ channel::oneshot, future::BoxFuture, stream::FuturesUnordered, FutureExt, SinkExt, StreamExt, TryStreamExt,};use prometheus::{register_int_gauge, IntGauge};use rpc::{ proto::{ self, Ack, AnyTypedEnvelope, EntityMessage, EnvelopedMessage, LiveKitConnectionInfo, RequestMessage, ShareProject, UpdateChannelBufferCollaborators, }, Connection, ConnectionId, ErrorCode, ErrorCodeExt, ErrorExt, Peer, Receipt, TypedEnvelope,};use semantic_version::SemanticVersion;use serde::{Serialize, Serializer};use std::{ any::TypeId, future::Future, marker::PhantomData, mem, net::SocketAddr, ops::{Deref, DerefMut}, rc::Rc, sync::{ atomic::{AtomicBool, Ordering::SeqCst}, Arc, OnceLock, }, time::{Duration, Instant},};use time::OffsetDateTime;use tokio::sync::{watch, MutexGuard, Semaphore};use tower::ServiceBuilder;use tracing::{ field::{self}, info_span, instrument, Instrument,}; pub const RECONNECT_TIMEOUT: Duration = Duration::from_secs(30); // kubernetes gives terminated pods 10s to shutdown gracefully. After they're gone, we can clean up old resources.pub const CLEANUP_TIMEOUT: Duration = Duration::from_secs(15); const MESSAGE_COUNT_PER_PAGE: usize = 100;const MAX_MESSAGE_LEN: usize = 1024;const NOTIFICATION_COUNT_PER_PAGE: usize = 50; type MessageHandler = Box<dyn Send + Sync + Fn(Box<dyn AnyTypedEnvelope>, Session) -> BoxFuture<'static, ()>>; struct Response<R> { peer: Arc<Peer>, receipt: Receipt<R>, responded: Arc<AtomicBool>,} impl<R: RequestMessage> Response<R> { fn send(self, payload: R::Response) -> Result<()> { self.responded.store(true, SeqCst); self.peer.respond(self.receipt, payload)?; Ok(()) }} #[derive(Clone, Debug)]pub enum Principal { User(User), Impersonated { user: User, admin: User },} impl Principal { fn update_span(&self, span: &tracing::Span) { match &self { Principal::User(user) => { span.record("user_id", user.id.0); span.record("login", &user.github_login); } Principal::Impersonated { user, admin } => { span.record("user_id", user.id.0); span.record("login", &user.github_login); span.record("impersonator", &admin.github_login); } } }} #[derive(Clone)]struct Session { principal: Principal, connection_id: ConnectionId, db: Arc<tokio::sync::Mutex<DbHandle>>, peer: Arc<Peer>, connection_pool: Arc<parking_lot::Mutex<ConnectionPool>>, app_state: Arc<AppState>, supermaven_client: Option<Arc<SupermavenAdminApi>>, http_client: Arc<dyn HttpClient>, /// The GeoIP country code for the user. #[allow(unused)] geoip_country_code: Option<String>, system_id: Option<String>, _executor: Executor,} impl Session { async fn db(&self) -> tokio::sync::MutexGuard<DbHandle> { #[cfg(test)] tokio::task::yield_now().await; let guard = self.db.lock().await; #[cfg(test)] tokio::task::yield_now().await; guard } async fn connection_pool(&self) -> ConnectionPoolGuard<'_> { #[cfg(test)] tokio::task::yield_now().await; let guard = self.connection_pool.lock(); ConnectionPoolGuard { guard, _not_send: PhantomData, } } fn is_staff(&self) -> bool { match &self.principal { Principal::User(user) => user.admin, Principal::Impersonated { .. } => true, } } pub async fn has_llm_subscription( &self, db: &MutexGuard<'_, DbHandle>, ) -> anyhow::Result<bool> { if self.is_staff() { return Ok(true); } let user_id = self.user_id(); Ok(db.has_active_billing_subscription(user_id).await?) } pub async fn current_plan( &self, _db: &MutexGuard<'_, DbHandle>, ) -> anyhow::Result<proto::Plan> { if self.is_staff() { Ok(proto::Plan::ZedPro) } else { Ok(proto::Plan::Free) } } fn user_id(&self) -> UserId { match &self.principal { Principal::User(user) => user.id, Principal::Impersonated { user, .. } => user.id, } } pub fn email(&self) -> Option<String> { match &self.principal { Principal::User(user) => user.email_address.clone(), Principal::Impersonated { user, .. } => user.email_address.clone(), } }} impl Debug for Session { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { let mut result = f.debug_struct("Session"); match &self.principal { Principal::User(user) => { result.field("user", &user.github_login); } Principal::Impersonated { user, admin } => { result.field("user", &user.github_login); result.field("impersonator", &admin.github_login); } } result.field("connection_id", &self.connection_id).finish() }} struct DbHandle(Arc<Database>); impl Deref for DbHandle { type Target = Database; fn deref(&self) -> &Self::Target { self.0.as_ref() }} pub struct Server { id: parking_lot::Mutex<ServerId>, peer: Arc<Peer>, pub(crate) connection_pool: Arc<parking_lot::Mutex<ConnectionPool>>, app_state: Arc<AppState>, handlers: HashMap<TypeId, MessageHandler>, teardown: watch::Sender<bool>,} pub(crate) struct ConnectionPoolGuard<'a> { guard: parking_lot::MutexGuard<'a, ConnectionPool>, _not_send: PhantomData<Rc<()>>,} #[derive(Serialize)]pub struct ServerSnapshot<'a> { peer: &'a Peer, #[serde(serialize_with = "serialize_deref")] connection_pool: ConnectionPoolGuard<'a>,} pub fn serialize_deref<S, T, U>(value: &T, serializer: S) -> Result<S::Ok, S::Error>where S: Serializer, T: Deref<Target = U>, U: Serialize,{ Serialize::serialize(value.deref(), serializer)} impl Server { pub fn new(id: ServerId, app_state: Arc<AppState>) -> Arc<Self> { let mut server = Self { id: parking_lot::Mutex::new(id), peer: Peer::new(id.0 as u32), app_state: app_state.clone(), connection_pool: Default::default(), handlers: Default::default(), teardown: watch::channel(false).0, }; server .add_request_handler(ping) .add_request_handler(create_room) .add_request_handler(join_room) .add_request_handler(rejoin_room) .add_request_handler(leave_room) .add_request_handler(set_room_participant_role) .add_request_handler(call) .add_request_handler(cancel_call) .add_message_handler(decline_call) .add_request_handler(update_participant_location) .add_request_handler(share_project) .add_message_handler(unshare_project) .add_request_handler(join_project) .add_message_handler(leave_project) .add_request_handler(update_project) .add_request_handler(update_worktree) .add_message_handler(start_language_server) .add_message_handler(update_language_server) .add_message_handler(update_diagnostic_summary) .add_message_handler(update_worktree_settings) .add_request_handler(forward_read_only_project_request::<proto::GetHover>) .add_request_handler(forward_read_only_project_request::<proto::GetDefinition>) .add_request_handler(forward_read_only_project_request::<proto::GetTypeDefinition>) .add_request_handler(forward_read_only_project_request::<proto::GetReferences>) .add_request_handler(forward_find_search_candidates_request) .add_request_handler(forward_read_only_project_request::<proto::GetDocumentHighlights>) .add_request_handler(forward_read_only_project_request::<proto::GetProjectSymbols>) .add_request_handler(forward_read_only_project_request::<proto::OpenBufferForSymbol>) .add_request_handler(forward_read_only_project_request::<proto::OpenBufferById>) .add_request_handler(forward_read_only_project_request::<proto::SynchronizeBuffers>) .add_request_handler(forward_read_only_project_request::<proto::InlayHints>) .add_request_handler(forward_read_only_project_request::<proto::ResolveInlayHint>) .add_request_handler(forward_read_only_project_request::<proto::OpenBufferByPath>) .add_request_handler(forward_read_only_project_request::<proto::GitBranches>) .add_request_handler(forward_read_only_project_request::<proto::OpenUnstagedDiff>) .add_request_handler(forward_read_only_project_request::<proto::OpenUncommittedDiff>) .add_request_handler( forward_mutating_project_request::<proto::RegisterBufferWithLanguageServers>, ) .add_request_handler(forward_mutating_project_request::<proto::UpdateGitBranch>) .add_request_handler(forward_mutating_project_request::<proto::GetCompletions>) .add_request_handler( forward_mutating_project_request::<proto::ApplyCompletionAdditionalEdits>, ) .add_request_handler(forward_mutating_project_request::<proto::OpenNewBuffer>) .add_request_handler( forward_mutating_project_request::<proto::ResolveCompletionDocumentation>, ) .add_request_handler(forward_mutating_project_request::<proto::GetCodeActions>) .add_request_handler(forward_mutating_project_request::<proto::ApplyCodeAction>) .add_request_handler(forward_mutating_project_request::<proto::PrepareRename>) .add_request_handler(forward_mutating_project_request::<proto::PerformRename>) .add_request_handler(forward_mutating_project_request::<proto::ReloadBuffers>) .add_request_handler(forward_mutating_project_request::<proto::FormatBuffers>) .add_request_handler(forward_mutating_project_request::<proto::CreateProjectEntry>) .add_request_handler(forward_mutating_project_request::<proto::RenameProjectEntry>) .add_request_handler(forward_mutating_project_request::<proto::CopyProjectEntry>) .add_request_handler(forward_mutating_project_request::<proto::DeleteProjectEntry>) .add_request_handler(forward_mutating_project_request::<proto::ExpandProjectEntry>) .add_request_handler( forward_mutating_project_request::<proto::ExpandAllForProjectEntry>, ) .add_request_handler(forward_mutating_project_request::<proto::OnTypeFormatting>) .add_request_handler(forward_mutating_project_request::<proto::SaveBuffer>) .add_request_handler(forward_mutating_project_request::<proto::BlameBuffer>) .add_request_handler(forward_mutating_project_request::<proto::MultiLspQuery>) .add_request_handler(forward_mutating_project_request::<proto::RestartLanguageServers>) .add_request_handler(forward_mutating_project_request::<proto::LinkedEditingRange>) .add_message_handler(create_buffer_for_peer) .add_request_handler(update_buffer) .add_message_handler(broadcast_project_message_from_host::<proto::RefreshInlayHints>) .add_message_handler(broadcast_project_message_from_host::<proto::UpdateBufferFile>) .add_message_handler(broadcast_project_message_from_host::<proto::BufferReloaded>) .add_message_handler(broadcast_project_message_from_host::<proto::BufferSaved>) .add_message_handler(broadcast_project_message_from_host::<proto::UpdateDiffBases>) .add_request_handler(get_users) .add_request_handler(fuzzy_search_users) .add_request_handler(request_contact) .add_request_handler(remove_contact) .add_request_handler(respond_to_contact_request) .add_message_handler(subscribe_to_channels) .add_request_handler(create_channel) .add_request_handler(delete_channel) .add_request_handler(invite_channel_member) .add_request_handler(remove_channel_member) .add_request_handler(set_channel_member_role) .add_request_handler(set_channel_visibility) .add_request_handler(rename_channel) .add_request_handler(join_channel_buffer) .add_request_handler(leave_channel_buffer) .add_message_handler(update_channel_buffer) .add_request_handler(rejoin_channel_buffers) .add_request_handler(get_channel_members) .add_request_handler(respond_to_channel_invite) .add_request_handler(join_channel) .add_request_handler(join_channel_chat) .add_message_handler(leave_channel_chat) .add_request_handler(send_channel_message) .add_request_handler(remove_channel_message) .add_request_handler(update_channel_message) .add_request_handler(get_channel_messages) .add_request_handler(get_channel_messages_by_id) .add_request_handler(get_notifications) .add_request_handler(mark_notification_as_read) .add_request_handler(move_channel) .add_request_handler(follow) .add_message_handler(unfollow) .add_message_handler(update_followers) .add_request_handler(get_private_user_info) .add_request_handler(get_llm_api_token) .add_request_handler(accept_terms_of_service) .add_message_handler(acknowledge_channel_message) .add_message_handler(acknowledge_buffer_version) .add_request_handler(get_supermaven_api_key) .add_request_handler(forward_mutating_project_request::<proto::OpenContext>) .add_request_handler(forward_mutating_project_request::<proto::CreateContext>) .add_request_handler(forward_mutating_project_request::<proto::SynchronizeContexts>) .add_request_handler(forward_mutating_project_request::<proto::Push>) .add_request_handler(forward_mutating_project_request::<proto::Pull>) .add_request_handler(forward_mutating_project_request::<proto::Fetch>) .add_request_handler(forward_mutating_project_request::<proto::Stage>) .add_request_handler(forward_mutating_project_request::<proto::Unstage>) .add_request_handler(forward_mutating_project_request::<proto::Commit>) .add_request_handler(forward_read_only_project_request::<proto::GetRemotes>) .add_request_handler(forward_read_only_project_request::<proto::GitShow>) .add_request_handler(forward_read_only_project_request::<proto::GitReset>) .add_request_handler(forward_read_only_project_request::<proto::GitCheckoutFiles>) .add_request_handler(forward_mutating_project_request::<proto::SetIndexText>) .add_request_handler(forward_mutating_project_request::<proto::OpenCommitMessageBuffer>) .add_message_handler(broadcast_project_message_from_host::<proto::AdvertiseContexts>) .add_message_handler(update_context) .add_request_handler({ let app_state = app_state.clone(); move |request, response, session| { let app_state = app_state.clone(); async move { count_language_model_tokens(request, response, session, &app_state.config) .await } } }) .add_request_handler(get_cached_embeddings) .add_request_handler({ let app_state = app_state.clone(); move |request, response, session| { compute_embeddings( request, response, session, app_state.config.openai_api_key.clone(), ) } }); Arc::new(server) } pub async fn start(&self) -> Result<()> { let server_id = *self.id.lock(); let app_state = self.app_state.clone(); let peer = self.peer.clone(); let timeout = self.app_state.executor.sleep(CLEANUP_TIMEOUT); let pool = self.connection_pool.clone(); let livekit_client = self.app_state.livekit_client.clone(); let span = info_span!("start server"); self.app_state.executor.spawn_detached( async move { tracing::info!("waiting for cleanup timeout"); timeout.await; tracing::info!("cleanup timeout expired, retrieving stale rooms"); if let Some((room_ids, channel_ids)) = app_state .db .stale_server_resource_ids(&app_state.config.zed_environment, server_id) .await .trace_err() { tracing::info!(stale_room_count = room_ids.len(), "retrieved stale rooms"); tracing::info!( stale_channel_buffer_count = channel_ids.len(), "retrieved stale channel buffers" ); for channel_id in channel_ids { if let Some(refreshed_channel_buffer) = app_state .db .clear_stale_channel_buffer_collaborators(channel_id, server_id) .await .trace_err() { for connection_id in refreshed_channel_buffer.connection_ids { peer.send( connection_id, proto::UpdateChannelBufferCollaborators { channel_id: channel_id.to_proto(), collaborators: refreshed_channel_buffer .collaborators .clone(), }, ) .trace_err(); } } } for room_id in room_ids { let mut contacts_to_update = HashSet::default(); let mut canceled_calls_to_user_ids = Vec::new(); let mut livekit_room = String::new(); let mut delete_livekit_room = false; if let Some(mut refreshed_room) = app_state .db .clear_stale_room_participants(room_id, server_id) .await .trace_err() { tracing::info!( room_id = room_id.0, new_participant_count = refreshed_room.room.participants.len(), "refreshed room" ); room_updated(&refreshed_room.room, &peer);Lines 1–500 of 4482. Load more
mod connection_pool; use crate::api::{CloudflareIpCountryHeader, SystemIdHeader}; use crate::llm::LlmTokenClaims; use crate::{ auth, db::{ self, BufferId, Capability, Channel, ChannelId, ChannelRole, ChannelsForUser, CreatedChannelMessage, Database, InviteMemberResult, MembershipUpdated, MessageId, NotificationId, Project, ProjectId, RejoinedProject, RemoveChannelMemberResult, ReplicaId, RespondToChannelInvite, RoomId, ServerId, UpdatedChannelMessage, User, UserId, }, executor::Executor, AppState, Config, Error, RateLimit, Result, }; use anyhow::{anyhow, bail, Context as _}; use async_tungstenite::tungstenite::{ protocol::CloseFrame as TungsteniteCloseFrame, Message as TungsteniteMessage, }; use axum::{ body::Body, extract::{ ws::{CloseFrame as AxumCloseFrame, Message as AxumMessage}, ConnectInfo, WebSocketUpgrade, }, headers::{Header, HeaderName}, http::StatusCode, middleware, response::IntoResponse, routing::get, Extension, Router, TypedHeader, }; use chrono::Utc; use collections::{HashMap, HashSet}; pub use connection_pool::{ConnectionPool, ZedVersion}; use core::fmt::{self, Debug, Formatter}; use http_client::HttpClient; use open_ai::{OpenAiEmbeddingModel, OPEN_AI_API_URL}; use reqwest_client::ReqwestClient; use sha2::Digest; use supermaven_api::{CreateExternalUserRequest, SupermavenAdminApi}; use futures::{ channel::oneshot, future::BoxFuture, stream::FuturesUnordered, FutureExt, SinkExt, StreamExt, TryStreamExt, }; use prometheus::{register_int_gauge, IntGauge}; use rpc::{ proto::{ self, Ack, AnyTypedEnvelope, EntityMessage, EnvelopedMessage, LiveKitConnectionInfo, RequestMessage, ShareProject, UpdateChannelBufferCollaborators, }, Connection, ConnectionId, ErrorCode, ErrorCodeExt, ErrorExt, Peer, Receipt, TypedEnvelope, }; use semantic_version::SemanticVersion; use serde::{Serialize, Serializer}; use std::{ any::TypeId, future::Future, marker::PhantomData, mem, net::SocketAddr, ops::{Deref, DerefMut}, rc::Rc, sync::{ atomic::{AtomicBool, Ordering::SeqCst}, Arc, OnceLock, }, time::{Duration, Instant}, }; use time::OffsetDateTime; use tokio::sync::{watch, MutexGuard, Semaphore}; use tower::ServiceBuilder; use tracing::{ field::{self}, info_span, instrument, Instrument, }; pub const RECONNECT_TIMEOUT: Duration = Duration::from_secs(30); // kubernetes gives terminated pods 10s to shutdown gracefully. After they're gone, we can clean up old resources. pub const CLEANUP_TIMEOUT: Duration = Duration::from_secs(15); const MESSAGE_COUNT_PER_PAGE: usize = 100; const MAX_MESSAGE_LEN: usize = 1024; const NOTIFICATION_COUNT_PER_PAGE: usize = 50; type MessageHandler = Box<dyn Send + Sync + Fn(Box<dyn AnyTypedEnvelope>, Session) -> BoxFuture<'static, ()>>; struct Response<R> { peer: Arc<Peer>, receipt: Receipt<R>, responded: Arc<AtomicBool>, } impl<R: RequestMessage> Response<R> { fn send(self, payload: R::Response) -> Result<()> { self.responded.store(true, SeqCst); self.peer.respond(self.receipt, payload)?; Ok(()) } } #[derive(Clone, Debug)] pub enum Principal { User(User), Impersonated { user: User, admin: User }, } impl Principal { fn update_span(&self, span: &tracing::Span) { match &self { Principal::User(user) => { span.record("user_id", user.id.0); span.record("login", &user.github_login); } Principal::Impersonated { user, admin } => { span.record("user_id", user.id.0); span.record("login", &user.github_login); span.record("impersonator", &admin.github_login); } } } } #[derive(Clone)] struct Session { principal: Principal, connection_id: ConnectionId, db: Arc<tokio::sync::Mutex<DbHandle>>, peer: Arc<Peer>, connection_pool: Arc<parking_lot::Mutex<ConnectionPool>>, app_state: Arc<AppState>, supermaven_client: Option<Arc<SupermavenAdminApi>>, http_client: Arc<dyn HttpClient>, /// The GeoIP country code for the user. #[allow(unused)] geoip_country_code: Option<String>, system_id: Option<String>, _executor: Executor, } impl Session { async fn db(&self) -> tokio::sync::MutexGuard<DbHandle> { #[cfg(test)] tokio::task::yield_now().await; let guard = self.db.lock().await; #[cfg(test)] tokio::task::yield_now().await; guard } async fn connection_pool(&self) -> ConnectionPoolGuard<'_> { #[cfg(test)] tokio::task::yield_now().await; let guard = self.connection_pool.lock(); ConnectionPoolGuard { guard, _not_send: PhantomData, } } fn is_staff(&self) -> bool { match &self.principal { Principal::User(user) => user.admin, Principal::Impersonated { .. } => true, } } pub async fn has_llm_subscription( &self, db: &MutexGuard<'_, DbHandle>, ) -> anyhow::Result<bool> { if self.is_staff() { return Ok(true); } let user_id = self.user_id(); Ok(db.has_active_billing_subscription(user_id).await?) } pub async fn current_plan( &self, _db: &MutexGuard<'_, DbHandle>, ) -> anyhow::Result<proto::Plan> { if self.is_staff() { Ok(proto::Plan::ZedPro) } else { Ok(proto::Plan::Free) } } fn user_id(&self) -> UserId { match &self.principal { Principal::User(user) => user.id, Principal::Impersonated { user, .. } => user.id, } } pub fn email(&self) -> Option<String> { match &self.principal { Principal::User(user) => user.email_address.clone(), Principal::Impersonated { user, .. } => user.email_address.clone(), } } } impl Debug for Session { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { let mut result = f.debug_struct("Session"); match &self.principal { Principal::User(user) => { result.field("user", &user.github_login); } Principal::Impersonated { user, admin } => { result.field("user", &user.github_login); result.field("impersonator", &admin.github_login); } } result.field("connection_id", &self.connection_id).finish() } } struct DbHandle(Arc<Database>); impl Deref for DbHandle { type Target = Database; fn deref(&self) -> &Self::Target { self.0.as_ref() } } pub struct Server { id: parking_lot::Mutex<ServerId>, peer: Arc<Peer>, pub(crate) connection_pool: Arc<parking_lot::Mutex<ConnectionPool>>, app_state: Arc<AppState>, handlers: HashMap<TypeId, MessageHandler>, teardown: watch::Sender<bool>, } pub(crate) struct ConnectionPoolGuard<'a> { guard: parking_lot::MutexGuard<'a, ConnectionPool>, _not_send: PhantomData<Rc<()>>, } #[derive(Serialize)] pub struct ServerSnapshot<'a> { peer: &'a Peer, #[serde(serialize_with = "serialize_deref")] connection_pool: ConnectionPoolGuard<'a>, } pub fn serialize_deref<S, T, U>(value: &T, serializer: S) -> Result<S::Ok, S::Error> where S: Serializer, T: Deref<Target = U>, U: Serialize, { Serialize::serialize(value.deref(), serializer) } impl Server { pub fn new(id: ServerId, app_state: Arc<AppState>) -> Arc<Self> { let mut server = Self { id: parking_lot::Mutex::new(id), peer: Peer::new(id.0 as u32), app_state: app_state.clone(), connection_pool: Default::default(), handlers: Default::default(), teardown: watch::channel(false).0, }; server .add_request_handler(ping) .add_request_handler(create_room) .add_request_handler(join_room) .add_request_handler(rejoin_room) .add_request_handler(leave_room) .add_request_handler(set_room_participant_role) .add_request_handler(call) .add_request_handler(cancel_call) .add_message_handler(decline_call) .add_request_handler(update_participant_location) .add_request_handler(share_project) .add_message_handler(unshare_project) .add_request_handler(join_project) .add_message_handler(leave_project) .add_request_handler(update_project) .add_request_handler(update_worktree) .add_message_handler(start_language_server) .add_message_handler(update_language_server) .add_message_handler(update_diagnostic_summary) .add_message_handler(update_worktree_settings) .add_request_handler(forward_read_only_project_request::<proto::GetHover>) .add_request_handler(forward_read_only_project_request::<proto::GetDefinition>) .add_request_handler(forward_read_only_project_request::<proto::GetTypeDefinition>) .add_request_handler(forward_read_only_project_request::<proto::GetReferences>) .add_request_handler(forward_find_search_candidates_request) .add_request_handler(forward_read_only_project_request::<proto::GetDocumentHighlights>) .add_request_handler(forward_read_only_project_request::<proto::GetProjectSymbols>) .add_request_handler(forward_read_only_project_request::<proto::OpenBufferForSymbol>) .add_request_handler(forward_read_only_project_request::<proto::OpenBufferById>) .add_request_handler(forward_read_only_project_request::<proto::SynchronizeBuffers>) .add_request_handler(forward_read_only_project_request::<proto::InlayHints>) .add_request_handler(forward_read_only_project_request::<proto::ResolveInlayHint>) .add_request_handler(forward_read_only_project_request::<proto::OpenBufferByPath>) .add_request_handler(forward_read_only_project_request::<proto::GitBranches>) .add_request_handler(forward_read_only_project_request::<proto::OpenUnstagedDiff>) .add_request_handler(forward_read_only_project_request::<proto::OpenUncommittedDiff>) .add_request_handler( forward_mutating_project_request::<proto::RegisterBufferWithLanguageServers>, ) .add_request_handler(forward_mutating_project_request::<proto::UpdateGitBranch>) .add_request_handler(forward_mutating_project_request::<proto::GetCompletions>) .add_request_handler( forward_mutating_project_request::<proto::ApplyCompletionAdditionalEdits>, ) .add_request_handler(forward_mutating_project_request::<proto::OpenNewBuffer>) .add_request_handler( forward_mutating_project_request::<proto::ResolveCompletionDocumentation>, ) .add_request_handler(forward_mutating_project_request::<proto::GetCodeActions>) .add_request_handler(forward_mutating_project_request::<proto::ApplyCodeAction>) .add_request_handler(forward_mutating_project_request::<proto::PrepareRename>) .add_request_handler(forward_mutating_project_request::<proto::PerformRename>) .add_request_handler(forward_mutating_project_request::<proto::ReloadBuffers>) .add_request_handler(forward_mutating_project_request::<proto::FormatBuffers>) .add_request_handler(forward_mutating_project_request::<proto::CreateProjectEntry>) .add_request_handler(forward_mutating_project_request::<proto::RenameProjectEntry>) .add_request_handler(forward_mutating_project_request::<proto::CopyProjectEntry>) .add_request_handler(forward_mutating_project_request::<proto::DeleteProjectEntry>) .add_request_handler(forward_mutating_project_request::<proto::ExpandProjectEntry>) .add_request_handler( forward_mutating_project_request::<proto::ExpandAllForProjectEntry>, ) .add_request_handler(forward_mutating_project_request::<proto::OnTypeFormatting>) .add_request_handler(forward_mutating_project_request::<proto::SaveBuffer>) .add_request_handler(forward_mutating_project_request::<proto::BlameBuffer>) .add_request_handler(forward_mutating_project_request::<proto::MultiLspQuery>) .add_request_handler(forward_mutating_project_request::<proto::RestartLanguageServers>) .add_request_handler(forward_mutating_project_request::<proto::LinkedEditingRange>) .add_message_handler(create_buffer_for_peer) .add_request_handler(update_buffer) .add_message_handler(broadcast_project_message_from_host::<proto::RefreshInlayHints>) .add_message_handler(broadcast_project_message_from_host::<proto::UpdateBufferFile>) .add_message_handler(broadcast_project_message_from_host::<proto::BufferReloaded>) .add_message_handler(broadcast_project_message_from_host::<proto::BufferSaved>) .add_message_handler(broadcast_project_message_from_host::<proto::UpdateDiffBases>) .add_request_handler(get_users) .add_request_handler(fuzzy_search_users) .add_request_handler(request_contact) .add_request_handler(remove_contact) .add_request_handler(respond_to_contact_request) .add_message_handler(subscribe_to_channels) .add_request_handler(create_channel) .add_request_handler(delete_channel) .add_request_handler(invite_channel_member) .add_request_handler(remove_channel_member) .add_request_handler(set_channel_member_role) .add_request_handler(set_channel_visibility) .add_request_handler(rename_channel) .add_request_handler(join_channel_buffer) .add_request_handler(leave_channel_buffer) .add_message_handler(update_channel_buffer) .add_request_handler(rejoin_channel_buffers) .add_request_handler(get_channel_members) .add_request_handler(respond_to_channel_invite) .add_request_handler(join_channel) .add_request_handler(join_channel_chat) .add_message_handler(leave_channel_chat) .add_request_handler(send_channel_message) .add_request_handler(remove_channel_message) .add_request_handler(update_channel_message) .add_request_handler(get_channel_messages) .add_request_handler(get_channel_messages_by_id) .add_request_handler(get_notifications) .add_request_handler(mark_notification_as_read) .add_request_handler(move_channel) .add_request_handler(follow) .add_message_handler(unfollow) .add_message_handler(update_followers) .add_request_handler(get_private_user_info) .add_request_handler(get_llm_api_token) .add_request_handler(accept_terms_of_service) .add_message_handler(acknowledge_channel_message) .add_message_handler(acknowledge_buffer_version) .add_request_handler(get_supermaven_api_key) .add_request_handler(forward_mutating_project_request::<proto::OpenContext>) .add_request_handler(forward_mutating_project_request::<proto::CreateContext>) .add_request_handler(forward_mutating_project_request::<proto::SynchronizeContexts>) .add_request_handler(forward_mutating_project_request::<proto::Push>) .add_request_handler(forward_mutating_project_request::<proto::Pull>) .add_request_handler(forward_mutating_project_request::<proto::Fetch>) .add_request_handler(forward_mutating_project_request::<proto::Stage>) .add_request_handler(forward_mutating_project_request::<proto::Unstage>) .add_request_handler(forward_mutating_project_request::<proto::Commit>) .add_request_handler(forward_read_only_project_request::<proto::GetRemotes>) .add_request_handler(forward_read_only_project_request::<proto::GitShow>) .add_request_handler(forward_read_only_project_request::<proto::GitReset>) .add_request_handler(forward_read_only_project_request::<proto::GitCheckoutFiles>) .add_request_handler(forward_mutating_project_request::<proto::SetIndexText>) .add_request_handler(forward_mutating_project_request::<proto::OpenCommitMessageBuffer>) .add_message_handler(broadcast_project_message_from_host::<proto::AdvertiseContexts>) .add_message_handler(update_context) .add_request_handler({ let app_state = app_state.clone(); move |request, response, session| { let app_state = app_state.clone(); async move { count_language_model_tokens(request, response, session, &app_state.config) .await } } }) .add_request_handler(get_cached_embeddings) .add_request_handler({ let app_state = app_state.clone(); move |request, response, session| { compute_embeddings( request, response, session, app_state.config.openai_api_key.clone(), ) } }); Arc::new(server) } pub async fn start(&self) -> Result<()> { let server_id = *self.id.lock(); let app_state = self.app_state.clone(); let peer = self.peer.clone(); let timeout = self.app_state.executor.sleep(CLEANUP_TIMEOUT); let pool = self.connection_pool.clone(); let livekit_client = self.app_state.livekit_client.clone(); let span = info_span!("start server"); self.app_state.executor.spawn_detached( async move { tracing::info!("waiting for cleanup timeout"); timeout.await; tracing::info!("cleanup timeout expired, retrieving stale rooms"); if let Some((room_ids, channel_ids)) = app_state .db .stale_server_resource_ids(&app_state.config.zed_environment, server_id) .await .trace_err() { tracing::info!(stale_room_count = room_ids.len(), "retrieved stale rooms"); tracing::info!( stale_channel_buffer_count = channel_ids.len(), "retrieved stale channel buffers" ); for channel_id in channel_ids { if let Some(refreshed_channel_buffer) = app_state .db .clear_stale_channel_buffer_collaborators(channel_id, server_id) .await .trace_err() { for connection_id in refreshed_channel_buffer.connection_ids { peer.send( connection_id, proto::UpdateChannelBufferCollaborators { channel_id: channel_id.to_proto(), collaborators: refreshed_channel_buffer .collaborators .clone(), }, ) .trace_err(); } } } for room_id in room_ids { let mut contacts_to_update = HashSet::default(); let mut canceled_calls_to_user_ids = Vec::new(); let mut livekit_room = String::new(); let mut delete_livekit_room = false; if let Some(mut refreshed_room) = app_state .db .clear_stale_room_participants(room_id, server_id) .await .trace_err() { tracing::info!( room_id = room_id.0, new_participant_count = refreshed_room.room.participants.len(), "refreshed room" ); room_updated(&refreshed_room.room, &peer);