mod connection_pool; use crate::api::CloudflareIpCountryHeader;use crate::llm::LlmTokenClaims;use crate::{    auth,    db::{        self, dev_server, BufferId, Capability, Channel, ChannelId, ChannelRole, ChannelsForUser,        CreatedChannelMessage, Database, DevServerId, DevServerProjectId, InviteMemberResult,        MembershipUpdated, MessageId, NotificationId, PrincipalId, 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::{self, 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 },    DevServer(dev_server::Model),} 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);            }            Principal::DevServer(dev_server) => {                span.record("dev_server_id", dev_server.id.0);            }        }    }} #[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>,    _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 for_user(self) -> Option<UserSession> {        UserSession::new(self)    }     fn for_dev_server(self) -> Option<DevServerSession> {        DevServerSession::new(self)    }     fn user_id(&self) -> Option<UserId> {        match &self.principal {            Principal::User(user) => Some(user.id),            Principal::Impersonated { user, .. } => Some(user.id),            Principal::DevServer(_) => None,        }    }     fn is_staff(&self) -> bool {        match &self.principal {            Principal::User(user) => user.admin,            Principal::Impersonated { .. } => true,            Principal::DevServer(_) => false,        }    }     pub async fn has_llm_subscription(        &self,        db: &MutexGuard<'_, DbHandle>,    ) -> anyhow::Result<bool> {        if self.is_staff() {            return Ok(true);        }         let Some(user_id) = self.user_id() else {            return Ok(false);        };         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 dev_server_id(&self) -> Option<DevServerId> {        match &self.principal {            Principal::User(_) | Principal::Impersonated { .. } => None,            Principal::DevServer(dev_server) => Some(dev_server.id),        }    }     fn principal_id(&self) -> PrincipalId {        match &self.principal {            Principal::User(user) => PrincipalId::UserId(user.id),            Principal::Impersonated { user, .. } => PrincipalId::UserId(user.id),            Principal::DevServer(dev_server) => PrincipalId::DevServerId(dev_server.id),        }    }} 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);            }            Principal::DevServer(dev_server) => {                result.field("dev_server", &dev_server.id);            }        }        result.field("connection_id", &self.connection_id).finish()    }} struct UserSession(Session); impl UserSession {    pub fn new(s: Session) -> Option<Self> {        s.user_id().map(|_| UserSession(s))    }    pub fn user_id(&self) -> UserId {        self.0.user_id().unwrap()    }     pub fn email(&self) -> Option<String> {        match &self.0.principal {            Principal::User(user) => user.email_address.clone(),            Principal::Impersonated { user, .. } => user.email_address.clone(),            Principal::DevServer(..) => None,        }    }} impl Deref for UserSession {    type Target = Session;     fn deref(&self) -> &Self::Target {        &self.0    }}impl DerefMut for UserSession {    fn deref_mut(&mut self) -> &mut Self::Target {        &mut self.0    }} struct DevServerSession(Session); impl DevServerSession {    pub fn new(s: Session) -> Option<Self> {        s.dev_server_id().map(|_| DevServerSession(s))    }    pub fn dev_server_id(&self) -> DevServerId {        self.0.dev_server_id().unwrap()    }     fn dev_server(&self) -> &dev_server::Model {        match &self.0.principal {            Principal::DevServer(dev_server) => dev_server,            _ => unreachable!(),        }    }} impl Deref for DevServerSession {    type Target = Session;     fn deref(&self) -> &Self::Target {        &self.0    }}impl DerefMut for DevServerSession {    fn deref_mut(&mut self) -> &mut Self::Target {        &mut self.0    }} fn user_handler<M: RequestMessage, Fut>(    handler: impl 'static + Send + Sync + Fn(M, Response<M>, UserSession) -> Fut,) -> impl 'static + Send + Sync + Fn(M, Response<M>, Session) -> BoxFuture<'static, Result<()>>where    Fut: Send + Future<Output = Result<()>>,{    let handler = Arc::new(handler);    move |message, response, session| {        let handler = handler.clone();        Box::pin(async move {            if let Some(user_session) = session.for_user() {                Ok(handler(message, response, user_session).await?)            } else {                Err(Error::Internal(anyhow!(                    "must be a user to call {}",                    M::NAME                )))            }        })    }} fn dev_server_handler<M: RequestMessage, Fut>(    handler: impl 'static + Send + Sync + Fn(M, Response<M>, DevServerSession) -> Fut,) -> impl 'static + Send + Sync + Fn(M, Response<M>, Session) -> BoxFuture<'static, Result<()>>where    Fut: Send + Future<Output = Result<()>>,{    let handler = Arc::new(handler);    move |message, response, session| {        let handler = handler.clone();        Box::pin(async move {            if let Some(dev_server_session) = session.for_dev_server() {                Ok(handler(message, response, dev_server_session).await?)            } else {                Err(Error::Internal(anyhow!(                    "must be a dev server to call {}",                    M::NAME                )))            }        })    }} fn user_message_handler<M: EnvelopedMessage, InnertRetFut>(    handler: impl 'static + Send + Sync + Fn(M, UserSession) -> InnertRetFut,) -> impl 'static + Send + Sync + Fn(M, Session) -> BoxFuture<'static, Result<()>>where    InnertRetFut: Send + Future<Output = Result<()>>,{    let handler = Arc::new(handler);    move |message, session| {        let handler = handler.clone();        Box::pin(async move {            if let Some(user_session) = session.for_user() {                Ok(handler(message, user_session).await?)            } else {                Err(Error::Internal(anyhow!(                    "must be a user to call {}",                    M::NAME                )))            }        })    }} 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(user_handler(create_room))            .add_request_handler(user_handler(join_room))            .add_request_handler(user_handler(rejoin_room))            .add_request_handler(user_handler(leave_room))            .add_request_handler(user_handler(set_room_participant_role))            .add_request_handler(user_handler(call))            .add_request_handler(user_handler(cancel_call))            .add_message_handler(user_message_handler(decline_call))            .add_request_handler(user_handler(update_participant_location))            .add_request_handler(user_handler(share_project))            .add_message_handler(unshare_project)            .add_request_handler(user_handler(join_project))            .add_request_handler(user_handler(join_hosted_project))            .add_request_handler(user_handler(rejoin_dev_server_projects))            .add_request_handler(user_handler(create_dev_server_project))            .add_request_handler(user_handler(update_dev_server_project))            .add_request_handler(user_handler(delete_dev_server_project))            .add_request_handler(user_handler(create_dev_server))            .add_request_handler(user_handler(regenerate_dev_server_token))            .add_request_handler(user_handler(rename_dev_server))            .add_request_handler(user_handler(delete_dev_server))            .add_request_handler(user_handler(list_remote_directory))            .add_request_handler(dev_server_handler(share_dev_server_project))            .add_request_handler(dev_server_handler(shutdown_dev_server))            .add_request_handler(dev_server_handler(reconnect_dev_server))            .add_message_handler(user_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(user_handler(                forward_project_request_for_owner::<proto::TaskContextForLocation>,            ))            .add_request_handler(user_handler(                forward_read_only_project_request::<proto::GetHover>,            ))            .add_request_handler(user_handler(                forward_read_only_project_request::<proto::GetDefinition>,            ))            .add_request_handler(user_handler(                forward_read_only_project_request::<proto::GetTypeDefinition>,            ))            .add_request_handler(user_handler(                forward_read_only_project_request::<proto::GetReferences>,            ))            .add_request_handler(user_handler(forward_find_search_candidates_request))            .add_request_handler(user_handler(                forward_read_only_project_request::<proto::GetDocumentHighlights>,            ))            .add_request_handler(user_handler(                forward_read_only_project_request::<proto::GetProjectSymbols>,            ))            .add_request_handler(user_handler(                forward_read_only_project_request::<proto::OpenBufferForSymbol>,            ))            .add_request_handler(user_handler(                forward_read_only_project_request::<proto::OpenBufferById>,            ))            .add_request_handler(user_handler(                forward_read_only_project_request::<proto::SynchronizeBuffers>,            ))            .add_request_handler(user_handler(

Lines 1–500 of 5374. Load more