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
359 let peer_identity = runtimelib::peer_identity_for_session(&session_id)?;
360 let shell_socket = runtimelib::create_client_shell_connection_with_identity(
361 &client_connection_info,
362 &session_id,
363 peer_identity.clone(),
364 )
365 .await?;
366
367 let control_socket =
368 runtimelib::create_client_control_connection(&client_connection_info, &session_id)
369 .await?;
370
371 let stdin_socket = runtimelib::create_client_stdin_connection_with_identity(
372 &client_connection_info,
373 &session_id,
374 peer_identity,
375 )
376 .await?;
377
378 let (request_tx, stdin_tx) = start_kernel_tasks(
379 session.clone(),
380 output_socket,
381 shell_socket,
382 control_socket,
383 stdin_socket,
384 cx,
385 );
386
387 let stderr = process.stderr.take();
388 cx.spawn(async move |_cx| {
389 if let Some(stderr) = stderr {
390 let reader = BufReader::new(stderr);
391 let mut lines = reader.lines();
392 while let Some(Ok(line)) = lines.next().await {
393 log::warn!("wsl kernel stderr: {}", line);
394 }
395 }
396 })
397 .detach();
398
399 let stdout = process.stdout.take();
400 cx.spawn(async move |_cx| {
401 if let Some(stdout) = stdout {
402 let reader = BufReader::new(stdout);
403 let mut lines = reader.lines();
404 while let Some(Ok(_line)) = lines.next().await {}
405 }
406 })
407 .detach();
408
409 let status = process.status();
410
411 let process_status_task = cx.spawn(async move |cx| {
412 let error_message = match status.await {
413 Ok(status) => {
414 if status.success() {
415 return;
416 }
417
418 format!("WSL kernel: kernel process exited with status: {:?}", status)
419 }
420 Err(err) => {
421 format!("WSL kernel: kernel process exited with error: {:?}", err)
422 }
423 };
424
425 session.update(cx, |session, cx| {
426 session.kernel_errored(error_message, cx);
427
428 cx.notify();
429 });
430 });
431
432 anyhow::Ok(Box::new(Self {
433 process,
434 request_tx,
435 stdin_tx,
436 working_directory,
437 _process_status_task: Some(process_status_task),
438 connection_path,
439 execution_state: ExecutionState::Idle,
440 kernel_info: None,
441 }) as Box<dyn RunningKernel>)
442 })
443 }
444}
445
446impl RunningKernel for WslRunningKernel {
447 fn request_tx(&self) -> mpsc::Sender<JupyterMessage> {
448 self.request_tx.clone()
449 }
450
451 fn stdin_tx(&self) -> mpsc::Sender<JupyterMessage> {
452 self.stdin_tx.clone()
453 }
454
455 fn working_directory(&self) -> &PathBuf {
456 &self.working_directory
457 }
458
459 fn execution_state(&self) -> &ExecutionState {
460 &self.execution_state
461 }
462
463 fn set_execution_state(&mut self, state: ExecutionState) {
464 self.execution_state = state;
465 }
466
467 fn kernel_info(&self) -> Option<&KernelInfoReply> {
468 self.kernel_info.as_ref()
469 }
470
471 fn set_kernel_info(&mut self, info: KernelInfoReply) {
472 self.kernel_info = Some(info);
473 }
474
475 fn force_shutdown(&mut self, _window: &mut Window, _cx: &mut App) -> Task<anyhow::Result<()>> {
476 self._process_status_task.take();
477 self.request_tx.close_channel();
478 self.process.kill().ok();
479 Task::ready(Ok(()))
480 }
481
482 fn kill(&mut self) {
483 self._process_status_task.take();
484 self.request_tx.close_channel();
485 self.process.kill().ok();
486 }
487}
488
489impl Drop for WslRunningKernel {
490 fn drop(&mut self) {
491 std::fs::remove_file(&self.connection_path).ok();
492 self.request_tx.close_channel();
493 self.process.kill().ok();
494 }
495}
496
497#[derive(serde::Deserialize)]
498struct LocalKernelSpecsResponse {
499 kernelspecs: std::collections::HashMap<String, LocalKernelSpec>,
500}
501
502#[derive(serde::Deserialize)]
503struct LocalKernelSpec {
504 spec: LocalKernelSpecContent,
505}
506
507#[derive(serde::Deserialize)]
508struct LocalKernelSpecContent {
509 argv: Vec<String>,
510 display_name: String,
511 language: String,
512 interrupt_mode: Option<String>,
513 env: Option<std::collections::HashMap<String, String>>,
514 metadata: Option<std::collections::HashMap<String, serde_json::Value>>,
515}
516
517pub async fn wsl_kernel_specifications(
518 background_executor: BackgroundExecutor,
519) -> Result<Vec<KernelSpecification>> {
520 let output = util::command::new_command("wsl")
521 .arg("-l")
522 .arg("-q")
523 .output()
524 .await;
525
526 if output.is_err() {
527 return Ok(Vec::new());
528 }
529
530 let output = output.unwrap();
531 if !output.status.success() {
532 return Ok(Vec::new());
533 }
534
535 // wsl output is often UTF-16LE, but -l -q might be simpler or just ASCII compatible if not using weird charsets.
536 // However, on Windows, wsl often outputs UTF-16LE.
537 // We can try to detect or use from_utf16 if valid, or just use String::from_utf8_lossy and see.
538 // Actually, `smol::process` on Windows might receive bytes that are UTF-16LE if wsl writes that.
539 // But typically terminal output for wsl is UTF-16.
540 // Let's try to parse as UTF-16LE if it looks like it (BOM or just 00 bytes).
541
542 let stdout = output.stdout;
543 let distros_str = if stdout.len() >= 2 && stdout[1] == 0 {
544 // likely UTF-16LE
545 let u16s: Vec<u16> = stdout
546 .chunks_exact(2)
547 .map(|c| u16::from_le_bytes([c[0], c[1]]))
548 .collect();
549 String::from_utf16_lossy(&u16s)
550 } else {
551 String::from_utf8_lossy(&stdout).to_string()
552 };
553
554 let distros: Vec<String> = distros_str
555 .lines()
556 .map(|line| line.trim().to_string())
557 .filter(|line| !line.is_empty())
558 .collect();
559
560 let tasks = distros.into_iter().map(|distro| {
561 background_executor.spawn(async move {
562 let output = util::command::new_command("wsl")
563 .arg("-d")
564 .arg(&distro)
565 .arg("bash")
566 .arg("-l")
567 .arg("-c")
568 .arg("jupyter kernelspec list --json")
569 .output()
570 .await;
571
572 if let Ok(output) = output {
573 if output.status.success() {
574 let json_str = String::from_utf8_lossy(&output.stdout);
575 // Use local permissive struct instead of strict KernelSpecsResponse from jupyter-protocol
576 if let Ok(specs_response) =
577 serde_json::from_str::<LocalKernelSpecsResponse>(&json_str)
578 {
579 return specs_response
580 .kernelspecs
581 .into_iter()
582 .map(|(name, spec)| {
583 KernelSpecification::WslRemote(WslKernelSpecification {
584 name,
585 kernelspec: jupyter_protocol::JupyterKernelspec {
586 argv: spec.spec.argv,
587 display_name: spec.spec.display_name,
588 language: spec.spec.language,
589 interrupt_mode: spec.spec.interrupt_mode,
590 env: spec.spec.env,
591 metadata: spec.spec.metadata,
592 },
593 distro: distro.clone(),
594 })
595 })
596 .collect::<Vec<_>>();
597 } else if let Err(e) =
598 serde_json::from_str::<LocalKernelSpecsResponse>(&json_str)
599 {
600 log::error!(
601 "wsl_kernel_specifications parse error: {} \nJSON: {}",
602 e,
603 json_str
604 );
605 }
606 } else {
607 log::error!("wsl_kernel_specifications command failed");
608 }
609 } else if let Err(e) = output {
610 log::error!("wsl_kernel_specifications command execution failed: {}", e);
611 }
612
613 Vec::new()
614 })
615 });
616
617 let specs: Vec<_> = futures::future::join_all(tasks)
618 .await
619 .into_iter()
620 .flatten()
621 .collect();
622
623 Ok(specs)
624}