1437 lines · 47.2 KB · Rust
RawDownload
pub mod extension;pub mod registry; use std::sync::Arc; use anyhow::{Context as _, Result};use collections::{HashMap, HashSet};use context_server::{ContextServer, ContextServerCommand, ContextServerId};use futures::{FutureExt as _, future::join_all};use gpui::{App, AsyncApp, Context, Entity, EventEmitter, Subscription, Task, WeakEntity, actions};use registry::ContextServerDescriptorRegistry;use settings::{Settings as _, SettingsStore};use util::{ResultExt as _, rel_path::RelPath}; use crate::{    Project,    project_settings::{ContextServerSettings, ProjectSettings},    worktree_store::WorktreeStore,}; pub fn init(cx: &mut App) {    extension::init(cx);} actions!(    context_server,    [        /// Restarts the context server.        Restart    ]); #[derive(Debug, Clone, PartialEq, Eq, Hash)]pub enum ContextServerStatus {    Starting,    Running,    Stopped,    Error(Arc<str>),} impl ContextServerStatus {    fn from_state(state: &ContextServerState) -> Self {        match state {            ContextServerState::Starting { .. } => ContextServerStatus::Starting,            ContextServerState::Running { .. } => ContextServerStatus::Running,            ContextServerState::Stopped { .. } => ContextServerStatus::Stopped,            ContextServerState::Error { error, .. } => ContextServerStatus::Error(error.clone()),        }    }} enum ContextServerState {    Starting {        server: Arc<ContextServer>,        configuration: Arc<ContextServerConfiguration>,        _task: Task<()>,    },    Running {        server: Arc<ContextServer>,        configuration: Arc<ContextServerConfiguration>,    },    Stopped {        server: Arc<ContextServer>,        configuration: Arc<ContextServerConfiguration>,    },    Error {        server: Arc<ContextServer>,        configuration: Arc<ContextServerConfiguration>,        error: Arc<str>,    },} impl ContextServerState {    pub fn server(&self) -> Arc<ContextServer> {        match self {            ContextServerState::Starting { server, .. } => server.clone(),            ContextServerState::Running { server, .. } => server.clone(),            ContextServerState::Stopped { server, .. } => server.clone(),            ContextServerState::Error { server, .. } => server.clone(),        }    }     pub fn configuration(&self) -> Arc<ContextServerConfiguration> {        match self {            ContextServerState::Starting { configuration, .. } => configuration.clone(),            ContextServerState::Running { configuration, .. } => configuration.clone(),            ContextServerState::Stopped { configuration, .. } => configuration.clone(),            ContextServerState::Error { configuration, .. } => configuration.clone(),        }    }} #[derive(Debug, PartialEq, Eq)]pub enum ContextServerConfiguration {    Custom {        command: ContextServerCommand,    },    Extension {        command: ContextServerCommand,        settings: serde_json::Value,    },    Http {        url: url::Url,        headers: HashMap<String, String>,    },} impl ContextServerConfiguration {    pub fn command(&self) -> Option<&ContextServerCommand> {        match self {            ContextServerConfiguration::Custom { command } => Some(command),            ContextServerConfiguration::Extension { command, .. } => Some(command),            ContextServerConfiguration::Http { .. } => None,        }    }     pub async fn from_settings(        settings: ContextServerSettings,        id: ContextServerId,        registry: Entity<ContextServerDescriptorRegistry>,        worktree_store: Entity<WorktreeStore>,        cx: &AsyncApp,    ) -> Option<Self> {        match settings {            ContextServerSettings::Stdio {                enabled: _,                command,            } => Some(ContextServerConfiguration::Custom { command }),            ContextServerSettings::Extension {                enabled: _,                settings,            } => {                let descriptor = cx                    .update(|cx| registry.read(cx).context_server_descriptor(&id.0))                    .ok()                    .flatten()?;                 match descriptor.command(worktree_store, cx).await {                    Ok(command) => {                        Some(ContextServerConfiguration::Extension { command, settings })                    }                    Err(e) => {                        log::error!(                            "Failed to create context server configuration from settings: {e:#}"                        );                        None                    }                }            }            ContextServerSettings::Http {                enabled: _,                url,                headers: auth,            } => {                let url = url::Url::parse(&url).log_err()?;                Some(ContextServerConfiguration::Http { url, headers: auth })            }        }    }} pub type ContextServerFactory =    Box<dyn Fn(ContextServerId, Arc<ContextServerConfiguration>) -> Arc<ContextServer>>; pub struct ContextServerStore {    context_server_settings: HashMap<Arc<str>, ContextServerSettings>,    servers: HashMap<ContextServerId, ContextServerState>,    worktree_store: Entity<WorktreeStore>,    project: WeakEntity<Project>,    registry: Entity<ContextServerDescriptorRegistry>,    update_servers_task: Option<Task<Result<()>>>,    context_server_factory: Option<ContextServerFactory>,    needs_server_update: bool,    _subscriptions: Vec<Subscription>,} pub enum Event {    ServerStatusChanged {        server_id: ContextServerId,        status: ContextServerStatus,    },} impl EventEmitter<Event> for ContextServerStore {} impl ContextServerStore {    pub fn new(        worktree_store: Entity<WorktreeStore>,        weak_project: WeakEntity<Project>,        cx: &mut Context<Self>,    ) -> Self {        Self::new_internal(            true,            None,            ContextServerDescriptorRegistry::default_global(cx),            worktree_store,            weak_project,            cx,        )    }     /// Returns all configured context server ids, excluding the ones that are disabled    pub fn configured_server_ids(&self) -> Vec<ContextServerId> {        self.context_server_settings            .iter()            .filter(|(_, settings)| settings.enabled())            .map(|(id, _)| ContextServerId(id.clone()))            .collect()    }     #[cfg(any(test, feature = "test-support"))]    pub fn test(        registry: Entity<ContextServerDescriptorRegistry>,        worktree_store: Entity<WorktreeStore>,        weak_project: WeakEntity<Project>,        cx: &mut Context<Self>,    ) -> Self {        Self::new_internal(false, None, registry, worktree_store, weak_project, cx)    }     #[cfg(any(test, feature = "test-support"))]    pub fn test_maintain_server_loop(        context_server_factory: Option<ContextServerFactory>,        registry: Entity<ContextServerDescriptorRegistry>,        worktree_store: Entity<WorktreeStore>,        weak_project: WeakEntity<Project>,        cx: &mut Context<Self>,    ) -> Self {        Self::new_internal(            true,            context_server_factory,            registry,            worktree_store,            weak_project,            cx,        )    }     fn new_internal(        maintain_server_loop: bool,        context_server_factory: Option<ContextServerFactory>,        registry: Entity<ContextServerDescriptorRegistry>,        worktree_store: Entity<WorktreeStore>,        weak_project: WeakEntity<Project>,        cx: &mut Context<Self>,    ) -> Self {        let subscriptions = if maintain_server_loop {            vec![                cx.observe(&registry, |this, _registry, cx| {                    this.available_context_servers_changed(cx);                }),                cx.observe_global::<SettingsStore>(|this, cx| {                    let settings = Self::resolve_context_server_settings(&this.worktree_store, cx);                    if &this.context_server_settings == settings {                        return;                    }                    this.context_server_settings = settings.clone();                    this.available_context_servers_changed(cx);                }),            ]        } else {            Vec::new()        };         let mut this = Self {            _subscriptions: subscriptions,            context_server_settings: Self::resolve_context_server_settings(&worktree_store, cx)                .clone(),            worktree_store,            project: weak_project,            registry,            needs_server_update: false,            servers: HashMap::default(),            update_servers_task: None,            context_server_factory,        };        if maintain_server_loop {            this.available_context_servers_changed(cx);        }        this    }     pub fn get_server(&self, id: &ContextServerId) -> Option<Arc<ContextServer>> {        self.servers.get(id).map(|state| state.server())    }     pub fn get_running_server(&self, id: &ContextServerId) -> Option<Arc<ContextServer>> {        if let Some(ContextServerState::Running { server, .. }) = self.servers.get(id) {            Some(server.clone())        } else {            None        }    }     pub fn status_for_server(&self, id: &ContextServerId) -> Option<ContextServerStatus> {        self.servers.get(id).map(ContextServerStatus::from_state)    }     pub fn configuration_for_server(        &self,        id: &ContextServerId,    ) -> Option<Arc<ContextServerConfiguration>> {        self.servers.get(id).map(|state| state.configuration())    }     pub fn server_ids(&self, cx: &App) -> HashSet<ContextServerId> {        self.servers            .keys()            .cloned()            .chain(                self.registry                    .read(cx)                    .context_server_descriptors()                    .into_iter()                    .map(|(id, _)| ContextServerId(id)),            )            .collect()    }     pub fn running_servers(&self) -> Vec<Arc<ContextServer>> {        self.servers            .values()            .filter_map(|state| {                if let ContextServerState::Running { server, .. } = state {                    Some(server.clone())                } else {                    None                }            })            .collect()    }     pub fn start_server(&mut self, server: Arc<ContextServer>, cx: &mut Context<Self>) {        cx.spawn(async move |this, cx| {            let this = this.upgrade().context("Context server store dropped")?;            let settings = this                .update(cx, |this, _| {                    this.context_server_settings.get(&server.id().0).cloned()                })                .ok()                .flatten()                .context("Failed to get context server settings")?;             if !settings.enabled() {                return Ok(());            }             let (registry, worktree_store) = this.update(cx, |this, _| {                (this.registry.clone(), this.worktree_store.clone())            })?;            let configuration = ContextServerConfiguration::from_settings(                settings,                server.id(),                registry,                worktree_store,                cx,            )            .await            .context("Failed to create context server configuration")?;             this.update(cx, |this, cx| {                this.run_server(server, Arc::new(configuration), cx)            })        })        .detach_and_log_err(cx);    }     pub fn stop_server(&mut self, id: &ContextServerId, cx: &mut Context<Self>) -> Result<()> {        if matches!(            self.servers.get(id),            Some(ContextServerState::Stopped { .. })        ) {            return Ok(());        }         let state = self            .servers            .remove(id)            .context("Context server not found")?;         let server = state.server();        let configuration = state.configuration();        let mut result = Ok(());        if let ContextServerState::Running { server, .. } = &state {            result = server.stop();        }        drop(state);         self.update_server_state(            id.clone(),            ContextServerState::Stopped {                configuration,                server,            },            cx,        );         result    }     fn run_server(        &mut self,        server: Arc<ContextServer>,        configuration: Arc<ContextServerConfiguration>,        cx: &mut Context<Self>,    ) {        let id = server.id();        if matches!(            self.servers.get(&id),            Some(ContextServerState::Starting { .. } | ContextServerState::Running { .. })        ) {            self.stop_server(&id, cx).log_err();        }        let task = cx.spawn({            let id = server.id();            let server = server.clone();            let configuration = configuration.clone();             async move |this, cx| {                match server.clone().start(cx).await {                    Ok(_) => {                        debug_assert!(server.client().is_some());                         this.update(cx, |this, cx| {                            this.update_server_state(                                id.clone(),                                ContextServerState::Running {                                    server,                                    configuration,                                },                                cx,                            )                        })                        .log_err()                    }                    Err(err) => {                        log::error!("{} context server failed to start: {}", id, err);                        this.update(cx, |this, cx| {                            this.update_server_state(                                id.clone(),                                ContextServerState::Error {                                    configuration,                                    server,                                    error: err.to_string().into(),                                },                                cx,                            )                        })                        .log_err()                    }                };            }        });         self.update_server_state(            id.clone(),            ContextServerState::Starting {                configuration,                _task: task,                server,            },            cx,        );    }     fn remove_server(&mut self, id: &ContextServerId, cx: &mut Context<Self>) -> Result<()> {        let state = self            .servers            .remove(id)            .context("Context server not found")?;        drop(state);        cx.emit(Event::ServerStatusChanged {            server_id: id.clone(),            status: ContextServerStatus::Stopped,        });        Ok(())    }     fn create_context_server(        &self,        id: ContextServerId,        configuration: Arc<ContextServerConfiguration>,        cx: &mut Context<Self>,    ) -> Result<Arc<ContextServer>> {        if let Some(factory) = self.context_server_factory.as_ref() {            return Ok(factory(id, configuration));        }         match configuration.as_ref() {            ContextServerConfiguration::Http { url, headers } => Ok(Arc::new(ContextServer::http(                id,                url,                headers.clone(),                cx.http_client(),                cx.background_executor().clone(),            )?)),            _ => {                let root_path = self                    .project                    .read_with(cx, |project, cx| project.active_project_directory(cx))

Lines 1–500 of 1437.