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