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 quote_posix_shell_arguments(&kernel_args)?
278 };
279
280 cmd.arg("bash")
281 .arg("-l")
282 .arg("-c")
283 .arg(&shell_command);
284
285 let mut process = cmd
286 .stdout(util::command::Stdio::piped())
287 .stderr(util::command::Stdio::piped())
288 .stdin(util::command::Stdio::piped())
289 .kill_on_drop(true)
290 .spawn()
291 .context("failed to start the kernel process")?;
292
293 let session_id = Uuid::new_v4().to_string();
294
295 let mut client_connection_info = connection_info.clone();
296 client_connection_info.ip = connect_ip.clone();
297
298 // Give the kernel a moment to start and bind to ports.
299 // WSL kernel startup can be slow, I am not sure if this is because of my testing environment
300 // or inherent to WSL. We can improve this later with better readiness checks.
301 cx.background_executor()
302 .timer(std::time::Duration::from_secs(2))
303 .await;
304
305 match process.try_status() {
306 Ok(Some(status)) => {
307 let mut stderr_content = String::new();
308 if let Some(mut stderr) = process.stderr.take() {
309 use futures::AsyncReadExt;
310 let mut buf = Vec::new();
311 if stderr.read_to_end(&mut buf).await.is_ok() {
312 stderr_content = String::from_utf8_lossy(&buf).to_string();
313 }
314 }
315
316 let mut stdout_content = String::new();
317 if let Some(mut stdout) = process.stdout.take() {
318 use futures::AsyncReadExt;
319 let mut buf = Vec::new();
320 if stdout.read_to_end(&mut buf).await.is_ok() {
321 stdout_content = String::from_utf8_lossy(&buf).to_string();
322 }
323 }
324
325 anyhow::bail!(
326 "WSL kernel process exited prematurely with status: {:?}\nstderr: {}\nstdout: {}",
327 status,
328 stderr_content,
329 stdout_content
330 );
331 }
332 Ok(None) => {}
333 Err(_) => {}
334 }
335
336 let output_socket = runtimelib::create_client_iopub_connection(
337 &client_connection_info,
338 "",
339 &session_id,
340 )
341 .await?;
342
343 let peer_identity = runtimelib::peer_identity_for_session(&session_id)?;
344 let shell_socket = runtimelib::create_client_shell_connection_with_identity(
345 &client_connection_info,
346 &session_id,
347 peer_identity.clone(),
348 )
349 .await?;
350
351 let control_socket =
352 runtimelib::create_client_control_connection(&client_connection_info, &session_id)
353 .await?;
354
355 let stdin_socket = runtimelib::create_client_stdin_connection_with_identity(
356 &client_connection_info,
357 &session_id,
358 peer_identity,
359 )
360 .await?;
361
362 let (request_tx, stdin_tx) = start_kernel_tasks(
363 session.clone(),
364 output_socket,
365 shell_socket,
366 control_socket,
367 stdin_socket,
368 cx,
369 );
370
371 let stderr = process.stderr.take();
372 cx.spawn(async move |_cx| {
373 if let Some(stderr) = stderr {
374 let reader = BufReader::new(stderr);
375 let mut lines = reader.lines();
376 while let Some(Ok(line)) = lines.next().await {
377 log::warn!("wsl kernel stderr: {}", line);
378 }
379 }
380 })
381 .detach();
382
383 let stdout = process.stdout.take();
384 cx.spawn(async move |_cx| {
385 if let Some(stdout) = stdout {
386 let reader = BufReader::new(stdout);
387 let mut lines = reader.lines();
388 while let Some(Ok(_line)) = lines.next().await {}
389 }
390 })
391 .detach();
392
393 let status = process.status();
394
395 let process_status_task = cx.spawn(async move |cx| {
396 let error_message = match status.await {
397 Ok(status) => {
398 if status.success() {
399 return;
400 }
401
402 format!("WSL kernel: kernel process exited with status: {:?}", status)
403 }
404 Err(err) => {
405 format!("WSL kernel: kernel process exited with error: {:?}", err)
406 }
407 };
408
409 session.update(cx, |session, cx| {
410 session.kernel_errored(error_message, cx);
411
412 cx.notify();
413 });
414 });
415
416 anyhow::Ok(Box::new(Self {
417 process,
418 request_tx,
419 stdin_tx,
420 working_directory,
421 _process_status_task: Some(process_status_task),
422 connection_path,
423 execution_state: ExecutionState::Idle,
424 kernel_info: None,
425 }) as Box<dyn RunningKernel>)
426 })
427 }
428}
429
430impl RunningKernel for WslRunningKernel {
431 fn request_tx(&self) -> mpsc::Sender<JupyterMessage> {
432 self.request_tx.clone()
433 }
434
435 fn stdin_tx(&self) -> mpsc::Sender<JupyterMessage> {
436 self.stdin_tx.clone()
437 }
438
439 fn working_directory(&self) -> &PathBuf {
440 &self.working_directory
441 }
442
443 fn execution_state(&self) -> &ExecutionState {
444 &self.execution_state
445 }
446
447 fn set_execution_state(&mut self, state: ExecutionState) {
448 self.execution_state = state;
449 }
450
451 fn kernel_info(&self) -> Option<&KernelInfoReply> {
452 self.kernel_info.as_ref()
453 }
454
455 fn set_kernel_info(&mut self, info: KernelInfoReply) {
456 self.kernel_info = Some(info);
457 }
458
459 fn force_shutdown(&mut self, _window: &mut Window, _cx: &mut App) -> Task<anyhow::Result<()>> {
460 self._process_status_task.take();
461 self.request_tx.close_channel();
462 self.process.kill().ok();
463 Task::ready(Ok(()))
464 }
465
466 fn kill(&mut self) {
467 self._process_status_task.take();
468 self.request_tx.close_channel();
469 self.process.kill().ok();
470 }
471}
472
473impl Drop for WslRunningKernel {
474 fn drop(&mut self) {
475 std::fs::remove_file(&self.connection_path).ok();
476 self.request_tx.close_channel();
477 self.process.kill().ok();
478 }
479}
480
481#[derive(serde::Deserialize)]
482struct LocalKernelSpecsResponse {
483 kernelspecs: std::collections::HashMap<String, LocalKernelSpec>,
484}
485
486#[derive(serde::Deserialize)]
487struct LocalKernelSpec {
488 spec: LocalKernelSpecContent,
489}
490
491#[derive(serde::Deserialize)]
492struct LocalKernelSpecContent {
493 argv: Vec<String>,
494 display_name: String,
495 language: String,
496 interrupt_mode: Option<String>,
497 env: Option<std::collections::HashMap<String, String>>,
498 metadata: Option<std::collections::HashMap<String, serde_json::Value>>,
499}
500
501pub async fn wsl_kernel_specifications(
502 background_executor: BackgroundExecutor,
503) -> Result<Vec<KernelSpecification>> {
504 let output = util::command::new_command("wsl")
505 .arg("-l")
506 .arg("-q")
507 .output()
508 .await;
509
510 if output.is_err() {
511 return Ok(Vec::new());
512 }
513
514 let output = output.unwrap();
515 if !output.status.success() {
516 return Ok(Vec::new());
517 }
518
519 // wsl output is often UTF-16LE, but -l -q might be simpler or just ASCII compatible if not using weird charsets.
520 // However, on Windows, wsl often outputs UTF-16LE.
521 // We can try to detect or use from_utf16 if valid, or just use String::from_utf8_lossy and see.
522 // Actually, `smol::process` on Windows might receive bytes that are UTF-16LE if wsl writes that.
523 // But typically terminal output for wsl is UTF-16.
524 // Let's try to parse as UTF-16LE if it looks like it (BOM or just 00 bytes).
525
526 let stdout = output.stdout;
527 let distros_str = if stdout.len() >= 2 && stdout[1] == 0 {
528 // likely UTF-16LE
529 let u16s: Vec<u16> = stdout
530 .chunks_exact(2)
531 .map(|c| u16::from_le_bytes([c[0], c[1]]))
532 .collect();
533 String::from_utf16_lossy(&u16s)
534 } else {
535 String::from_utf8_lossy(&stdout).to_string()
536 };
537
538 let distros: Vec<String> = distros_str
539 .lines()
540 .map(|line| line.trim().to_string())
541 .filter(|line| !line.is_empty())
542 .collect();
543
544 let tasks = distros.into_iter().map(|distro| {
545 background_executor.spawn(async move {
546 let output = util::command::new_command("wsl")
547 .arg("-d")
548 .arg(&distro)
549 .arg("bash")
550 .arg("-l")
551 .arg("-c")
552 .arg("jupyter kernelspec list --json")
553 .output()
554 .await;
555
556 if let Ok(output) = output {
557 if output.status.success() {
558 let json_str = String::from_utf8_lossy(&output.stdout);
559 // Use local permissive struct instead of strict KernelSpecsResponse from jupyter-protocol
560 if let Ok(specs_response) =
561 serde_json::from_str::<LocalKernelSpecsResponse>(&json_str)
562 {
563 return specs_response
564 .kernelspecs
565 .into_iter()
566 .map(|(name, spec)| {
567 KernelSpecification::WslRemote(WslKernelSpecification {
568 name,
569 kernelspec: jupyter_protocol::JupyterKernelspec {
570 argv: spec.spec.argv,
571 display_name: spec.spec.display_name,
572 language: spec.spec.language,
573 interrupt_mode: spec.spec.interrupt_mode,
574 env: spec.spec.env,
575 metadata: spec.spec.metadata,
576 },
577 distro: distro.clone(),
578 })
579 })
580 .collect::<Vec<_>>();
581 }
582 }
583 }
584
585 Vec::new()
586 })
587 });
588
589 let specs: Vec<_> = futures::future::join_all(tasks)
590 .await
591 .into_iter()
592 .flatten()
593 .collect();
594
595 Ok(specs)
596}