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

### 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
This commit is contained in:
MostlyK 2026-04-21 12:54:09 +05:30 committed by GitHub
parent eb6e7d7b64
commit 040b03b34d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 119 additions and 63 deletions

View file

@ -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
}

View file

@ -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,
};
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())
};
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 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(),
},

View file

@ -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)?;