pub mod extension;pub mod registry; use std::path::Path;use std::sync::Arc;use std::time::Duration; 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 remote::RemoteClient;use rpc::{AnyProtoClient, TypedEnvelope, proto};use settings::{Settings as _, SettingsStore};use util::{ResultExt as _, rel_path::RelPath}; use crate::{    Project,    project_settings::{ContextServerSettings, ProjectSettings},    worktree_store::WorktreeStore,}; /// Maximum timeout for context server requests/// Prevents extremely large timeout values from tying up resources indefinitely.const MAX_TIMEOUT_SECS: u64 = 600; // 10 minutes 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,        remote: bool,    },    Extension {        command: ContextServerCommand,        settings: serde_json::Value,        remote: bool,    },    Http {        url: url::Url,        headers: HashMap<String, String>,        timeout: Option<u64>,    },} impl ContextServerConfiguration {    pub fn command(&self) -> Option<&ContextServerCommand> {        match self {            ContextServerConfiguration::Custom { command, .. } => Some(command),            ContextServerConfiguration::Extension { command, .. } => Some(command),            ContextServerConfiguration::Http { .. } => None,        }    }     pub fn remote(&self) -> bool {        match self {            ContextServerConfiguration::Custom { remote, .. } => *remote,            ContextServerConfiguration::Extension { remote, .. } => *remote,            ContextServerConfiguration::Http { .. } => false,        }    }     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,                remote,            } => Some(ContextServerConfiguration::Custom { command, remote }),            ContextServerSettings::Extension {                enabled: _,                settings,                remote,            } => {                let descriptor =                    cx.update(|cx| registry.read(cx).context_server_descriptor(&id.0))?;                 match descriptor.command(worktree_store, cx).await {                    Ok(command) => Some(ContextServerConfiguration::Extension {                        command,                        settings,                        remote,                    }),                    Err(e) => {                        log::error!(                            "Failed to create context server configuration from settings: {e:#}"                        );                        None                    }                }            }            ContextServerSettings::Http {                enabled: _,                url,                headers: auth,                timeout,            } => {                let url = url::Url::parse(&url).log_err()?;                Some(ContextServerConfiguration::Http {                    url,                    headers: auth,                    timeout,                })            }        }    }} pub type ContextServerFactory =    Box<dyn Fn(ContextServerId, Arc<ContextServerConfiguration>) -> Arc<ContextServer>>; enum ContextServerStoreState {    Local {        downstream_client: Option<(u64, AnyProtoClient)>,        is_headless: bool,    },    Remote {        project_id: u64,        upstream_client: Entity<RemoteClient>,    },} pub struct ContextServerStore {    state: ContextServerStoreState,    context_server_settings: HashMap<Arc<str>, ContextServerSettings>,    servers: HashMap<ContextServerId, ContextServerState>,    worktree_store: Entity<WorktreeStore>,    project: Option<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 local(        worktree_store: Entity<WorktreeStore>,        weak_project: Option<WeakEntity<Project>>,        headless: bool,        cx: &mut Context<Self>,    ) -> Self {        Self::new_internal(            !headless,            None,            ContextServerDescriptorRegistry::default_global(cx),            worktree_store,            weak_project,            ContextServerStoreState::Local {                downstream_client: None,                is_headless: headless,            },            cx,        )    }     pub fn remote(        project_id: u64,        upstream_client: Entity<RemoteClient>,        worktree_store: Entity<WorktreeStore>,        weak_project: Option<WeakEntity<Project>>,        cx: &mut Context<Self>,    ) -> Self {        Self::new_internal(            true,            None,            ContextServerDescriptorRegistry::default_global(cx),            worktree_store,            weak_project,            ContextServerStoreState::Remote {                project_id,                upstream_client,            },            cx,        )    }     pub fn init_headless(session: &AnyProtoClient) {        session.add_entity_request_handler(Self::handle_get_context_server_command);    }     pub fn shared(&mut self, project_id: u64, client: AnyProtoClient) {        if let ContextServerStoreState::Local {            downstream_client, ..        } = &mut self.state        {            *downstream_client = Some((project_id, client));        }    }     pub fn is_remote_project(&self) -> bool {        matches!(self.state, ContextServerStoreState::Remote { .. })    }     /// 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: Option<WeakEntity<Project>>,        cx: &mut Context<Self>,    ) -> Self {        Self::new_internal(            false,            None,            registry,            worktree_store,            weak_project,            ContextServerStoreState::Local {                downstream_client: None,                is_headless: false,            },            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: Option<WeakEntity<Project>>,        cx: &mut Context<Self>,    ) -> Self {        Self::new_internal(            true,            context_server_factory,            registry,            worktree_store,            weak_project,            ContextServerStoreState::Local {                downstream_client: None,                is_headless: false,            },            cx,        )    }     #[cfg(any(test, feature = "test-support"))]    pub fn set_context_server_factory(&mut self, factory: ContextServerFactory) {        self.context_server_factory = Some(factory);    }     #[cfg(any(test, feature = "test-support"))]    pub fn registry(&self) -> &Entity<ContextServerDescriptorRegistry> {        &self.registry    }     #[cfg(any(test, feature = "test-support"))]    pub fn test_start_server(&mut self, server: Arc<ContextServer>, cx: &mut Context<Self>) {        let configuration = Arc::new(ContextServerConfiguration::Custom {            command: ContextServerCommand {                path: "test".into(),                args: vec![],                env: None,                timeout: None,            },            remote: false,        });        self.run_server(server, configuration, cx);    }     fn new_internal(        maintain_server_loop: bool,        context_server_factory: Option<ContextServerFactory>,        registry: Entity<ContextServerDescriptorRegistry>,        worktree_store: Entity<WorktreeStore>,        weak_project: Option<WeakEntity<Project>>,        state: ContextServerStoreState,        cx: &mut Context<Self>,    ) -> Self {        let mut subscriptions = vec![cx.observe_global::<SettingsStore>(move |this, cx| {            let settings =                &Self::resolve_project_settings(&this.worktree_store, cx).context_servers;            if &this.context_server_settings == settings {                return;            }            this.context_server_settings = settings.clone();            if maintain_server_loop {                this.available_context_servers_changed(cx);            }        })];         if maintain_server_loop {            subscriptions.push(cx.observe(&registry, |this, _registry, cx| {                this.available_context_servers_changed(cx);            }));        }         let mut this = Self {            state,            _subscriptions: subscriptions,            context_server_settings: Self::resolve_project_settings(&worktree_store, cx)                .context_servers                .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()                })                .context("Failed to get context server settings")?;             if !settings.enabled() {                return anyhow::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)            });            Ok(())        })        .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)

Lines 1–500 of 1811. Load more