1use super::{
2 KernelSession, KernelSpecification, RunningKernel, WslKernelSpecification, start_kernel_tasks,
3};
4use anyhow::{Context as _, Result};
5use futures::{
6 AsyncBufReadExt as _, StreamExt as _,
7 channel::mpsc::{self},
8 io::BufReader,
9};
10use gpui::{App, BackgroundExecutor, Entity, EntityId, Task, Window};
11use jupyter_protocol::{
12 ExecutionState, JupyterMessage, KernelInfoReply,
13 connection_info::{ConnectionInfo, Transport},
14};
15use project::Fs;
16use runtimelib::dirs;
17use smol::net::TcpListener;
18use std::{
19 fmt::Debug,
20 net::{IpAddr, Ipv4Addr, SocketAddr},
21 path::PathBuf,
22 sync::Arc,
23};
24use uuid::Uuid;
25
26// Find a set of open ports. This creates a listener with port set to 0. The listener will be closed at the end when it goes out of scope.
27// There's a race condition between closing the ports and usage by a kernel, but it's inherent to the Jupyter protocol.
28async fn peek_ports(ip: IpAddr) -> Result<[u16; 5]> {
29 let mut addr_zeroport: SocketAddr = SocketAddr::new(ip, 0);
30 addr_zeroport.set_port(0);
31 let mut ports: [u16; 5] = [0; 5];
32 for i in 0..5 {
33 let listener = TcpListener::bind(addr_zeroport).await?;
34 let addr = listener.local_addr()?;
35 ports[i] = addr.port();
36 }
37 Ok(ports)
38}
39
40pub struct WslRunningKernel {
41 pub process: util::command::Child,
42 connection_path: PathBuf,
43 _process_status_task: Option<Task<()>>,
44 pub working_directory: PathBuf,
45 pub request_tx: mpsc::Sender<JupyterMessage>,
46 pub stdin_tx: mpsc::Sender<JupyterMessage>,
47 pub execution_state: ExecutionState,
48 pub kernel_info: Option<KernelInfoReply>,
49}
50
51impl Debug for WslRunningKernel {
52 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
53 f.debug_struct("WslRunningKernel")
54 .field("process", &self.process)
55 .finish()
56 }
57}
58
59impl WslRunningKernel {
60 pub fn new<S: KernelSession + 'static>(
61 kernel_specification: WslKernelSpecification,
62 entity_id: EntityId,
63 working_directory: PathBuf,
64 fs: Arc<dyn Fs>,
65 session: Entity<S>,
66 window: &mut Window,
67 cx: &mut App,
68 ) -> Task<Result<Box<dyn RunningKernel>>> {
69 window.spawn(cx, async move |cx| {
70 // For WSL2, we need to get the WSL VM's IP address to connect to it
71 // because WSL2 runs in a lightweight VM with its own network namespace.
72 // The kernel will bind to 127.0.0.1 inside WSL, and we connect to localhost.
73 // WSL2 localhost forwarding handles the rest.
74 let bind_ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
75
76 // Use 127.0.0.1 and rely on WSL 2 localhost forwarding.
77 // This avoids issues where the VM IP is unreachable or binding fails on Windows.
78 let connect_ip = "127.0.0.1".to_string();
79
80 let ports = peek_ports(bind_ip).await?;
81
82 let connection_info = ConnectionInfo {
83 transport: Transport::TCP,
84 ip: bind_ip.to_string(),
85 stdin_port: ports[0],
86 control_port: ports[1],
87 hb_port: ports[2],
88 shell_port: ports[3],
89 iopub_port: ports[4],
90 signature_scheme: "hmac-sha256".to_string(),
91 key: uuid::Uuid::new_v4().to_string(),
92 kernel_name: Some(format!("zed-wsl-{}", kernel_specification.name)),
93 };
94
95 let runtime_dir = dirs::runtime_dir();
96 fs.create_dir(&runtime_dir)
97 .await
98 .with_context(|| format!("Failed to create jupyter runtime dir {runtime_dir:?}"))?;
99 let connection_path = runtime_dir.join(format!("kernel-zed-wsl-{entity_id}.json"));
100 let content = serde_json::to_string(&connection_info)?;
101 fs.atomic_write(connection_path.clone(), content).await?;
102
103 // Convert connection_path to WSL path
104 // yeah we can't assume this is available on WSL.
105 // running `wsl -d <distro> wslpath -u <windows_path>`
106 let mut wslpath_cmd = util::command::new_command("wsl");
107
108 // On Windows, passing paths with backslashes to wsl.exe can sometimes cause
109 // escaping issues or be misinterpreted. Converting to forward slashes is safer
110 // and often accepted by wslpath.
111 let connection_path_str = connection_path.to_string_lossy().replace('\\', "/");
112
113 wslpath_cmd
114 .arg("-d")
115 .arg(&kernel_specification.distro)
116 .arg("wslpath")
117 .arg("-u")
118 .arg(&connection_path_str);
119
120 let output = wslpath_cmd.output().await?;
121 if !output.status.success() {
122 anyhow::bail!("Failed to convert path to WSL path: {:?}", output);
123 }
124 let wsl_connection_path = String::from_utf8_lossy(&output.stdout).trim().to_string();
125
126 // Construct the kernel command
127 // The kernel spec argv might have absolute paths valid INSIDE WSL.
128 // We need to run inside WSL.
129 // `wsl -d <distro> --exec <argv0> <argv1> ...`
130 // But we need to replace {connection_file} with wsl_connection_path.
131
132 let argv = kernel_specification.kernelspec.argv;
133 anyhow::ensure!(
134 !argv.is_empty(),
135 "Empty argv in kernelspec {}",
136 kernel_specification.name
137 );
138
139 let working_directory_str = working_directory.to_string_lossy().replace('\\', "/");
140
141 let wsl_working_directory = if working_directory_str.starts_with('/') {
142 // If path starts with /, assume it is already a WSL path (e.g. /home/user)
143 Some(working_directory_str)
144 } else {
145 let mut wslpath_wd_cmd = util::command::new_command("wsl");
146 wslpath_wd_cmd
147 .arg("-d")
148 .arg(&kernel_specification.distro)
149 .arg("wslpath")
150 .arg("-u")
151 .arg(&working_directory_str);
152
153 let wd_output = wslpath_wd_cmd.output().await;
154 if let Ok(output) = wd_output {
155 if output.status.success() {
156 Some(String::from_utf8_lossy(&output.stdout).trim().to_string())
157 } else {
158 None
159 }
160 } else {
161 None
162 }
163 };
164
165 // If we couldn't convert the working directory or it's a temp directory,
166 // and the kernel spec uses a relative path (like .venv/bin/python),
167 // we need to handle this better. For now, let's use the converted path
168 // if available, otherwise we'll rely on WSL's default home directory.
169
170 let mut cmd = util::command::new_command("wsl");
171 cmd.arg("-d").arg(&kernel_specification.distro);
172
173 // Set CWD for the host process to a safe location to avoid "Directory name is invalid"
174 // if the project root is a path not supported by Windows CWD (e.g. UNC path for some tools).
175 cmd.current_dir(std::env::temp_dir());
176
177 if let Some(wd) = wsl_working_directory.as_ref() {
178 cmd.arg("--cd").arg(wd);
179 }
180
181 // Build the command to run inside WSL
182 // We use bash -lc to run in a login shell for proper environment setup
183 let mut kernel_args: Vec<String> = Vec::new();
184
185 if let Some(env) = &kernel_specification.kernelspec.env {
186 if !env.is_empty() {
187 kernel_args.push("env".to_string());
188 for (k, v) in env {
189 kernel_args.push(format!("{}={}", k, v));
190 }
191 }
192 }
193
194 for arg in argv {
195 if arg == "{connection_file}" {
196 kernel_args.push(wsl_connection_path.clone());
197 } else {
198 kernel_args.push(arg.clone());
199 }
200 }
201
202 // because first command is python/python3 we need make sure it's present in the env
203 let first_cmd = kernel_args.first().map(|arg| {
204 arg.split_whitespace().next().unwrap_or(arg)
205 });
206
207 let needs_python_resolution = first_cmd.map_or(false, |cmd| {
208 cmd == "python" || cmd == "python3" || !cmd.starts_with('/')
209 });
210
211 let shell_command = if needs_python_resolution {
212 // 1. Check for .venv/bin/python or .venv/bin/python3 in working directory
213 // 2. Fall back to system python3 or python
214 let rest_args: Vec<String> = kernel_args.iter().skip(1).cloned().collect();
215 let rest_string = rest_args
216 .iter()
217 .map(|arg| {
218 if arg.contains(' ') || arg.contains('\'') || arg.contains('"') {
219 format!("'{}'", arg.replace('\'', "'\\''"))
220 } else {
221 arg.clone()
222 }
223 })
224 .collect::<Vec<_>>()
225 .join(" ");
226
227 let cd_command = if let Some(wd) = wsl_working_directory.as_ref() {
228 format!("cd '{}' && ", wd.replace('\'', "'\\''"))
229 } else {
230 String::new()
231 };
232 // TODO: find a better way to debug missing python issues in WSL
233
234 format!(
235 "set -e; \
236 {} \
237 echo \"Working directory: $(pwd)\" >&2; \
238 if [ -x .venv/bin/python ]; then \
239 echo \"Found .venv/bin/python\" >&2; \
240 exec .venv/bin/python {}; \
241 elif [ -x .venv/bin/python3 ]; then \
242 echo \"Found .venv/bin/python3\" >&2; \
243 exec .venv/bin/python3 {}; \
244 elif command -v python3 >/dev/null 2>&1; then \
245 echo \"Found system python3\" >&2; \
246 exec python3 {}; \
247 elif command -v python >/dev/null 2>&1; then \
248 echo \"Found system python\" >&2; \
249 exec python {}; \
250 else \
251 echo 'Error: Python not found in .venv or PATH' >&2; \
252 echo 'Contents of current directory:' >&2; \
253 ls -la >&2; \
254 echo 'PATH:' \"$PATH\" >&2; \
255 exit 127; \
256 fi",
257 cd_command, rest_string, rest_string, rest_string, rest_string
258 )
259 } else {
260 kernel_args
261 .iter()
262 .map(|arg| {
263 if arg.contains(' ') || arg.contains('\'') || arg.contains('"') {
264 format!("'{}'", arg.replace('\'', "'\\''"))
265 } else {
266 arg.clone()
267 }
268 })
269 .collect::<Vec<_>>()
270 .join(" ")
271 };
272
273 cmd.arg("bash")
274 .arg("-l")
275 .arg("-c")
276 .arg(&shell_command);
277
278 let mut process = cmd
279 .stdout(util::command::Stdio::piped())
280 .stderr(util::command::Stdio::piped())
281 .stdin(util::command::Stdio::piped())
282 .kill_on_drop(true)
283 .spawn()
284 .context("failed to start the kernel process")?;
285
286 let session_id = Uuid::new_v4().to_string();
287
288 let mut client_connection_info = connection_info.clone();
289 client_connection_info.ip = connect_ip.clone();
290
291 // Give the kernel a moment to start and bind to ports.
292 // WSL kernel startup can be slow, I am not sure if this is because of my testing environment
293 // or inherent to WSL. We can improve this later with better readiness checks.
294 cx.background_executor()
295 .timer(std::time::Duration::from_secs(2))
296 .await;
297
298 match process.try_status() {
299 Ok(Some(status)) => {
300 let mut stderr_content = String::new();
301 if let Some(mut stderr) = process.stderr.take() {
302 use futures::AsyncReadExt;
303 let mut buf = Vec::new();
304 if stderr.read_to_end(&mut buf).await.is_ok() {
305 stderr_content = String::from_utf8_lossy(&buf).to_string();
306 }
307 }
308
309 let mut stdout_content = String::new();
310 if let Some(mut stdout) = process.stdout.take() {
311 use futures::AsyncReadExt;
312 let mut buf = Vec::new();
313 if stdout.read_to_end(&mut buf).await.is_ok() {
314 stdout_content = String::from_utf8_lossy(&buf).to_string();
315 }
316 }
317
318 anyhow::bail!(
319 "WSL kernel process exited prematurely with status: {:?}\nstderr: {}\nstdout: {}",
320 status,
321 stderr_content,
322 stdout_content
323 );
324 }
325 Ok(None) => {}
326 Err(_) => {}
327 }
328
329 let output_socket = runtimelib::create_client_iopub_connection(
330 &client_connection_info,
331 "",
332 &session_id,
333 )
334 .await?;
335
336 let peer_identity = runtimelib::peer_identity_for_session(&session_id)?;
337 let shell_socket = runtimelib::create_client_shell_connection_with_identity(
338 &client_connection_info,
339 &session_id,
340 peer_identity.clone(),
341 )
342 .await?;
343
344 let control_socket =
345 runtimelib::create_client_control_connection(&client_connection_info, &session_id)
346 .await?;
347
348 let stdin_socket = runtimelib::create_client_stdin_connection_with_identity(
349 &client_connection_info,
350 &session_id,
351 peer_identity,
352 )
353 .await?;
354
355 let (request_tx, stdin_tx) = start_kernel_tasks(
356 session.clone(),
357 output_socket,
358 shell_socket,
359 control_socket,
360 stdin_socket,
361 cx,
362 );
363
364 let stderr = process.stderr.take();
365 cx.spawn(async move |_cx| {
366 if let Some(stderr) = stderr {
367 let reader = BufReader::new(stderr);
368 let mut lines = reader.lines();
369 while let Some(Ok(line)) = lines.next().await {
370 log::warn!("wsl kernel stderr: {}", line);
371 }
372 }
373 })
374 .detach();
375
376 let stdout = process.stdout.take();
377 cx.spawn(async move |_cx| {
378 if let Some(stdout) = stdout {
379 let reader = BufReader::new(stdout);
380 let mut lines = reader.lines();
381 while let Some(Ok(_line)) = lines.next().await {}
382 }
383 })
384 .detach();
385
386 let status = process.status();
387
388 let process_status_task = cx.spawn(async move |cx| {
389 let error_message = match status.await {
390 Ok(status) => {
391 if status.success() {
392 return;
393 }
394
395 format!("WSL kernel: kernel process exited with status: {:?}", status)
396 }
397 Err(err) => {
398 format!("WSL kernel: kernel process exited with error: {:?}", err)
399 }
400 };
401
402 session.update(cx, |session, cx| {
403 session.kernel_errored(error_message, cx);
404
405 cx.notify();
406 });
407 });
408
409 anyhow::Ok(Box::new(Self {
410 process,
411 request_tx,
412 stdin_tx,
413 working_directory,
414 _process_status_task: Some(process_status_task),
415 connection_path,
416 execution_state: ExecutionState::Idle,
417 kernel_info: None,
418 }) as Box<dyn RunningKernel>)
419 })
420 }
421}
422
423impl RunningKernel for WslRunningKernel {
424 fn request_tx(&self) -> mpsc::Sender<JupyterMessage> {
425 self.request_tx.clone()
426 }
427
428 fn stdin_tx(&self) -> mpsc::Sender<JupyterMessage> {
429 self.stdin_tx.clone()
430 }
431
432 fn working_directory(&self) -> &PathBuf {
433 &self.working_directory
434 }
435
436 fn execution_state(&self) -> &ExecutionState {
437 &self.execution_state
438 }
439
440 fn set_execution_state(&mut self, state: ExecutionState) {
441 self.execution_state = state;
442 }
443
444 fn kernel_info(&self) -> Option<&KernelInfoReply> {
445 self.kernel_info.as_ref()
446 }
447
448 fn set_kernel_info(&mut self, info: KernelInfoReply) {
449 self.kernel_info = Some(info);
450 }
451
452 fn force_shutdown(&mut self, _window: &mut Window, _cx: &mut App) -> Task<anyhow::Result<()>> {
453 self._process_status_task.take();
454 self.request_tx.close_channel();
455 self.process.kill().ok();
456 Task::ready(Ok(()))
457 }
458
459 fn kill(&mut self) {
460 self._process_status_task.take();
461 self.request_tx.close_channel();
462 self.process.kill().ok();
463 }
464}
465
466impl Drop for WslRunningKernel {
467 fn drop(&mut self) {
468 std::fs::remove_file(&self.connection_path).ok();
469 self.request_tx.close_channel();
470 self.process.kill().ok();
471 }
472}
473
474#[derive(serde::Deserialize)]
475struct LocalKernelSpecsResponse {
476 kernelspecs: std::collections::HashMap<String, LocalKernelSpec>,
477}
478
479#[derive(serde::Deserialize)]
480struct LocalKernelSpec {
481 spec: LocalKernelSpecContent,
482}
483
484#[derive(serde::Deserialize)]
485struct LocalKernelSpecContent {
486 argv: Vec<String>,
487 display_name: String,
488 language: String,
489 interrupt_mode: Option<String>,
490 env: Option<std::collections::HashMap<String, String>>,
491 metadata: Option<std::collections::HashMap<String, serde_json::Value>>,
492}
493
494pub async fn wsl_kernel_specifications(
495 background_executor: BackgroundExecutor,
496) -> Result<Vec<KernelSpecification>> {
497 let output = util::command::new_command("wsl")
498 .arg("-l")
499 .arg("-q")
500 .output()
501 .await;
502
503 if output.is_err() {
504 return Ok(Vec::new());
505 }
506
507 let output = output.unwrap();
508 if !output.status.success() {
509 return Ok(Vec::new());
510 }
511
512 // wsl output is often UTF-16LE, but -l -q might be simpler or just ASCII compatible if not using weird charsets.
513 // However, on Windows, wsl often outputs UTF-16LE.
514 // We can try to detect or use from_utf16 if valid, or just use String::from_utf8_lossy and see.
515 // Actually, `smol::process` on Windows might receive bytes that are UTF-16LE if wsl writes that.
516 // But typically terminal output for wsl is UTF-16.
517 // Let's try to parse as UTF-16LE if it looks like it (BOM or just 00 bytes).
518
519 let stdout = output.stdout;
520 let distros_str = if stdout.len() >= 2 && stdout[1] == 0 {
521 // likely UTF-16LE
522 let u16s: Vec<u16> = stdout
523 .chunks_exact(2)
524 .map(|c| u16::from_le_bytes([c[0], c[1]]))
525 .collect();
526 String::from_utf16_lossy(&u16s)
527 } else {
528 String::from_utf8_lossy(&stdout).to_string()
529 };
530
531 let distros: Vec<String> = distros_str
532 .lines()
533 .map(|line| line.trim().to_string())
534 .filter(|line| !line.is_empty())
535 .collect();
536
537 let tasks = distros.into_iter().map(|distro| {
538 background_executor.spawn(async move {
539 let output = util::command::new_command("wsl")
540 .arg("-d")
541 .arg(&distro)
542 .arg("bash")
543 .arg("-l")
544 .arg("-c")
545 .arg("jupyter kernelspec list --json")
546 .output()
547 .await;
548
549 if let Ok(output) = output {
550 if output.status.success() {
551 let json_str = String::from_utf8_lossy(&output.stdout);
552 // Use local permissive struct instead of strict KernelSpecsResponse from jupyter-protocol
553 if let Ok(specs_response) =
554 serde_json::from_str::<LocalKernelSpecsResponse>(&json_str)
555 {
556 return specs_response
557 .kernelspecs
558 .into_iter()
559 .map(|(name, spec)| {
560 KernelSpecification::WslRemote(WslKernelSpecification {
561 name,
562 kernelspec: jupyter_protocol::JupyterKernelspec {
563 argv: spec.spec.argv,
564 display_name: spec.spec.display_name,
565 language: spec.spec.language,
566 interrupt_mode: spec.spec.interrupt_mode,
567 env: spec.spec.env,
568 metadata: spec.spec.metadata,
569 },
570 distro: distro.clone(),
571 })
572 })
573 .collect::<Vec<_>>();
574 }
575 }
576 }
577
578 Vec::new()
579 })
580 });
581
582 let specs: Vec<_> = futures::future::join_all(tasks)
583 .await
584 .into_iter()
585 .flatten()
586 .collect();
587
588 Ok(specs)
589}