1pub mod wit;
2
3use crate::ExtensionManifest;
4use anyhow::{Context as _, Result, anyhow, bail};
5use async_trait::async_trait;
6use extension::{
7 CodeLabel, Command, Completion, ContextServerConfiguration, ExtensionHostProxy,
8 KeyValueStoreDelegate, ProjectDelegate, SlashCommand, SlashCommandArgumentCompletion,
9 SlashCommandOutput, Symbol, WorktreeDelegate,
10};
11use fs::{Fs, normalize_path};
12use futures::future::LocalBoxFuture;
13use futures::{
14 Future, FutureExt, StreamExt as _,
15 channel::{
16 mpsc::{self, UnboundedSender},
17 oneshot,
18 },
19 future::BoxFuture,
20};
21use gpui::{App, AsyncApp, BackgroundExecutor, Task};
22use http_client::HttpClient;
23use language::LanguageName;
24use lsp::LanguageServerName;
25use node_runtime::NodeRuntime;
26use release_channel::ReleaseChannel;
27use semantic_version::SemanticVersion;
28use std::{
29 path::{Path, PathBuf},
30 sync::{Arc, OnceLock},
31};
32use wasmtime::{
33 Engine, Store,
34 component::{Component, ResourceTable},
35};
36use wasmtime_wasi::{self as wasi, WasiView};
37use wit::Extension;
38
39pub struct WasmHost {
40 engine: Engine,
41 release_channel: ReleaseChannel,
42 http_client: Arc<dyn HttpClient>,
43 node_runtime: NodeRuntime,
44 pub(crate) proxy: Arc<ExtensionHostProxy>,
45 fs: Arc<dyn Fs>,
46 pub work_dir: PathBuf,
47 _main_thread_message_task: Task<()>,
48 main_thread_message_tx: mpsc::UnboundedSender<MainThreadCall>,
49}
50
51#[derive(Clone)]
52pub struct WasmExtension {
53 tx: UnboundedSender<ExtensionCall>,
54 pub manifest: Arc<ExtensionManifest>,
55 pub work_dir: Arc<Path>,
56 #[allow(unused)]
57 pub zed_api_version: SemanticVersion,
58}
59
60#[async_trait]
61impl extension::Extension for WasmExtension {
62 fn manifest(&self) -> Arc<ExtensionManifest> {
63 self.manifest.clone()
64 }
65
66 fn work_dir(&self) -> Arc<Path> {
67 self.work_dir.clone()
68 }
69
70 async fn language_server_command(
71 &self,
72 language_server_id: LanguageServerName,
73 language_name: LanguageName,
74 worktree: Arc<dyn WorktreeDelegate>,
75 ) -> Result<Command> {
76 self.call(|extension, store| {
77 async move {
78 let resource = store.data_mut().table().push(worktree)?;
79 let command = extension
80 .call_language_server_command(
81 store,
82 &language_server_id,
83 &language_name,
84 resource,
85 )
86 .await?
87 .map_err(|err| anyhow!("{err}"))?;
88
89 Ok(command.into())
90 }
91 .boxed()
92 })
93 .await
94 }
95
96 async fn language_server_initialization_options(
97 &self,
98 language_server_id: LanguageServerName,
99 language_name: LanguageName,
100 worktree: Arc<dyn WorktreeDelegate>,
101 ) -> Result<Option<String>> {
102 self.call(|extension, store| {
103 async move {
104 let resource = store.data_mut().table().push(worktree)?;
105 let options = extension
106 .call_language_server_initialization_options(
107 store,
108 &language_server_id,
109 &language_name,
110 resource,
111 )
112 .await?
113 .map_err(|err| anyhow!("{err}"))?;
114 anyhow::Ok(options)
115 }
116 .boxed()
117 })
118 .await
119 }
120
121 async fn language_server_workspace_configuration(
122 &self,
123 language_server_id: LanguageServerName,
124 worktree: Arc<dyn WorktreeDelegate>,
125 ) -> Result<Option<String>> {
126 self.call(|extension, store| {
127 async move {
128 let resource = store.data_mut().table().push(worktree)?;
129 let options = extension
130 .call_language_server_workspace_configuration(
131 store,
132 &language_server_id,
133 resource,
134 )
135 .await?
136 .map_err(|err| anyhow!("{err}"))?;
137 anyhow::Ok(options)
138 }
139 .boxed()
140 })
141 .await
142 }
143
144 async fn language_server_additional_initialization_options(
145 &self,
146 language_server_id: LanguageServerName,
147 target_language_server_id: LanguageServerName,
148 worktree: Arc<dyn WorktreeDelegate>,
149 ) -> Result<Option<String>> {
150 self.call(|extension, store| {
151 async move {
152 let resource = store.data_mut().table().push(worktree)?;
153 let options = extension
154 .call_language_server_additional_initialization_options(
155 store,
156 &language_server_id,
157 &target_language_server_id,
158 resource,
159 )
160 .await?
161 .map_err(|err| anyhow!("{err}"))?;
162 anyhow::Ok(options)
163 }
164 .boxed()
165 })
166 .await
167 }
168
169 async fn language_server_additional_workspace_configuration(
170 &self,
171 language_server_id: LanguageServerName,
172 target_language_server_id: LanguageServerName,
173 worktree: Arc<dyn WorktreeDelegate>,
174 ) -> Result<Option<String>> {
175 self.call(|extension, store| {
176 async move {
177 let resource = store.data_mut().table().push(worktree)?;
178 let options = extension
179 .call_language_server_additional_workspace_configuration(
180 store,
181 &language_server_id,
182 &target_language_server_id,
183 resource,
184 )
185 .await?
186 .map_err(|err| anyhow!("{err}"))?;
187 anyhow::Ok(options)
188 }
189 .boxed()
190 })
191 .await
192 }
193
194 async fn labels_for_completions(
195 &self,
196 language_server_id: LanguageServerName,
197 completions: Vec<Completion>,
198 ) -> Result<Vec<Option<CodeLabel>>> {
199 self.call(|extension, store| {
200 async move {
201 let labels = extension
202 .call_labels_for_completions(
203 store,
204 &language_server_id,
205 completions.into_iter().map(Into::into).collect(),
206 )
207 .await?
208 .map_err(|err| anyhow!("{err}"))?;
209
210 Ok(labels
211 .into_iter()
212 .map(|label| label.map(Into::into))
213 .collect())
214 }
215 .boxed()
216 })
217 .await
218 }
219
220 async fn labels_for_symbols(
221 &self,
222 language_server_id: LanguageServerName,
223 symbols: Vec<Symbol>,
224 ) -> Result<Vec<Option<CodeLabel>>> {
225 self.call(|extension, store| {
226 async move {
227 let labels = extension
228 .call_labels_for_symbols(
229 store,
230 &language_server_id,
231 symbols.into_iter().map(Into::into).collect(),
232 )
233 .await?
234 .map_err(|err| anyhow!("{err}"))?;
235
236 Ok(labels
237 .into_iter()
238 .map(|label| label.map(Into::into))
239 .collect())
240 }
241 .boxed()
242 })
243 .await
244 }
245
246 async fn complete_slash_command_argument(
247 &self,
248 command: SlashCommand,
249 arguments: Vec<String>,
250 ) -> Result<Vec<SlashCommandArgumentCompletion>> {
251 self.call(|extension, store| {
252 async move {
253 let completions = extension
254 .call_complete_slash_command_argument(store, &command.into(), &arguments)
255 .await?
256 .map_err(|err| anyhow!("{err}"))?;
257
258 Ok(completions.into_iter().map(Into::into).collect())
259 }
260 .boxed()
261 })
262 .await
263 }
264
265 async fn run_slash_command(
266 &self,
267 command: SlashCommand,
268 arguments: Vec<String>,
269 delegate: Option<Arc<dyn WorktreeDelegate>>,
270 ) -> Result<SlashCommandOutput> {
271 self.call(|extension, store| {
272 async move {
273 let resource = if let Some(delegate) = delegate {
274 Some(store.data_mut().table().push(delegate)?)
275 } else {
276 None
277 };
278
279 let output = extension
280 .call_run_slash_command(store, &command.into(), &arguments, resource)
281 .await?
282 .map_err(|err| anyhow!("{err}"))?;
283
284 Ok(output.into())
285 }
286 .boxed()
287 })
288 .await
289 }
290
291 async fn context_server_command(
292 &self,
293 context_server_id: Arc<str>,
294 project: Arc<dyn ProjectDelegate>,
295 ) -> Result<Command> {
296 self.call(|extension, store| {
297 async move {
298 let project_resource = store.data_mut().table().push(project)?;
299 let command = extension
300 .call_context_server_command(store, context_server_id.clone(), project_resource)
301 .await?
302 .map_err(|err| anyhow!("{err}"))?;
303 anyhow::Ok(command.into())
304 }
305 .boxed()
306 })
307 .await
308 }
309
310 async fn context_server_configuration(
311 &self,
312 context_server_id: Arc<str>,
313 project: Arc<dyn ProjectDelegate>,
314 ) -> Result<Option<ContextServerConfiguration>> {
315 self.call(|extension, store| {
316 async move {
317 let project_resource = store.data_mut().table().push(project)?;
318 let Some(configuration) = extension
319 .call_context_server_configuration(
320 store,
321 context_server_id.clone(),
322 project_resource,
323 )
324 .await?
325 .map_err(|err| anyhow!("{err}"))?
326 else {
327 return Ok(None);
328 };
329
330 Ok(Some(configuration.try_into()?))
331 }
332 .boxed()
333 })
334 .await
335 }
336
337 async fn suggest_docs_packages(&self, provider: Arc<str>) -> Result<Vec<String>> {
338 self.call(|extension, store| {
339 async move {
340 let packages = extension
341 .call_suggest_docs_packages(store, provider.as_ref())
342 .await?
343 .map_err(|err| anyhow!("{err:?}"))?;
344
345 Ok(packages)
346 }
347 .boxed()
348 })
349 .await
350 }
351
352 async fn index_docs(
353 &self,
354 provider: Arc<str>,
355 package_name: Arc<str>,
356 kv_store: Arc<dyn KeyValueStoreDelegate>,
357 ) -> Result<()> {
358 self.call(|extension, store| {
359 async move {
360 let kv_store_resource = store.data_mut().table().push(kv_store)?;
361 extension
362 .call_index_docs(
363 store,
364 provider.as_ref(),
365 package_name.as_ref(),
366 kv_store_resource,
367 )
368 .await?
369 .map_err(|err| anyhow!("{err:?}"))?;
370
371 anyhow::Ok(())
372 }
373 .boxed()
374 })
375 .await
376 }
377}
378
379pub struct WasmState {
380 manifest: Arc<ExtensionManifest>,
381 pub table: ResourceTable,
382 ctx: wasi::WasiCtx,
383 pub host: Arc<WasmHost>,
384}
385
386type MainThreadCall = Box<dyn Send + for<'a> FnOnce(&'a mut AsyncApp) -> LocalBoxFuture<'a, ()>>;
387
388type ExtensionCall = Box<
389 dyn Send + for<'a> FnOnce(&'a mut Extension, &'a mut Store<WasmState>) -> BoxFuture<'a, ()>,
390>;
391
392fn wasm_engine() -> wasmtime::Engine {
393 static WASM_ENGINE: OnceLock<wasmtime::Engine> = OnceLock::new();
394
395 WASM_ENGINE
396 .get_or_init(|| {
397 let mut config = wasmtime::Config::new();
398 config.wasm_component_model(true);
399 config.async_support(true);
400 wasmtime::Engine::new(&config).unwrap()
401 })
402 .clone()
403}
404
405impl WasmHost {
406 pub fn new(
407 fs: Arc<dyn Fs>,
408 http_client: Arc<dyn HttpClient>,
409 node_runtime: NodeRuntime,
410 proxy: Arc<ExtensionHostProxy>,
411 work_dir: PathBuf,
412 cx: &mut App,
413 ) -> Arc<Self> {
414 let (tx, mut rx) = mpsc::unbounded::<MainThreadCall>();
415 let task = cx.spawn(async move |cx| {
416 while let Some(message) = rx.next().await {
417 message(cx).await;
418 }
419 });
420 Arc::new(Self {
421 engine: wasm_engine(),
422 fs,
423 work_dir,
424 http_client,
425 node_runtime,
426 proxy,
427 release_channel: ReleaseChannel::global(cx),
428 _main_thread_message_task: task,
429 main_thread_message_tx: tx,
430 })
431 }
432
433 pub fn load_extension(
434 self: &Arc<Self>,
435 wasm_bytes: Vec<u8>,
436 manifest: &Arc<ExtensionManifest>,
437 executor: BackgroundExecutor,
438 ) -> Task<Result<WasmExtension>> {
439 let this = self.clone();
440 let manifest = manifest.clone();
441 executor.clone().spawn(async move {
442 let zed_api_version = parse_wasm_extension_version(&manifest.id, &wasm_bytes)?;
443
444 let component = Component::from_binary(&this.engine, &wasm_bytes)
445 .context("failed to compile wasm component")?;
446
447 let mut store = wasmtime::Store::new(
448 &this.engine,
449 WasmState {
450 ctx: this.build_wasi_ctx(&manifest).await?,
451 manifest: manifest.clone(),
452 table: ResourceTable::new(),
453 host: this.clone(),
454 },
455 );
456
457 let mut extension = Extension::instantiate_async(
458 &mut store,
459 this.release_channel,
460 zed_api_version,
461 &component,
462 )
463 .await?;
464
465 extension
466 .call_init_extension(&mut store)
467 .await
468 .context("failed to initialize wasm extension")?;
469
470 let (tx, mut rx) = mpsc::unbounded::<ExtensionCall>();
471 executor
472 .spawn(async move {
473 while let Some(call) = rx.next().await {
474 (call)(&mut extension, &mut store).await;
475 }
476 })
477 .detach();
478
479 Ok(WasmExtension {
480 manifest: manifest.clone(),
481 work_dir: this.work_dir.join(manifest.id.as_ref()).into(),
482 tx,
483 zed_api_version,
484 })
485 })
486 }
487
488 async fn build_wasi_ctx(&self, manifest: &Arc<ExtensionManifest>) -> Result<wasi::WasiCtx> {
489 let extension_work_dir = self.work_dir.join(manifest.id.as_ref());
490 self.fs
491 .create_dir(&extension_work_dir)
492 .await
493 .context("failed to create extension work dir")?;
494
495 let file_perms = wasi::FilePerms::all();
496 let dir_perms = wasi::DirPerms::all();
497
498 Ok(wasi::WasiCtxBuilder::new()
499 .inherit_stdio()
500 .preopened_dir(&extension_work_dir, ".", dir_perms, file_perms)?
501 .preopened_dir(
502 &extension_work_dir,
503 extension_work_dir.to_string_lossy(),
504 dir_perms,
505 file_perms,
506 )?
507 .env("PWD", extension_work_dir.to_string_lossy())
508 .env("RUST_BACKTRACE", "full")
509 .build())
510 }
511
512 pub fn writeable_path_from_extension(&self, id: &Arc<str>, path: &Path) -> Result<PathBuf> {
513 let extension_work_dir = self.work_dir.join(id.as_ref());
514 let path = normalize_path(&extension_work_dir.join(path));
515 if path.starts_with(&extension_work_dir) {
516 Ok(path)
517 } else {
518 Err(anyhow!("cannot write to path {}", path.display()))
519 }
520 }
521}
522
523pub fn parse_wasm_extension_version(
524 extension_id: &str,
525 wasm_bytes: &[u8],
526) -> Result<SemanticVersion> {
527 let mut version = None;
528
529 for part in wasmparser::Parser::new(0).parse_all(wasm_bytes) {
530 if let wasmparser::Payload::CustomSection(s) =
531 part.context("error parsing wasm extension")?
532 {
533 if s.name() == "zed:api-version" {
534 version = parse_wasm_extension_version_custom_section(s.data());
535 if version.is_none() {
536 bail!(
537 "extension {} has invalid zed:api-version section: {:?}",
538 extension_id,
539 s.data()
540 );
541 }
542 }
543 }
544 }
545
546 // The reason we wait until we're done parsing all of the Wasm bytes to return the version
547 // is to work around a panic that can happen inside of Wasmtime when the bytes are invalid.
548 //
549 // By parsing the entirety of the Wasm bytes before we return, we're able to detect this problem
550 // earlier as an `Err` rather than as a panic.
551 version.ok_or_else(|| anyhow!("extension {} has no zed:api-version section", extension_id))
552}
553
554fn parse_wasm_extension_version_custom_section(data: &[u8]) -> Option<SemanticVersion> {
555 if data.len() == 6 {
556 Some(SemanticVersion::new(
557 u16::from_be_bytes([data[0], data[1]]) as _,
558 u16::from_be_bytes([data[2], data[3]]) as _,
559 u16::from_be_bytes([data[4], data[5]]) as _,
560 ))
561 } else {
562 None
563 }
564}
565
566impl WasmExtension {
567 pub async fn load(
568 extension_dir: PathBuf,
569 manifest: &Arc<ExtensionManifest>,
570 wasm_host: Arc<WasmHost>,
571 cx: &AsyncApp,
572 ) -> Result<Self> {
573 let path = extension_dir.join("extension.wasm");
574
575 let mut wasm_file = wasm_host
576 .fs
577 .open_sync(&path)
578 .await
579 .context("failed to open wasm file")?;
580
581 let mut wasm_bytes = Vec::new();
582 wasm_file
583 .read_to_end(&mut wasm_bytes)
584 .context("failed to read wasm")?;
585
586 wasm_host
587 .load_extension(wasm_bytes, manifest, cx.background_executor().clone())
588 .await
589 .with_context(|| format!("failed to load wasm extension {}", manifest.id))
590 }
591
592 pub async fn call<T, Fn>(&self, f: Fn) -> T
593 where
594 T: 'static + Send,
595 Fn: 'static
596 + Send
597 + for<'a> FnOnce(&'a mut Extension, &'a mut Store<WasmState>) -> BoxFuture<'a, T>,
598 {
599 let (return_tx, return_rx) = oneshot::channel();
600 self.tx
601 .clone()
602 .unbounded_send(Box::new(move |extension, store| {
603 async {
604 let result = f(extension, store).await;
605 return_tx.send(result).ok();
606 }
607 .boxed()
608 }))
609 .expect("wasm extension channel should not be closed yet");
610 return_rx.await.expect("wasm extension channel")
611 }
612}
613
614impl WasmState {
615 fn on_main_thread<T, Fn>(&self, f: Fn) -> impl 'static + Future<Output = T>
616 where
617 T: 'static + Send,
618 Fn: 'static + Send + for<'a> FnOnce(&'a mut AsyncApp) -> LocalBoxFuture<'a, T>,
619 {
620 let (return_tx, return_rx) = oneshot::channel();
621 self.host
622 .main_thread_message_tx
623 .clone()
624 .unbounded_send(Box::new(move |cx| {
625 async {
626 let result = f(cx).await;
627 return_tx.send(result).ok();
628 }
629 .boxed_local()
630 }))
631 .expect("main thread message channel should not be closed yet");
632 async move { return_rx.await.expect("main thread message channel") }
633 }
634
635 fn work_dir(&self) -> PathBuf {
636 self.host.work_dir.join(self.manifest.id.as_ref())
637 }
638}
639
640impl wasi::WasiView for WasmState {
641 fn table(&mut self) -> &mut ResourceTable {
642 &mut self.table
643 }
644
645 fn ctx(&mut self) -> &mut wasi::WasiCtx {
646 &mut self.ctx
647 }
648}