From 278032c6b815aea8ce749af5e1580ddd13043a5c Mon Sep 17 00:00:00 2001
From: Lukas Wirth <lukas@zed.dev>
Date: Thu, 23 Oct 2025 09:41:05 +0200
Subject: [PATCH] extension_host: Run extensions on the tokio threadpool
 (#40936)

Fixes ZED-12D

`wasmtime_wasi` might call into tokio futures (to sleep for example)
which requires access to the tokio runtime. So we are required to run
these extensions in the tokio thread pool

Release Notes:

- Fixed extensions causing zed to occasionally panic
---
 Cargo.lock                                    |  1 +
 crates/extension_host/Cargo.toml              |  1 +
 .../extension_compilation_benchmark.rs        |  3 +-
 .../src/extension_store_test.rs               |  1 +
 crates/extension_host/src/wasm_host.rs        | 42 ++++++++++++-------
 5 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 48db1977ef..eb5df52724 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5883,6 +5883,7 @@ dependencies = [
  "fs",
  "futures 0.3.31",
  "gpui",
+ "gpui_tokio",
  "http_client",
  "language",
  "language_extension",
diff --git a/crates/extension_host/Cargo.toml b/crates/extension_host/Cargo.toml
index 42dcdd3a6f..16cbd9ac0c 100644
--- a/crates/extension_host/Cargo.toml
+++ b/crates/extension_host/Cargo.toml
@@ -27,6 +27,7 @@ extension.workspace = true
 fs.workspace = true
 futures.workspace = true
 gpui.workspace = true
+gpui_tokio.workspace = true
 http_client.workspace = true
 language.workspace = true
 log.workspace = true
diff --git a/crates/extension_host/benches/extension_compilation_benchmark.rs b/crates/extension_host/benches/extension_compilation_benchmark.rs
index 6f0897af6e..309e089758 100644
--- a/crates/extension_host/benches/extension_compilation_benchmark.rs
+++ b/crates/extension_host/benches/extension_compilation_benchmark.rs
@@ -19,6 +19,7 @@ use util::test::TempTree;
 
 fn extension_benchmarks(c: &mut Criterion) {
     let cx = init();
+    cx.update(gpui_tokio::init);
 
     let mut group = c.benchmark_group("load");
 
@@ -37,7 +38,7 @@ fn extension_benchmarks(c: &mut Criterion) {
             |wasm_bytes| {
                 let _extension = cx
                     .executor()
-                    .block(wasm_host.load_extension(wasm_bytes, &manifest, cx.executor()))
+                    .block(wasm_host.load_extension(wasm_bytes, &manifest, &cx.to_async()))
                     .unwrap();
             },
             BatchSize::SmallInput,
diff --git a/crates/extension_host/src/extension_store_test.rs b/crates/extension_host/src/extension_store_test.rs
index 855077bcf8..509edc6845 100644
--- a/crates/extension_host/src/extension_store_test.rs
+++ b/crates/extension_host/src/extension_store_test.rs
@@ -868,5 +868,6 @@ fn init_test(cx: &mut TestAppContext) {
         Project::init_settings(cx);
         ExtensionSettings::register(cx);
         language::init(cx);
+        gpui_tokio::init(cx);
     });
 }
diff --git a/crates/extension_host/src/wasm_host.rs b/crates/extension_host/src/wasm_host.rs
index f77258e895..22d11732a7 100644
--- a/crates/extension_host/src/wasm_host.rs
+++ b/crates/extension_host/src/wasm_host.rs
@@ -591,11 +591,12 @@ impl WasmHost {
         self: &Arc<Self>,
         wasm_bytes: Vec<u8>,
         manifest: &Arc<ExtensionManifest>,
-        executor: BackgroundExecutor,
+        cx: &AsyncApp,
     ) -> Task<Result<WasmExtension>> {
         let this = self.clone();
         let manifest = manifest.clone();
-        executor.clone().spawn(async move {
+        let executor = cx.background_executor().clone();
+        let load_extension_task = async move {
             let zed_api_version = parse_wasm_extension_version(&manifest.id, &wasm_bytes)?;
 
             let component = Component::from_binary(&this.engine, &wasm_bytes)
@@ -632,20 +633,29 @@ impl WasmHost {
                 .context("failed to initialize wasm extension")?;
 
             let (tx, mut rx) = mpsc::unbounded::<ExtensionCall>();
-            executor
-                .spawn(async move {
-                    while let Some(call) = rx.next().await {
-                        (call)(&mut extension, &mut store).await;
-                    }
-                })
-                .detach();
+            let extension_task = async move {
+                while let Some(call) = rx.next().await {
+                    (call)(&mut extension, &mut store).await;
+                }
+            };
 
-            Ok(WasmExtension {
-                manifest: manifest.clone(),
-                work_dir: this.work_dir.join(manifest.id.as_ref()).into(),
-                tx,
-                zed_api_version,
-            })
+            anyhow::Ok((
+                extension_task,
+                WasmExtension {
+                    manifest: manifest.clone(),
+                    work_dir: this.work_dir.join(manifest.id.as_ref()).into(),
+                    tx,
+                    zed_api_version,
+                },
+            ))
+        };
+        cx.spawn(async move |cx| {
+            let (extension_task, extension) = load_extension_task.await?;
+            // we need to run run the task in an extension context as wasmtime_wasi may
+            // call into tokio, accessing its runtime handle
+            gpui_tokio::Tokio::spawn(cx, extension_task)?.detach();
+
+            Ok(extension)
         })
     }
 
@@ -747,7 +757,7 @@ impl WasmExtension {
             .context("failed to read wasm")?;
 
         wasm_host
-            .load_extension(wasm_bytes, manifest, cx.background_executor().clone())
+            .load_extension(wasm_bytes, manifest, cx)
             .await
             .with_context(|| format!("failed to load wasm extension {}", manifest.id))
     }