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