repl: Unify kernel searching in remote and wsl (#53049)

MostlyK created

### Context:

- Having a unified way of searching would allow for better debugging as
we move forward here. Right now we have remote/headless specific
searching and it's getting messy. This should allow for a more intuitive
function graph in the head for debugging environment related issues in
remote repl. The implementation mirrors python.rs approach.

Self-Review Checklist:

- [x] I've reviewed my own diff for quality, security, and reliability
- [x] Unsafe blocks (if any) have justifying comments
- [x] The content is consistent with the [UI/UX
checklist](https://github.com/zed-industries/zed/blob/main/CONTRIBUTING.md#uiux-checklist)
- [x] Tests cover the new/changed behavior
- [x] Performance impact has been considered and is acceptable

Closes #50892

Release Notes:

- N/A

Change summary

crates/remote_server/src/headless_project.rs |  29 +++++
crates/repl/src/kernels/mod.rs               | 120 +++++++++++++++------
crates/repl/src/kernels/wsl_kernel.rs        |  35 -----
3 files changed, 120 insertions(+), 64 deletions(-)

Detailed changes

crates/remote_server/src/headless_project.rs 🔗

@@ -1001,6 +1001,15 @@ impl HeadlessProject {
                 "failed to spawn kernel process (command: {})",
                 envelope.payload.command
             ))?
+        } else if let Some(venv_python) = working_directory
+            .as_ref()
+            .and_then(|wd| find_venv_python(wd))
+        {
+            let path_str = venv_python.to_string_lossy().to_string();
+            spawn_kernel(&path_str, &[]).context(format!(
+                "failed to spawn kernel process (venv: {})",
+                path_str
+            ))?
         } else {
             spawn_kernel("python3", &[])
                 .or_else(|_| spawn_kernel("python", &[]))
@@ -1325,3 +1334,23 @@ fn prompt_to_proto(
         ),
     }
 }
+
+fn find_venv_python(working_directory: &str) -> Option<std::path::PathBuf> {
+    let wd = std::path::Path::new(working_directory);
+    for dir_name in &[".venv", "venv", ".env", "env"] {
+        let venv_dir = wd.join(dir_name);
+        let has_pyvenv_cfg = venv_dir.join("pyvenv.cfg").is_file();
+        let has_activate = venv_dir.join("bin").join("activate").is_file();
+        if has_pyvenv_cfg || has_activate {
+            let python = venv_dir.join("bin").join("python");
+            if python.is_file() {
+                return Some(python);
+            }
+            let python3 = venv_dir.join("bin").join("python3");
+            if python3.is_file() {
+                return Some(python3);
+            }
+        }
+    }
+    None
+}

crates/repl/src/kernels/mod.rs 🔗

