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, ContextServerId};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::{ContextServerConfiguration, 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 { error, .. } => { if let Some(error) = error { ContextServerStatus::Error(error.clone()) } else { ContextServerStatus::Stopped } } } }} 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: Option<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(), } } pub fn configuration(&self) -> Arc<ContextServerConfiguration> { match self { ContextServerState::Starting { configuration, .. } => configuration.clone(), ContextServerState::Running { configuration, .. } => configuration.clone(), ContextServerState::Stopped { configuration, .. } => configuration.clone(), } }} 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(®istry, |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 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>, ) -> Result<()> { let location = self .worktree_store .read(cx) .visible_worktrees(cx) .next() .map(|worktree| settings::SettingsLocation { worktree_id: worktree.read(cx).id(), path: Path::new(""), }); let settings = ProjectSettings::get(location, cx); let configuration = settings .context_servers .get(&server.id().0) .context("Failed to load context server configuration from settings")? .clone(); self.run_server(server, Arc::new(configuration), cx); Ok(()) } pub fn stop_server(&mut self, id: &ContextServerId, cx: &mut Context<Self>) -> Result<()> { 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, error: None, }, 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::Stopped { configuration, server, error: Some(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 is_configuration_valid(&self, configuration: &ContextServerConfiguration) -> bool { // Command must be some when we are running in stdio mode. self.context_server_factory.as_ref().is_some() || configuration.command.is_some() } 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 { let command = configuration .command .clone() .context("Missing command to run context server")?; Ok(Arc::new(ContextServer::stdio(id, command))) } } 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 desired_servers = HashMap::default(); let (registry, worktree_store) = this.update(cx, |this, cx| { let location = this .worktree_store .read(cx) .visible_worktrees(cx) .next() .map(|worktree| settings::SettingsLocation { worktree_id: worktree.read(cx).id(), path: Path::new(""), }); let settings = ProjectSettings::get(location, cx); desired_servers = settings.context_servers.clone(); (this.registry.clone(), this.worktree_store.clone()) })?; for (id, descriptor) in registry.read_with(cx, |registry, _| registry.context_server_descriptors())? { let config = desired_servers.entry(id.clone()).or_default(); if config.command.is_none() { if let Some(extension_command) = descriptor .command(worktree_store.clone(), &cx) .await .log_err() { config.command = Some(extension_command); } } } this.update(cx, |this, _| { // Filter out configurations without commands, the user uninstalled an extension. desired_servers.retain(|_, configuration| this.is_configuration_valid(configuration)); })?; let mut servers_to_start = Vec::new(); let mut servers_to_remove = HashSet::default(); let mut servers_to_stop = HashSet::default(); this.update(cx, |this, _cx| { for server_id in this.servers.keys() { // All servers that are not in desired_servers should be removed from the store. // E.g. this can happen if the user removed a server from the configuration, // or the user uninstalled an extension. if !desired_servers.contains_key(&server_id.0) { servers_to_remove.insert(server_id.clone()); } } for (id, config) in desired_servers { let id = ContextServerId(id.clone()); let existing_config = this.servers.get(&id).map(|state| state.configuration()); if existing_config.as_deref() != Some(&config) { let config = Arc::new(config); if let Some(server) = this .create_context_server(id.clone(), config.clone()) .log_err() { servers_to_start.push((server, config)); if this.servers.contains_key(&id) { servers_to_stop.insert(id); } } } } })?; for id in servers_to_stop { this.update(cx, |this, cx| this.stop_server(&id, cx).ok())?; } for id in servers_to_remove { this.update(cx, |this, cx| this.remove_server(&id, cx).ok())?; } for (server, config) in servers_to_start { this.update(cx, |this, cx| this.run_server(server, config, cx)) .log_err(); } Ok(()) }} #[cfg(test)]mod tests { use super::*;Lines 1–500 of 1131. Load more
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, ContextServerId}; 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::{ContextServerConfiguration, 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 { error, .. } => { if let Some(error) = error { ContextServerStatus::Error(error.clone()) } else { ContextServerStatus::Stopped } } } } } 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: Option<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(), } } pub fn configuration(&self) -> Arc<ContextServerConfiguration> { match self { ContextServerState::Starting { configuration, .. } => configuration.clone(), ContextServerState::Running { configuration, .. } => configuration.clone(), ContextServerState::Stopped { configuration, .. } => configuration.clone(), } } } 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(®istry, |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 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>, ) -> Result<()> { let location = self .worktree_store .read(cx) .visible_worktrees(cx) .next() .map(|worktree| settings::SettingsLocation { worktree_id: worktree.read(cx).id(), path: Path::new(""), }); let settings = ProjectSettings::get(location, cx); let configuration = settings .context_servers .get(&server.id().0) .context("Failed to load context server configuration from settings")? .clone(); self.run_server(server, Arc::new(configuration), cx); Ok(()) } pub fn stop_server(&mut self, id: &ContextServerId, cx: &mut Context<Self>) -> Result<()> { 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, error: None, }, 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::Stopped { configuration, server, error: Some(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 is_configuration_valid(&self, configuration: &ContextServerConfiguration) -> bool { // Command must be some when we are running in stdio mode. self.context_server_factory.as_ref().is_some() || configuration.command.is_some() } 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 { let command = configuration .command .clone() .context("Missing command to run context server")?; Ok(Arc::new(ContextServer::stdio(id, command))) } } 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 desired_servers = HashMap::default(); let (registry, worktree_store) = this.update(cx, |this, cx| { let location = this .worktree_store .read(cx) .visible_worktrees(cx) .next() .map(|worktree| settings::SettingsLocation { worktree_id: worktree.read(cx).id(), path: Path::new(""), }); let settings = ProjectSettings::get(location, cx); desired_servers = settings.context_servers.clone(); (this.registry.clone(), this.worktree_store.clone()) })?; for (id, descriptor) in registry.read_with(cx, |registry, _| registry.context_server_descriptors())? { let config = desired_servers.entry(id.clone()).or_default(); if config.command.is_none() { if let Some(extension_command) = descriptor .command(worktree_store.clone(), &cx) .await .log_err() { config.command = Some(extension_command); } } } this.update(cx, |this, _| { // Filter out configurations without commands, the user uninstalled an extension. desired_servers.retain(|_, configuration| this.is_configuration_valid(configuration)); })?; let mut servers_to_start = Vec::new(); let mut servers_to_remove = HashSet::default(); let mut servers_to_stop = HashSet::default(); this.update(cx, |this, _cx| { for server_id in this.servers.keys() { // All servers that are not in desired_servers should be removed from the store. // E.g. this can happen if the user removed a server from the configuration, // or the user uninstalled an extension. if !desired_servers.contains_key(&server_id.0) { servers_to_remove.insert(server_id.clone()); } } for (id, config) in desired_servers { let id = ContextServerId(id.clone()); let existing_config = this.servers.get(&id).map(|state| state.configuration()); if existing_config.as_deref() != Some(&config) { let config = Arc::new(config); if let Some(server) = this .create_context_server(id.clone(), config.clone()) .log_err() { servers_to_start.push((server, config)); if this.servers.contains_key(&id) { servers_to_stop.insert(id); } } } } })?; for id in servers_to_stop { this.update(cx, |this, cx| this.stop_server(&id, cx).ok())?; } for id in servers_to_remove { this.update(cx, |this, cx| this.remove_server(&id, cx).ok())?; } for (server, config) in servers_to_start { this.update(cx, |this, cx| this.run_server(server, config, cx)) .log_err(); } Ok(()) } } #[cfg(test)] mod tests { use super::*;