mod connection_pool; use crate::api::billing::find_or_create_billing_customer;use crate::api::{CloudflareIpCountryHeader, SystemIdHeader};use crate::db::billing_subscription::SubscriptionKind;use crate::llm::db::LlmDatabase;use crate::llm::{ AGENT_EXTENDED_TRIAL_FEATURE_FLAG, BYPASS_ACCOUNT_AGE_CHECK_FEATURE_FLAG, LlmTokenClaims, MIN_ACCOUNT_AGE_FOR_LLM_USE,};use crate::stripe_client::StripeCustomerId;use crate::{ AppState, Error, Result, auth, db::{ self, BufferId, Capability, Channel, ChannelId, ChannelRole, ChannelsForUser, CreatedChannelMessage, Database, InviteMemberResult, MembershipUpdated, MessageId, NotificationId, ProjectId, RejoinedProject, RemoveChannelMemberResult, RespondToChannelInvite, RoomId, ServerId, UpdatedChannelMessage, User, UserId, }, executor::Executor,};use anyhow::{Context as _, anyhow, bail};use async_tungstenite::tungstenite::{ Message as TungsteniteMessage, protocol::CloseFrame as TungsteniteCloseFrame,};use axum::{ Extension, Router, TypedHeader, body::Body, extract::{ ConnectInfo, WebSocketUpgrade, ws::{CloseFrame as AxumCloseFrame, Message as AxumMessage}, }, headers::{Header, HeaderName}, http::StatusCode, middleware, response::IntoResponse, routing::get,};use chrono::Utc;use collections::{HashMap, HashSet};pub use connection_pool::{ConnectionPool, ZedVersion};use core::fmt::{self, Debug, Formatter};use reqwest_client::ReqwestClient;use rpc::proto::split_repository_update;use supermaven_api::{CreateExternalUserRequest, SupermavenAdminApi}; use futures::{ FutureExt, SinkExt, StreamExt, TryStreamExt, channel::oneshot, future::BoxFuture, stream::FuturesUnordered,};use prometheus::{IntGauge, register_int_gauge};use rpc::{ Connection, ConnectionId, ErrorCode, ErrorCodeExt, ErrorExt, Peer, Receipt, TypedEnvelope, proto::{ self, Ack, AnyTypedEnvelope, EntityMessage, EnvelopedMessage, LiveKitConnectionInfo, RequestMessage, ShareProject, UpdateChannelBufferCollaborators, },};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::{ Arc, OnceLock, atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst}, }, time::{Duration, Instant},};use time::OffsetDateTime;use tokio::sync::{Semaphore, watch};use tower::ServiceBuilder;use tracing::{ Instrument, field::{self}, info_span, 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;const MAX_CONCURRENT_CONNECTIONS: usize = 512; static CONCURRENT_CONNECTIONS: AtomicUsize = AtomicUsize::new(0); type MessageHandler = Box<dyn Send + Sync + Fn(Box<dyn AnyTypedEnvelope>, Session) -> BoxFuture<'static, ()>>; pub struct ConnectionGuard; impl ConnectionGuard { pub fn try_acquire() -> Result<Self, ()> { let current_connections = CONCURRENT_CONNECTIONS.fetch_add(1, SeqCst); if current_connections >= MAX_CONCURRENT_CONNECTIONS { CONCURRENT_CONNECTIONS.fetch_sub(1, SeqCst); tracing::error!( "too many concurrent connections: {}", current_connections + 1 ); return Err(()); } Ok(ConnectionGuard) }} impl Drop for ConnectionGuard { fn drop(&mut self) { CONCURRENT_CONNECTIONS.fetch_sub(1, SeqCst); }} 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 user(&self) -> &User { match self { Principal::User(user) => user, Principal::Impersonated { user, .. } => user, } } 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>>, /// 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, } } 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_request_handler(update_repository) .add_request_handler(remove_repository) .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::GetDocumentSymbols>) .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::GetColorPresentation>) .add_request_handler(forward_mutating_project_request::<proto::GetCodeLens>) .add_request_handler(forward_read_only_project_request::<proto::OpenBufferByPath>) .add_request_handler(forward_read_only_project_request::<proto::GitGetBranches>) .add_request_handler(forward_read_only_project_request::<proto::OpenUnstagedDiff>) .add_request_handler(forward_read_only_project_request::<proto::OpenUncommittedDiff>) .add_request_handler(forward_read_only_project_request::<proto::LspExtExpandMacro>) .add_request_handler(forward_read_only_project_request::<proto::LspExtOpenDocs>) .add_request_handler(forward_mutating_project_request::<proto::LspExtRunnables>) .add_request_handler( forward_read_only_project_request::<proto::LspExtSwitchSourceHeader>, ) .add_request_handler(forward_read_only_project_request::<proto::LspExtGoToParentModule>) .add_request_handler(forward_read_only_project_request::<proto::LspExtCancelFlycheck>) .add_request_handler(forward_read_only_project_request::<proto::LspExtRunFlycheck>) .add_request_handler(forward_read_only_project_request::<proto::LspExtClearFlycheck>) .add_request_handler( forward_read_only_project_request::<proto::LanguageServerIdForName>, ) .add_request_handler(forward_read_only_project_request::<proto::GetDocumentDiagnostics>) .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::ApplyCodeActionKind>) .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::StopLanguageServers>) .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::RefreshCodeLens>) .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_message_handler( broadcast_project_message_from_host::<proto::PullWorkspaceDiagnostics>, ) .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(reorder_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::Stage>) .add_request_handler(forward_mutating_project_request::<proto::Unstage>) .add_request_handler(forward_mutating_project_request::<proto::Commit>) .add_request_handler(forward_mutating_project_request::<proto::GitInit>) .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::LoadCommitDiff>) .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::ToggleBreakpoint>) .add_message_handler(broadcast_project_message_from_host::<proto::BreakpointsForFile>) .add_request_handler(forward_mutating_project_request::<proto::OpenCommitMessageBuffer>) .add_request_handler(forward_mutating_project_request::<proto::GitDiff>) .add_request_handler(forward_mutating_project_request::<proto::GitCreateBranch>) .add_request_handler(forward_mutating_project_request::<proto::GitChangeBranch>) .add_request_handler(forward_mutating_project_request::<proto::CheckForPushedCommits>) .add_message_handler(broadcast_project_message_from_host::<proto::AdvertiseContexts>) .add_message_handler(update_context); 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"); app_state .db .delete_stale_channel_chat_participants( &app_state.config.zed_environment, server_id, ) .await .trace_err(); 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 {Lines 1–500 of 4589. Load more
mod connection_pool; use crate::api::billing::find_or_create_billing_customer; use crate::api::{CloudflareIpCountryHeader, SystemIdHeader}; use crate::db::billing_subscription::SubscriptionKind; use crate::llm::db::LlmDatabase; use crate::llm::{ AGENT_EXTENDED_TRIAL_FEATURE_FLAG, BYPASS_ACCOUNT_AGE_CHECK_FEATURE_FLAG, LlmTokenClaims, MIN_ACCOUNT_AGE_FOR_LLM_USE, }; use crate::stripe_client::StripeCustomerId; use crate::{ AppState, Error, Result, auth, db::{ self, BufferId, Capability, Channel, ChannelId, ChannelRole, ChannelsForUser, CreatedChannelMessage, Database, InviteMemberResult, MembershipUpdated, MessageId, NotificationId, ProjectId, RejoinedProject, RemoveChannelMemberResult, RespondToChannelInvite, RoomId, ServerId, UpdatedChannelMessage, User, UserId, }, executor::Executor, }; use anyhow::{Context as _, anyhow, bail}; use async_tungstenite::tungstenite::{ Message as TungsteniteMessage, protocol::CloseFrame as TungsteniteCloseFrame, }; use axum::{ Extension, Router, TypedHeader, body::Body, extract::{ ConnectInfo, WebSocketUpgrade, ws::{CloseFrame as AxumCloseFrame, Message as AxumMessage}, }, headers::{Header, HeaderName}, http::StatusCode, middleware, response::IntoResponse, routing::get, }; use chrono::Utc; use collections::{HashMap, HashSet}; pub use connection_pool::{ConnectionPool, ZedVersion}; use core::fmt::{self, Debug, Formatter}; use reqwest_client::ReqwestClient; use rpc::proto::split_repository_update; use supermaven_api::{CreateExternalUserRequest, SupermavenAdminApi}; use futures::{ FutureExt, SinkExt, StreamExt, TryStreamExt, channel::oneshot, future::BoxFuture, stream::FuturesUnordered, }; use prometheus::{IntGauge, register_int_gauge}; use rpc::{ Connection, ConnectionId, ErrorCode, ErrorCodeExt, ErrorExt, Peer, Receipt, TypedEnvelope, proto::{ self, Ack, AnyTypedEnvelope, EntityMessage, EnvelopedMessage, LiveKitConnectionInfo, RequestMessage, ShareProject, UpdateChannelBufferCollaborators, }, }; 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::{ Arc, OnceLock, atomic::{AtomicBool, AtomicUsize, Ordering::SeqCst}, }, time::{Duration, Instant}, }; use time::OffsetDateTime; use tokio::sync::{Semaphore, watch}; use tower::ServiceBuilder; use tracing::{ Instrument, field::{self}, info_span, 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; const MAX_CONCURRENT_CONNECTIONS: usize = 512; static CONCURRENT_CONNECTIONS: AtomicUsize = AtomicUsize::new(0); type MessageHandler = Box<dyn Send + Sync + Fn(Box<dyn AnyTypedEnvelope>, Session) -> BoxFuture<'static, ()>>; pub struct ConnectionGuard; impl ConnectionGuard { pub fn try_acquire() -> Result<Self, ()> { let current_connections = CONCURRENT_CONNECTIONS.fetch_add(1, SeqCst); if current_connections >= MAX_CONCURRENT_CONNECTIONS { CONCURRENT_CONNECTIONS.fetch_sub(1, SeqCst); tracing::error!( "too many concurrent connections: {}", current_connections + 1 ); return Err(()); } Ok(ConnectionGuard) } } impl Drop for ConnectionGuard { fn drop(&mut self) { CONCURRENT_CONNECTIONS.fetch_sub(1, SeqCst); } } 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 user(&self) -> &User { match self { Principal::User(user) => user, Principal::Impersonated { user, .. } => user, } } 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>>, /// 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, } } 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_request_handler(update_repository) .add_request_handler(remove_repository) .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::GetDocumentSymbols>) .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::GetColorPresentation>) .add_request_handler(forward_mutating_project_request::<proto::GetCodeLens>) .add_request_handler(forward_read_only_project_request::<proto::OpenBufferByPath>) .add_request_handler(forward_read_only_project_request::<proto::GitGetBranches>) .add_request_handler(forward_read_only_project_request::<proto::OpenUnstagedDiff>) .add_request_handler(forward_read_only_project_request::<proto::OpenUncommittedDiff>) .add_request_handler(forward_read_only_project_request::<proto::LspExtExpandMacro>) .add_request_handler(forward_read_only_project_request::<proto::LspExtOpenDocs>) .add_request_handler(forward_mutating_project_request::<proto::LspExtRunnables>) .add_request_handler( forward_read_only_project_request::<proto::LspExtSwitchSourceHeader>, ) .add_request_handler(forward_read_only_project_request::<proto::LspExtGoToParentModule>) .add_request_handler(forward_read_only_project_request::<proto::LspExtCancelFlycheck>) .add_request_handler(forward_read_only_project_request::<proto::LspExtRunFlycheck>) .add_request_handler(forward_read_only_project_request::<proto::LspExtClearFlycheck>) .add_request_handler( forward_read_only_project_request::<proto::LanguageServerIdForName>, ) .add_request_handler(forward_read_only_project_request::<proto::GetDocumentDiagnostics>) .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::ApplyCodeActionKind>) .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::StopLanguageServers>) .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::RefreshCodeLens>) .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_message_handler( broadcast_project_message_from_host::<proto::PullWorkspaceDiagnostics>, ) .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(reorder_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::Stage>) .add_request_handler(forward_mutating_project_request::<proto::Unstage>) .add_request_handler(forward_mutating_project_request::<proto::Commit>) .add_request_handler(forward_mutating_project_request::<proto::GitInit>) .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::LoadCommitDiff>) .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::ToggleBreakpoint>) .add_message_handler(broadcast_project_message_from_host::<proto::BreakpointsForFile>) .add_request_handler(forward_mutating_project_request::<proto::OpenCommitMessageBuffer>) .add_request_handler(forward_mutating_project_request::<proto::GitDiff>) .add_request_handler(forward_mutating_project_request::<proto::GitCreateBranch>) .add_request_handler(forward_mutating_project_request::<proto::GitChangeBranch>) .add_request_handler(forward_mutating_project_request::<proto::CheckForPushedCommits>) .add_message_handler(broadcast_project_message_from_host::<proto::AdvertiseContexts>) .add_message_handler(update_context); 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"); app_state .db .delete_stale_channel_chat_participants( &app_state.config.zed_environment, server_id, ) .await .trace_err(); 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 {