@@ -31,6 +31,62 @@ use runtimelib::{
 use ui::{Icon, IconName, SharedString};
 use util::rel_path::RelPath;
 
+pub(crate) const VENV_DIR_NAMES: &[&str] = &[".venv", "venv", ".env", "env"];
+
+// Build a POSIX shell script that attempts to find and exec the best Python binary to run with the given arguments.
+pub(crate) fn build_python_exec_shell_script(
+    python_args: &str,
+    cd_command: &str,
+    env_command: &str,
+) -> String {
+    let venv_dirs = VENV_DIR_NAMES.join(" ");
+    format!(
+        "set -e; \
+         {cd_command}\
+         {env_command}\
+         for venv_dir in {venv_dirs}; do \
+           if [ -f \"$venv_dir/pyvenv.cfg\" ] || [ -f \"$venv_dir/bin/activate\" ]; then \
+             if [ -x \"$venv_dir/bin/python\" ]; then \
+               exec \"$venv_dir/bin/python\" {python_args}; \
+             elif [ -x \"$venv_dir/bin/python3\" ]; then \
+               exec \"$venv_dir/bin/python3\" {python_args}; \
+             fi; \
+           fi; \
+         done; \
+         if command -v python3 >/dev/null 2>&1; then \
+           exec python3 {python_args}; \
+         elif command -v python >/dev/null 2>&1; then \
+           exec python {python_args}; \
+         else \
+           echo 'Error: Python not found in virtual environment or PATH' >&2; \
+           exit 127; \
+         fi"
+    )
+}
+
+/// Build a POSIX shell script that outputs the best Python binary.
+#[cfg(target_os = "windows")]
+pub(crate) fn build_python_discovery_shell_script() -> String {
+    let venv_dirs = VENV_DIR_NAMES.join(" ");
+    format!(
+        "for venv_dir in {venv_dirs}; do \
+           if [ -f \"$venv_dir/pyvenv.cfg\" ] || [ -f \"$venv_dir/bin/activate\" ]; then \
+             if [ -x \"$venv_dir/bin/python\" ]; then \
+               echo \"$venv_dir/bin/python\"; exit 0; \
+             elif [ -x \"$venv_dir/bin/python3\" ]; then \
+               echo \"$venv_dir/bin/python3\"; exit 0; \
+             fi; \
+           fi; \
+         done; \
+         if command -v python3 >/dev/null 2>&1; then \
+           echo python3; exit 0; \
+         elif command -v python >/dev/null 2>&1; then \
+           echo python; exit 0; \
+         fi; \
+         exit 1"
+    )
+}
+
 pub fn start_kernel_tasks<S: KernelSession + 'static>(
     session: Entity<S>,
     iopub_socket: ClientIoPubConnection,
@@ -542,49 +598,47 @@ pub fn python_env_kernel_specifications(
                     };
 
                 if let (Some(distro), Some(internal_path)) = (distro, internal_path) {
-                    let python_path = format!("{}/.venv/bin/python", internal_path);
-                    let check = util::command::new_command("wsl")
-                        .args(&["-d", distro, "test", "-f", &python_path])
+                    let discovery_script = build_python_discovery_shell_script();
+                    let script = format!(
+                        "cd {} && {}",
+                        shlex::try_quote(&internal_path)
+                            .unwrap_or(std::borrow::Cow::Borrowed(&internal_path)),
+                        discovery_script
+                    );
+                    let output = util::command::new_command("wsl")
+                        .arg("-d")
+                        .arg(distro)
+                        .arg("bash")
+                        .arg("-l")
+                        .arg("-c")
+                        .arg(&script)
                         .output()
                         .await;
 
-                    if check.is_ok() && check.unwrap().status.success() {
-                        let default_kernelspec = JupyterKernelspec {
-                            argv: vec![
-                                python_path.clone(),
-                                "-m".to_string(),
-                                "ipykernel_launcher".to_string(),
-                                "-f".to_string(),
-                                "{connection_file}".to_string(),
-                            ],
-                            display_name: format!("WSL: {} (.venv)", distro),
-                            language: "python".to_string(),
-                            interrupt_mode: None,
-                            metadata: None,
-                            env: None,
-                        };
-
-                        kernel_specs.push(KernelSpecification::WslRemote(WslKernelSpecification {
-                            name: format!("WSL: {} (.venv)", distro),
-                            kernelspec: default_kernelspec,
-                            distro: distro.to_string(),
-                        }));
-                    } else {
-                        let check_system = util::command::new_command("wsl")
-                            .args(&["-d", distro, "command", "-v", "python3"])
-                            .output()
-                            .await;
+                    if let Ok(output) = output {
+                        if output.status.success() {
+                            let python_cmd =
+                                String::from_utf8_lossy(&output.stdout).trim().to_string();
+                            let (python_path, display_suffix) = if python_cmd.contains('/') {
+                                let venv_name = python_cmd.split('/').next().unwrap_or("venv");
+                                (
+                                    format!("{}/{}", internal_path, python_cmd),
+                                    format!("({})", venv_name),
+                                )
+                            } else {
+                                (python_cmd, "(System)".to_string())
+                            };
 
-                        if check_system.is_ok() && check_system.unwrap().status.success() {
+                            let display_name = format!("WSL: {} {}", distro, display_suffix);
                             let default_kernelspec = JupyterKernelspec {
                                 argv: vec![
-                                    "python3".to_string(),
+                                    python_path,
                                     "-m".to_string(),
                                     "ipykernel_launcher".to_string(),
                                     "-f".to_string(),
                                     "{connection_file}".to_string(),
                                 ],
-                                display_name: format!("WSL: {} (System)", distro),
+                                display_name: display_name.clone(),
                                 language: "python".to_string(),
                                 interrupt_mode: None,
                                 metadata: None,
@@ -593,7 +647,7 @@ pub fn python_env_kernel_specifications(
 
                             kernel_specs.push(KernelSpecification::WslRemote(
                                 WslKernelSpecification {
-                                    name: format!("WSL: {} (System)", distro),
+                                    name: display_name,
                                     kernelspec: default_kernelspec,
                                     distro: distro.to_string(),
                                 },

crates/repl/src/kernels/wsl_kernel.rs 🔗

@@ -1,5 +1,6 @@
 use super::{
-    KernelSession, KernelSpecification, RunningKernel, WslKernelSpecification, start_kernel_tasks,
+    KernelSession, KernelSpecification, RunningKernel, WslKernelSpecification,
+    build_python_exec_shell_script, start_kernel_tasks,
 };
 use anyhow::{Context as _, Result};
 use futures::{
@@ -228,8 +229,6 @@ impl WslRunningKernel {
             kernel_args.extend(resolved_argv.iter().cloned());
 
             let shell_command = if needs_python_resolution {
-                // 1. Check for .venv/bin/python or .venv/bin/python3 in working directory
-                // 2. Fall back to system python3 or python
                 let rest_args: Vec<String> = resolved_argv.iter().skip(1).cloned().collect();
                 let arg_string = quote_posix_shell_arguments(&rest_args)?;
                 let set_env_command = if env_assignments.is_empty() {
@@ -245,34 +244,8 @@ impl WslRunningKernel {
                 } else {
                     String::new()
                 };
-                // TODO: find a better way to debug missing python issues in WSL
-
-                format!(
-                    "set -e; \
-                     {} \
-                     {} \
-                     echo \"Working directory: $(pwd)\" >&2; \
-                     if [ -x .venv/bin/python ]; then \
-                       echo \"Found .venv/bin/python\" >&2; \
-                       exec .venv/bin/python {}; \
-                     elif [ -x .venv/bin/python3 ]; then \
-                       echo \"Found .venv/bin/python3\" >&2; \
-                       exec .venv/bin/python3 {}; \
-                     elif command -v python3 >/dev/null 2>&1; then \
-                       echo \"Found system python3\" >&2; \
-                       exec python3 {}; \
-                     elif command -v python >/dev/null 2>&1; then \
-                       echo \"Found system python\" >&2; \
-                       exec python {}; \
-                     else \
-                       echo 'Error: Python not found in .venv or PATH' >&2; \
-                       echo 'Contents of current directory:' >&2; \
-                       ls -la >&2; \
-                       echo 'PATH:' \"$PATH\" >&2; \
-                       exit 127; \
-                     fi",
-                    cd_command, set_env_command, arg_string, arg_string, arg_string, arg_string
-                )
+
+                build_python_exec_shell_script(&arg_string, &cd_command, &set_env_command)
             } else {
                 let args_string = quote_posix_shell_arguments(&resolved_argv)?;