oak-gpui/crates/extension_host/src/headless_host.rs

use std::{path::PathBuf, sync::Arc};

use anyhow::{anyhow, Context as _, Result};
use client::{proto, TypedEnvelope};
use collections::{HashMap, HashSet};
use extension::{
    Extension, ExtensionHostProxy, ExtensionLanguageProxy, ExtensionLanguageServerProxy,
    ExtensionManifest,
};
use fs::{Fs, RemoveOptions, RenameOptions};
use gpui::{AppContext, AsyncAppContext, Context, Model, ModelContext, Task, WeakModel};
use http_client::HttpClient;
use language::{LanguageConfig, LanguageName, LanguageQueries, LoadedLanguage};
use lsp::LanguageServerName;
use node_runtime::NodeRuntime;

use crate::wasm_host::{WasmExtension, WasmHost};

#[derive(Clone, Debug)]
pub struct ExtensionVersion {
    pub id: String,
    pub version: String,
    pub dev: bool,
}

pub struct HeadlessExtensionStore {
    pub fs: Arc<dyn Fs>,
    pub extension_dir: PathBuf,
    pub proxy: Arc<ExtensionHostProxy>,
    pub wasm_host: Arc<WasmHost>,
    pub loaded_extensions: HashMap<Arc<str>, Arc<str>>,
    pub loaded_languages: HashMap<Arc<str>, Vec<LanguageName>>,
    pub loaded_language_servers: HashMap<Arc<str>, Vec<(LanguageServerName, LanguageName)>>,
}

impl HeadlessExtensionStore {
    pub fn new(
        fs: Arc<dyn Fs>,
        http_client: Arc<dyn HttpClient>,
        extension_dir: PathBuf,
        extension_host_proxy: Arc<ExtensionHostProxy>,
        node_runtime: NodeRuntime,
        cx: &mut AppContext,
    ) -> Model<Self> {
        cx.new_model(|cx| Self {
            fs: fs.clone(),
            wasm_host: WasmHost::new(
                fs.clone(),
                http_client.clone(),
                node_runtime,
                extension_host_proxy.clone(),
                extension_dir.join("work"),
                cx,
            ),
            extension_dir,
            proxy: extension_host_proxy,
            loaded_extensions: Default::default(),
            loaded_languages: Default::default(),
            loaded_language_servers: Default::default(),
        })
    }

    pub fn sync_extensions(
        &mut self,
        extensions: Vec<ExtensionVersion>,
        cx: &ModelContext<Self>,
    ) -> Task<Result<Vec<ExtensionVersion>>> {
        let on_client = HashSet::from_iter(extensions.iter().map(|e| e.id.as_str()));
        let to_remove: Vec<Arc<str>> = self
            .loaded_extensions
            .keys()
            .filter(|id| !on_client.contains(id.as_ref()))
            .cloned()
            .collect();
        let to_load: Vec<ExtensionVersion> = extensions
            .into_iter()
            .filter(|e| {
                if e.dev {
                    return true;
                }
                !self
                    .loaded_extensions
                    .get(e.id.as_str())
                    .is_some_and(|loaded| loaded.as_ref() == e.version.as_str())
            })
            .collect();

        cx.spawn(|this, mut cx| async move {
            let mut missing = Vec::new();

            for extension_id in to_remove {
                log::info!("removing extension: {}", extension_id);
                this.update(&mut cx, |this, cx| {
                    this.uninstall_extension(&extension_id, cx)
                })?
                .await?;
            }

            for extension in to_load {
                if let Err(e) = Self::load_extension(this.clone(), extension.clone(), &mut cx).await
                {
                    log::info!("failed to load extension: {}, {:?}", extension.id, e);
                    missing.push(extension)
                } else if extension.dev {
                    missing.push(extension)
                }
            }

            Ok(missing)
        })
    }

