pub mod extension;pub mod registry; use std::{path::Path, 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 _; use crate::{    project_settings::{ContextServerSettings, ProjectSettings},    worktree_store::WorktreeStore,}; pub fn init(cx: &mut App) {    extension::init(cx);} actions!(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,    },} impl ContextServerConfiguration {    pub fn command(&self) -> &ContextServerCommand {        match self {            ContextServerConfiguration::Custom { command } => command,            ContextServerConfiguration::Extension { command, .. } => command,        }    }     pub async fn from_settings(        settings: ContextServerSettings,        id: ContextServerId,        registry: Entity<ContextServerDescriptorRegistry>,        worktree_store: Entity<WorktreeStore>,        cx: &AsyncApp,    ) -> Option<Self> {        match settings {            ContextServerSettings::Custom {                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()?;                 let command = descriptor.command(worktree_store, cx).await.log_err()?;                 Some(ContextServerConfiguration::Extension { command, settings })            }        }    }} pub type ContextServerFactory =    Box<dyn Fn(ContextServerId, Arc<ContextServerConfiguration>) -> Arc<ContextServer>>; pub struct ContextServerStore {    servers: HashMap<ContextServerId, ContextServerState>,    worktree_store: Entity<WorktreeStore>,    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>, cx: &mut Context<Self>) -> Self {        Self::new_internal(            true,            None,            ContextServerDescriptorRegistry::default_global(cx),            worktree_store,            cx,        )    }     #[cfg(any(test, feature = "test-support"))]    pub fn test(        registry: Entity<ContextServerDescriptorRegistry>,        worktree_store: Entity<WorktreeStore>,        cx: &mut Context<Self>,    ) -> Self {        Self::new_internal(false, None, registry, worktree_store, cx)    }     #[cfg(any(test, feature = "test-support"))]    pub fn test_maintain_server_loop(        context_server_factory: ContextServerFactory,        registry: Entity<ContextServerDescriptorRegistry>,        worktree_store: Entity<WorktreeStore>,        cx: &mut Context<Self>,    ) -> Self {        Self::new_internal(            true,            Some(context_server_factory),            registry,            worktree_store,            cx,        )    }     fn new_internal(        maintain_server_loop: bool,        context_server_factory: Option<ContextServerFactory>,        registry: Entity<ContextServerDescriptorRegistry>,        worktree_store: Entity<WorktreeStore>,        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| {                    this.available_context_servers_changed(cx);                }),            ]        } else {            Vec::new()        };         let mut this = Self {            _subscriptions: subscriptions,            worktree_store,            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 all_server_ids(&self) -> Vec<ContextServerId> {        self.servers.keys().cloned().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, cx| {                    this.context_server_settings(cx)                        .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    }     pub fn restart_server(&mut self, id: &ContextServerId, cx: &mut Context<Self>) -> Result<()> {        if let Some(state) = self.servers.get(&id) {            let configuration = state.configuration();             self.stop_server(&state.server().id(), cx)?;            let new_server = self.create_context_server(id.clone(), configuration.clone())?;            self.run_server(new_server, configuration, cx);        }        Ok(())    }     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(_) => {                        log::info!("Started {} context server", id);                        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>,    ) -> Result<Arc<ContextServer>> {        if let Some(factory) = self.context_server_factory.as_ref() {            Ok(factory(id, configuration))        } else {            Ok(Arc::new(ContextServer::stdio(                id,                configuration.command().clone(),            )))        }    }     fn context_server_settings<'a>(        &'a self,        cx: &'a App,    ) -> &'a HashMap<Arc<str>, ContextServerSettings> {        let location = self            .worktree_store            .read(cx)            .visible_worktrees(cx)            .next()            .map(|worktree| settings::SettingsLocation {                worktree_id: worktree.read(cx).id(),                path: Path::new(""),            });        &ProjectSettings::get(location, cx).context_servers    }     fn update_server_state(        &mut self,        id: ContextServerId,        state: ContextServerState,        cx: &mut Context<Self>,    ) {        let status = ContextServerStatus::from_state(&state);        self.servers.insert(id.clone(), state);        cx.emit(Event::ServerStatusChanged {            server_id: id,            status,        });    }     fn available_context_servers_changed(&mut self, cx: &mut Context<Self>) {        if self.update_servers_task.is_some() {            self.needs_server_update = true;        } else {            self.needs_server_update = false;            self.update_servers_task = Some(cx.spawn(async move |this, cx| {                if let Err(err) = Self::maintain_servers(this.clone(), cx).await {                    log::error!("Error maintaining context servers: {}", err);                }                 this.update(cx, |this, cx| {                    this.update_servers_task.take();                    if this.needs_server_update {                        this.available_context_servers_changed(cx);                    }                })?;                 Ok(())            }));        }    }     async fn maintain_servers(this: WeakEntity<Self>, cx: &mut AsyncApp) -> Result<()> {        let (mut configured_servers, registry, worktree_store) = this.update(cx, |this, cx| {            (                this.context_server_settings(cx).clone(),                this.registry.clone(),                this.worktree_store.clone(),            )

Lines 1–500 of 1242. Load more