1use std::{path::PathBuf, sync::Arc};
2
3use anyhow::{anyhow, Context as _, Result};
4use client::{proto, TypedEnvelope};
5use collections::{HashMap, HashSet};
6use extension::{Extension, ExtensionManifest};
7use fs::{Fs, RemoveOptions, RenameOptions};
8use gpui::{AppContext, AsyncAppContext, Context, Model, ModelContext, Task, WeakModel};
9use http_client::HttpClient;
10use language::{LanguageConfig, LanguageName, LanguageQueries, LanguageRegistry, LoadedLanguage};
11use lsp::LanguageServerName;
12use node_runtime::NodeRuntime;
13
14use crate::{
15 extension_lsp_adapter::ExtensionLspAdapter,
16 wasm_host::{WasmExtension, WasmHost},
17 ExtensionRegistrationHooks,
18};
19
20pub struct HeadlessExtensionStore {
21 pub registration_hooks: Arc<dyn ExtensionRegistrationHooks>,
22 pub fs: Arc<dyn Fs>,
23 pub extension_dir: PathBuf,
24 pub wasm_host: Arc<WasmHost>,
25 pub loaded_extensions: HashMap<Arc<str>, Arc<str>>,
26 pub loaded_languages: HashMap<Arc<str>, Vec<LanguageName>>,
27 pub loaded_language_servers: HashMap<Arc<str>, Vec<(LanguageServerName, LanguageName)>>,
28}
29
30#[derive(Clone, Debug)]
31pub struct ExtensionVersion {
32 pub id: String,
33 pub version: String,
34 pub dev: bool,
35}
36
37impl HeadlessExtensionStore {
38 pub fn new(
39 fs: Arc<dyn Fs>,
40 http_client: Arc<dyn HttpClient>,
41 languages: Arc<LanguageRegistry>,
42 extension_dir: PathBuf,
43 node_runtime: NodeRuntime,
44 cx: &mut AppContext,
45 ) -> Model<Self> {
46 let registration_hooks = Arc::new(HeadlessRegistrationHooks::new(languages.clone()));
47 cx.new_model(|cx| Self {
48 registration_hooks: registration_hooks.clone(),
49 fs: fs.clone(),
50 wasm_host: WasmHost::new(
51 fs.clone(),
52 http_client.clone(),
53 node_runtime,
54 registration_hooks,
55 extension_dir.join("work"),
56 cx,
57 ),
58 extension_dir,
59 loaded_extensions: Default::default(),
60 loaded_languages: Default::default(),
61 loaded_language_servers: Default::default(),
62 })
63 }
64
65 pub fn sync_extensions(
66 &mut self,
67 extensions: Vec<ExtensionVersion>,
68 cx: &ModelContext<Self>,
69 ) -> Task<Result<Vec<ExtensionVersion>>> {
70 let on_client = HashSet::from_iter(extensions.iter().map(|e| e.id.as_str()));
71 let to_remove: Vec<Arc<str>> = self
72 .loaded_extensions
73 .keys()
74 .filter(|id| !on_client.contains(id.as_ref()))
75 .cloned()
76 .collect();
77 let to_load: Vec<ExtensionVersion> = extensions
78 .into_iter()
79 .filter(|e| {
80 if e.dev {
81 return true;
82 }
83 !self
84 .loaded_extensions
85 .get(e.id.as_str())
86 .is_some_and(|loaded| loaded.as_ref() == e.version.as_str())
87 })
88 .collect();
89
90 cx.spawn(|this, mut cx| async move {
91 let mut missing = Vec::new();
92
93 for extension_id in to_remove {
94 log::info!("removing extension: {}", extension_id);
95 this.update(&mut cx, |this, cx| {
96 this.uninstall_extension(&extension_id, cx)
97 })?
98 .await?;
99 }
100
101 for extension in to_load {
102 if let Err(e) = Self::load_extension(this.clone(), extension.clone(), &mut cx).await
103 {
104 log::info!("failed to load extension: {}, {:?}", extension.id, e);
105 missing.push(extension)
106 } else if extension.dev {
107 missing.push(extension)
108 }
109 }
110
111 Ok(missing)
112 })
113 }
114
115 pub async fn load_extension(
116 this: WeakModel<Self>,
117 extension: ExtensionVersion,
118 cx: &mut AsyncAppContext,
119 ) -> Result<()> {
120 let (fs, wasm_host, extension_dir) = this.update(cx, |this, _cx| {
121 this.loaded_extensions.insert(
122 extension.id.clone().into(),
123 extension.version.clone().into(),
124 );
125 (
126 this.fs.clone(),
127 this.wasm_host.clone(),
128 this.extension_dir.join(&extension.id),
129 )
130 })?;
131
132 let manifest = Arc::new(ExtensionManifest::load(fs.clone(), &extension_dir).await?);
133
134 debug_assert!(!manifest.languages.is_empty() || !manifest.language_servers.is_empty());
135
136 if manifest.version.as_ref() != extension.version.as_str() {
137 anyhow::bail!(
138 "mismatched versions: ({}) != ({})",
139 manifest.version,
140 extension.version
141 )
142 }
143
144 for language_path in &manifest.languages {
145 let language_path = extension_dir.join(language_path);
146 let config = fs.load(&language_path.join("config.toml")).await?;
147 let mut config = ::toml::from_str::<LanguageConfig>(&config)?;
148
149 this.update(cx, |this, _cx| {
150 this.loaded_languages
151 .entry(manifest.id.clone())
152 .or_default()
153 .push(config.name.clone());
154
155 config.grammar = None;
156
157 this.registration_hooks.register_language(
158 config.name.clone(),
159 None,
160 config.matcher.clone(),
161 Arc::new(move || {
162 Ok(LoadedLanguage {
163 config: config.clone(),
164 queries: LanguageQueries::default(),
165 context_provider: None,
166 toolchain_provider: None,
167 })
168 }),
169 );
170 })?;
171 }
172
173 if manifest.language_servers.is_empty() {
174 return Ok(());
175 }
176
177 let wasm_extension: Arc<dyn Extension> =
178 Arc::new(WasmExtension::load(extension_dir, &manifest, wasm_host.clone(), &cx).await?);
179
180 for (language_server_id, language_server_config) in &manifest.language_servers {
181 for language in language_server_config.languages() {
182 this.update(cx, |this, _cx| {
183 this.loaded_language_servers
184 .entry(manifest.id.clone())
185 .or_default()
186 .push((language_server_id.clone(), language.clone()));
187 this.registration_hooks.register_lsp_adapter(
188 wasm_extension.clone(),
189 language_server_id.clone(),
190 language.clone(),
191 );
192 })?;
193 }
194 }
195
196 Ok(())
197 }
198
199 fn uninstall_extension(
200 &mut self,
201 extension_id: &Arc<str>,
202 cx: &mut ModelContext<Self>,
203 ) -> Task<Result<()>> {
204 self.loaded_extensions.remove(extension_id);
205 let languages_to_remove = self
206 .loaded_languages
207 .remove(extension_id)
208 .unwrap_or_default();
209 self.registration_hooks
210 .remove_languages(&languages_to_remove, &[]);
211 for (language_server_name, language) in self
212 .loaded_language_servers
213 .remove(extension_id)
214 .unwrap_or_default()
215 {
216 self.registration_hooks
217 .remove_lsp_adapter(&language, &language_server_name);
218 }
219
220 let path = self.extension_dir.join(&extension_id.to_string());
221 let fs = self.fs.clone();
222 cx.spawn(|_, _| async move {
223 fs.remove_dir(
224 &path,
225 RemoveOptions {
226 recursive: true,
227 ignore_if_not_exists: true,
228 },
229 )
230 .await
231 })
232 }
233
234 pub fn install_extension(
235 &mut self,
236 extension: ExtensionVersion,
237 tmp_path: PathBuf,
238 cx: &mut ModelContext<Self>,
239 ) -> Task<Result<()>> {
240 let path = self.extension_dir.join(&extension.id);
241 let fs = self.fs.clone();
242
243 cx.spawn(|this, mut cx| async move {
244 if fs.is_dir(&path).await {
245 this.update(&mut cx, |this, cx| {
246 this.uninstall_extension(&extension.id.clone().into(), cx)
247 })?
248 .await?;
249 }
250
251 fs.rename(&tmp_path, &path, RenameOptions::default())
252 .await?;
253
254 Self::load_extension(this, extension, &mut cx).await
255 })
256 }
257
258 pub async fn handle_sync_extensions(
259 extension_store: Model<HeadlessExtensionStore>,
260 envelope: TypedEnvelope<proto::SyncExtensions>,
261 mut cx: AsyncAppContext,
262 ) -> Result<proto::SyncExtensionsResponse> {
263 let requested_extensions =
264 envelope
265 .payload
266 .extensions
267 .into_iter()
268 .map(|p| ExtensionVersion {
269 id: p.id,
270 version: p.version,
271 dev: p.dev,
272 });
273 let missing_extensions = extension_store
274 .update(&mut cx, |extension_store, cx| {
275 extension_store.sync_extensions(requested_extensions.collect(), cx)
276 })?
277 .await?;
278
279 Ok(proto::SyncExtensionsResponse {
280 missing_extensions: missing_extensions
281 .into_iter()
282 .map(|e| proto::Extension {
283 id: e.id,
284 version: e.version,
285 dev: e.dev,
286 })
287 .collect(),
288 tmp_dir: paths::remote_extensions_uploads_dir()
289 .to_string_lossy()
290 .to_string(),
291 })
292 }
293
294 pub async fn handle_install_extension(
295 extensions: Model<HeadlessExtensionStore>,
296 envelope: TypedEnvelope<proto::InstallExtension>,
297 mut cx: AsyncAppContext,
298 ) -> Result<proto::Ack> {
299 let extension = envelope
300 .payload
301 .extension
302 .with_context(|| anyhow!("Invalid InstallExtension request"))?;
303
304 extensions
305 .update(&mut cx, |extensions, cx| {
306 extensions.install_extension(
307 ExtensionVersion {
308 id: extension.id,
309 version: extension.version,
310 dev: extension.dev,
311 },
312 PathBuf::from(envelope.payload.tmp_dir),
313 cx,
314 )
315 })?
316 .await?;
317
318 Ok(proto::Ack {})
319 }
320}
321
322struct HeadlessRegistrationHooks {
323 language_registry: Arc<LanguageRegistry>,
324}
325
326impl HeadlessRegistrationHooks {
327 fn new(language_registry: Arc<LanguageRegistry>) -> Self {
328 Self { language_registry }
329 }
330}
331
332impl ExtensionRegistrationHooks for HeadlessRegistrationHooks {
333 fn register_language(
334 &self,
335 language: LanguageName,
336 _grammar: Option<Arc<str>>,
337 matcher: language::LanguageMatcher,
338 load: Arc<dyn Fn() -> Result<LoadedLanguage> + 'static + Send + Sync>,
339 ) {
340 log::info!("registering language: {:?}", language);
341 self.language_registry
342 .register_language(language, None, matcher, load)
343 }
344
345 fn register_lsp_adapter(
346 &self,
347 extension: Arc<dyn Extension>,
348 language_server_id: LanguageServerName,
349 language: LanguageName,
350 ) {
351 log::info!("registering lsp adapter {:?}", language);
352 self.language_registry.register_lsp_adapter(
353 language.clone(),
354 Arc::new(ExtensionLspAdapter::new(
355 extension,
356 language_server_id,
357 language,
358 )),
359 );
360 }
361
362 fn register_wasm_grammars(&self, grammars: Vec<(Arc<str>, PathBuf)>) {
363 self.language_registry.register_wasm_grammars(grammars)
364 }
365
366 fn remove_lsp_adapter(&self, language: &LanguageName, server_name: &LanguageServerName) {
367 self.language_registry
368 .remove_lsp_adapter(language, server_name)
369 }
370
371 fn remove_languages(
372 &self,
373 languages_to_remove: &[LanguageName],
374 _grammars_to_remove: &[Arc<str>],
375 ) {
376 self.language_registry
377 .remove_languages(languages_to_remove, &[])
378 }
379
380 fn update_lsp_status(
381 &self,
382 server_name: LanguageServerName,
383 status: language::LanguageServerBinaryStatus,
384 ) {
385 self.language_registry
386 .update_lsp_status(server_name, status)
387 }
388}