    pub async fn load_extension(
        this: WeakModel<Self>,
        extension: ExtensionVersion,
        cx: &mut AsyncAppContext,
    ) -> Result<()> {
        let (fs, wasm_host, extension_dir) = this.update(cx, |this, _cx| {
            this.loaded_extensions.insert(
                extension.id.clone().into(),
                extension.version.clone().into(),
            );
            (
                this.fs.clone(),
                this.wasm_host.clone(),
                this.extension_dir.join(&extension.id),
            )
        })?;

        let manifest = Arc::new(ExtensionManifest::load(fs.clone(), &extension_dir).await?);

        debug_assert!(!manifest.languages.is_empty() || !manifest.language_servers.is_empty());

        if manifest.version.as_ref() != extension.version.as_str() {
            anyhow::bail!(
                "mismatched versions: ({}) != ({})",
                manifest.version,
                extension.version
            )
        }

        for language_path in &manifest.languages {
            let language_path = extension_dir.join(language_path);
            let config = fs.load(&language_path.join("config.toml")).await?;
            let mut config = ::toml::from_str::<LanguageConfig>(&config)?;

            this.update(cx, |this, _cx| {
                this.loaded_languages
                    .entry(manifest.id.clone())
                    .or_default()
                    .push(config.name.clone());

                config.grammar = None;

                this.proxy.register_language(
                    config.name.clone(),
                    None,
                    config.matcher.clone(),
                    Arc::new(move || {
                        Ok(LoadedLanguage {
                            config: config.clone(),
                            queries: LanguageQueries::default(),
                            context_provider: None,
                            toolchain_provider: None,
                        })
                    }),
                );
            })?;
        }

        if manifest.language_servers.is_empty() {
            return Ok(());
        }

        let wasm_extension: Arc<dyn Extension> =
            Arc::new(WasmExtension::load(extension_dir, &manifest, wasm_host.clone(), &cx).await?);

        for (language_server_id, language_server_config) in &manifest.language_servers {
            for language in language_server_config.languages() {
                this.update(cx, |this, _cx| {
                    this.loaded_language_servers
                        .entry(manifest.id.clone())
                        .or_default()
                        .push((language_server_id.clone(), language.clone()));
                    this.proxy.register_language_server(
                        wasm_extension.clone(),
                        language_server_id.clone(),
                        language.clone(),
                    );
                })?;
            }
        }

        Ok(())
    }

    fn uninstall_extension(
        &mut self,
        extension_id: &Arc<str>,
        cx: &mut ModelContext<Self>,
    ) -> Task<Result<()>> {
        self.loaded_extensions.remove(extension_id);

        let languages_to_remove = self
            .loaded_languages
            .remove(extension_id)
            .unwrap_or_default();
        self.proxy.remove_languages(&languages_to_remove, &[]);

        for (language_server_name, language) in self
            .loaded_language_servers
            .remove(extension_id)
            .unwrap_or_default()
        {
            self.proxy
                .remove_language_server(&language, &language_server_name);
        }

        let path = self.extension_dir.join(&extension_id.to_string());
        let fs = self.fs.clone();
        cx.spawn(|_, _| async move {
            fs.remove_dir(
                &path,
                RemoveOptions {
                    recursive: true,
                    ignore_if_not_exists: true,
                },
            )
            .await
        })
    }

    pub fn install_extension(
        &mut self,
        extension: ExtensionVersion,
        tmp_path: PathBuf,
        cx: &mut ModelContext<Self>,
    ) -> Task<Result<()>> {
        let path = self.extension_dir.join(&extension.id);
        let fs = self.fs.clone();

        cx.spawn(|this, mut cx| async move {
            if fs.is_dir(&path).await {
                this.update(&mut cx, |this, cx| {
                    this.uninstall_extension(&extension.id.clone().into(), cx)
                })?
                .await?;
            }

            fs.rename(&tmp_path, &path, RenameOptions::default())
                .await?;

            Self::load_extension(this, extension, &mut cx).await
        })
    }

    pub async fn handle_sync_extensions(
        extension_store: Model<HeadlessExtensionStore>,
        envelope: TypedEnvelope<proto::SyncExtensions>,
        mut cx: AsyncAppContext,
    ) -> Result<proto::SyncExtensionsResponse> {
        let requested_extensions =
            envelope
                .payload
                .extensions
                .into_iter()
                .map(|p| ExtensionVersion {
                    id: p.id,
                    version: p.version,
                    dev: p.dev,
                });
        let missing_extensions = extension_store
            .update(&mut cx, |extension_store, cx| {
                extension_store.sync_extensions(requested_extensions.collect(), cx)
            })?
            .await?;

        Ok(proto::SyncExtensionsResponse {
            missing_extensions: missing_extensions
                .into_iter()
                .map(|e| proto::Extension {
                    id: e.id,
                    version: e.version,
                    dev: e.dev,
                })
                .collect(),
            tmp_dir: paths::remote_extensions_uploads_dir()
                .to_string_lossy()
                .to_string(),
        })
    }

    pub async fn handle_install_extension(
        extensions: Model<HeadlessExtensionStore>,
        envelope: TypedEnvelope<proto::InstallExtension>,
        mut cx: AsyncAppContext,
    ) -> Result<proto::Ack> {
        let extension = envelope
            .payload
            .extension
            .with_context(|| anyhow!("Invalid InstallExtension request"))?;

        extensions
            .update(&mut cx, |extensions, cx| {
                extensions.install_extension(
                    ExtensionVersion {
                        id: extension.id,
                        version: extension.version,
                        dev: extension.dev,
                    },
                    PathBuf::from(envelope.payload.tmp_dir),
                    cx,
                )
            })?
            .await?;

        Ok(proto::Ack {})
    }
}