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