624 lines · 23.4 KB · Rust
RawDownload
use super::{    KernelSession, KernelSpecification, RunningKernel, WslKernelSpecification, start_kernel_tasks,};use anyhow::{Context as _, Result};use futures::{    AsyncBufReadExt as _, StreamExt as _,    channel::mpsc::{self},    io::BufReader,};use gpui::{App, BackgroundExecutor, Entity, EntityId, Task, Window};use jupyter_protocol::{    ExecutionState, JupyterMessage, KernelInfoReply,    connection_info::{ConnectionInfo, Transport},};use project::Fs;use runtimelib::dirs;use smol::net::TcpListener;use std::{    fmt::Debug,    net::{IpAddr, Ipv4Addr, SocketAddr},    path::PathBuf,    sync::Arc,}; use uuid::Uuid; // 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.// There's a race condition between closing the ports and usage by a kernel, but it's inherent to the Jupyter protocol.async fn peek_ports(ip: IpAddr) -> Result<[u16; 5]> {    let mut addr_zeroport: SocketAddr = SocketAddr::new(ip, 0);    addr_zeroport.set_port(0);    let mut ports: [u16; 5] = [0; 5];    for i in 0..5 {        let listener = TcpListener::bind(addr_zeroport).await?;        let addr = listener.local_addr()?;        ports[i] = addr.port();    }    Ok(ports)} pub struct WslRunningKernel {    pub process: util::command::Child,    connection_path: PathBuf,    _process_status_task: Option<Task<()>>,    pub working_directory: PathBuf,    pub request_tx: mpsc::Sender<JupyterMessage>,    pub stdin_tx: mpsc::Sender<JupyterMessage>,    pub execution_state: ExecutionState,    pub kernel_info: Option<KernelInfoReply>,} impl Debug for WslRunningKernel {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        f.debug_struct("WslRunningKernel")            .field("process", &self.process)            .finish()    }} fn quote_posix_shell_arguments(arguments: &[String]) -> Result<String> {    let mut quoted_arguments = Vec::with_capacity(arguments.len());    for argument in arguments {        let quoted = shlex::try_quote(argument).map(|quoted| quoted.into_owned())?;        quoted_arguments.push(quoted);    }    Ok(quoted_arguments.join(" "))} impl WslRunningKernel {    pub fn new<S: KernelSession + 'static>(        kernel_specification: WslKernelSpecification,        entity_id: EntityId,        working_directory: PathBuf,        fs: Arc<dyn Fs>,        session: Entity<S>,        window: &mut Window,        cx: &mut App,    ) -> Task<Result<Box<dyn RunningKernel>>> {        window.spawn(cx, async move |cx| {            // For WSL2, we need to get the WSL VM's IP address to connect to it            // because WSL2 runs in a lightweight VM with its own network namespace.            // The kernel will bind to 127.0.0.1 inside WSL, and we connect to localhost.            // WSL2 localhost forwarding handles the rest.            let bind_ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));             // Use 127.0.0.1 and rely on WSL 2 localhost forwarding.            // This avoids issues where the VM IP is unreachable or binding fails on Windows.            let connect_ip = "127.0.0.1".to_string();             let ports = peek_ports(bind_ip).await?;             let connection_info = ConnectionInfo {                transport: Transport::TCP,                ip: bind_ip.to_string(),                stdin_port: ports[0],                control_port: ports[1],                hb_port: ports[2],                shell_port: ports[3],                iopub_port: ports[4],                signature_scheme: "hmac-sha256".to_string(),                key: uuid::Uuid::new_v4().to_string(),                kernel_name: Some(format!("zed-wsl-{}", kernel_specification.name)),            };             let runtime_dir = dirs::runtime_dir();            fs.create_dir(&runtime_dir)                .await                .with_context(|| format!("Failed to create jupyter runtime dir {runtime_dir:?}"))?;            let connection_path = runtime_dir.join(format!("kernel-zed-wsl-{entity_id}.json"));            let content = serde_json::to_string(&connection_info)?;            fs.atomic_write(connection_path.clone(), content).await?;             // Convert connection_path to WSL path            // yeah we can't assume this is available on WSL.            // running `wsl -d <distro> wslpath -u <windows_path>`            let mut wslpath_cmd = util::command::new_command("wsl");             // On Windows, passing paths with backslashes to wsl.exe can sometimes cause            // escaping issues or be misinterpreted. Converting to forward slashes is safer            // and often accepted by wslpath.            let connection_path_str = connection_path.to_string_lossy().replace('\\', "/");             wslpath_cmd                .arg("-d")                .arg(&kernel_specification.distro)                .arg("wslpath")                .arg("-u")                .arg(&connection_path_str);             let output = wslpath_cmd.output().await?;            if !output.status.success() {                anyhow::bail!("Failed to convert path to WSL path: {:?}", output);            }            let wsl_connection_path = String::from_utf8_lossy(&output.stdout).trim().to_string();             // Construct the kernel command            // The kernel spec argv might have absolute paths valid INSIDE WSL.            // We need to run inside WSL.            // `wsl -d <distro> --exec <argv0> <argv1> ...`            // But we need to replace {connection_file} with wsl_connection_path.             anyhow::ensure!(                !kernel_specification.kernelspec.argv.is_empty(),                "Empty argv in kernelspec {}",                kernel_specification.name            );             let working_directory_str = working_directory.to_string_lossy().replace('\\', "/");             let wsl_working_directory = if working_directory_str.starts_with('/') {                // If path starts with /, assume it is already a WSL path (e.g. /home/user)                Some(working_directory_str)            } else {                let mut wslpath_wd_cmd = util::command::new_command("wsl");                wslpath_wd_cmd                    .arg("-d")                    .arg(&kernel_specification.distro)                    .arg("wslpath")                    .arg("-u")                    .arg(&working_directory_str);                 let wd_output = wslpath_wd_cmd.output().await;                if let Ok(output) = wd_output {                    if output.status.success() {                        Some(String::from_utf8_lossy(&output.stdout).trim().to_string())                    } else {                        None                    }                } else {                    None                }            };             // If we couldn't convert the working directory or it's a temp directory,            // and the kernel spec uses a relative path (like .venv/bin/python),            // we need to handle this better. For now, let's use the converted path            // if available, otherwise we'll rely on WSL's default home directory.             let mut cmd = util::command::new_command("wsl");            cmd.arg("-d").arg(&kernel_specification.distro);             // Set CWD for the host process to a safe location to avoid "Directory name is invalid"            // if the project root is a path not supported by Windows CWD (e.g. UNC path for some tools).            cmd.current_dir(std::env::temp_dir());             if let Some(wd) = wsl_working_directory.as_ref() {                cmd.arg("--cd").arg(wd);            }             // Build the command to run inside WSL            // We use bash -lc to run in a login shell for proper environment setup            let mut kernel_args: Vec<String> = Vec::new();             let resolved_argv: Vec<String> = kernel_specification                .kernelspec                .argv                .iter()                .map(|arg| {                    if arg == "{connection_file}" {                        wsl_connection_path.clone()                    } else {                        arg.clone()                    }                })                .collect();             let executable = resolved_argv.first().map(String::as_str);            let needs_python_resolution = executable.map_or(false, |executable| {                executable == "python" || executable == "python3" || !executable.starts_with('/')            });             let mut env_assignments: Vec<String> = Vec::new();            if let Some(env) = &kernel_specification.kernelspec.env {                env_assignments.reserve(env.len());                for (key, value) in env {                    let assignment = format!("{key}={value}");                    let assignment = shlex::try_quote(&assignment)                        .map(|quoted| quoted.into_owned())?;                    env_assignments.push(assignment);                }                 if !env_assignments.is_empty() {                    kernel_args.push("env".to_string());                    kernel_args.extend(env_assignments.iter().cloned());                }            }             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() {                    String::new()                } else {                    format!("export {}; ", env_assignments.join(" "))                };                 let cd_command = if let Some(wd) = wsl_working_directory.as_ref() {                    let quoted_wd = shlex::try_quote(wd)                        .map(|quoted| quoted.into_owned())?;                    format!("cd {quoted_wd} && ")                } 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                )            } else {                let args_string = quote_posix_shell_arguments(&resolved_argv)?;                 let cd_command = if let Some(wd) = wsl_working_directory.as_ref() {                    let quoted_wd = shlex::try_quote(wd)                        .map(|quoted| quoted.into_owned())?;                    format!("cd {quoted_wd} && ")                } else {                    String::new()                };                 let env_prefix_inline = if !env_assignments.is_empty() {                    format!("env {} ", env_assignments.join(" "))                } else {                    String::new()                };                 format!("{cd_command}exec {env_prefix_inline}{args_string}")            };             cmd.arg("bash")                .arg("-l")                .arg("-c")                .arg(&shell_command);             let mut process = cmd                .stdout(util::command::Stdio::piped())                .stderr(util::command::Stdio::piped())                .stdin(util::command::Stdio::piped())                .kill_on_drop(true)                .spawn()                .context("failed to start the kernel process")?;             let session_id = Uuid::new_v4().to_string();             let mut client_connection_info = connection_info.clone();            client_connection_info.ip = connect_ip.clone();             // Give the kernel a moment to start and bind to ports.            // WSL kernel startup can be slow, I am not sure if this is because of my testing environment            // or inherent to WSL. We can improve this later with better readiness checks.            cx.background_executor()                .timer(std::time::Duration::from_secs(2))                .await;             match process.try_status() {                Ok(Some(status)) => {                    let mut stderr_content = String::new();                    if let Some(mut stderr) = process.stderr.take() {                        use futures::AsyncReadExt;                        let mut buf = Vec::new();                        if stderr.read_to_end(&mut buf).await.is_ok() {                            stderr_content = String::from_utf8_lossy(&buf).to_string();                        }                    }                     let mut stdout_content = String::new();                    if let Some(mut stdout) = process.stdout.take() {                        use futures::AsyncReadExt;                        let mut buf = Vec::new();                        if stdout.read_to_end(&mut buf).await.is_ok() {                            stdout_content = String::from_utf8_lossy(&buf).to_string();                        }                    }                     anyhow::bail!(                        "WSL kernel process exited prematurely with status: {:?}\nstderr: {}\nstdout: {}",                        status,                        stderr_content,                        stdout_content                    );                }                Ok(None) => {}                Err(_) => {}            }             let output_socket = runtimelib::create_client_iopub_connection(                &client_connection_info,                "",                &session_id,            )            .await?;             let peer_identity = runtimelib::peer_identity_for_session(&session_id)?;            let shell_socket = runtimelib::create_client_shell_connection_with_identity(                &client_connection_info,                &session_id,                peer_identity.clone(),            )            .await?;             let control_socket =                runtimelib::create_client_control_connection(&client_connection_info, &session_id)                    .await?;             let stdin_socket = runtimelib::create_client_stdin_connection_with_identity(                &client_connection_info,                &session_id,                peer_identity,            )            .await?;             let (request_tx, stdin_tx) = start_kernel_tasks(                session.clone(),                output_socket,                shell_socket,                control_socket,                stdin_socket,                cx,            );             let stderr = process.stderr.take();            cx.spawn(async move |_cx| {                if let Some(stderr) = stderr {                    let reader = BufReader::new(stderr);                    let mut lines = reader.lines();                    while let Some(Ok(line)) = lines.next().await {                        log::warn!("wsl kernel stderr: {}", line);                    }                }            })            .detach();             let stdout = process.stdout.take();            cx.spawn(async move |_cx| {                if let Some(stdout) = stdout {                    let reader = BufReader::new(stdout);                    let mut lines = reader.lines();                    while let Some(Ok(_line)) = lines.next().await {}                }            })            .detach();             let status = process.status();             let process_status_task = cx.spawn(async move |cx| {                let error_message = match status.await {                    Ok(status) => {                        if status.success() {                            return;                        }                         format!("WSL kernel: kernel process exited with status: {:?}", status)                    }                    Err(err) => {                        format!("WSL kernel: kernel process exited with error: {:?}", err)                    }                };                 session.update(cx, |session, cx| {                    session.kernel_errored(error_message, cx);                     cx.notify();                });            });             anyhow::Ok(Box::new(Self {                process,                request_tx,                stdin_tx,                working_directory,                _process_status_task: Some(process_status_task),                connection_path,                execution_state: ExecutionState::Idle,                kernel_info: None,            }) as Box<dyn RunningKernel>)        })    }} impl RunningKernel for WslRunningKernel {    fn request_tx(&self) -> mpsc::Sender<JupyterMessage> {        self.request_tx.clone()    }     fn stdin_tx(&self) -> mpsc::Sender<JupyterMessage> {        self.stdin_tx.clone()    }     fn working_directory(&self) -> &PathBuf {        &self.working_directory    }     fn execution_state(&self) -> &ExecutionState {        &self.execution_state    }     fn set_execution_state(&mut self, state: ExecutionState) {        self.execution_state = state;    }     fn kernel_info(&self) -> Option<&KernelInfoReply> {        self.kernel_info.as_ref()    }     fn set_kernel_info(&mut self, info: KernelInfoReply) {        self.kernel_info = Some(info);    }     fn force_shutdown(&mut self, _window: &mut Window, _cx: &mut App) -> Task<anyhow::Result<()>> {        self._process_status_task.take();        self.request_tx.close_channel();        self.process.kill().ok();        Task::ready(Ok(()))    }     fn kill(&mut self) {        self._process_status_task.take();        self.request_tx.close_channel();        self.process.kill().ok();    }} impl Drop for WslRunningKernel {    fn drop(&mut self) {        std::fs::remove_file(&self.connection_path).ok();        self.request_tx.close_channel();        self.process.kill().ok();    }} #[derive(serde::Deserialize)]struct LocalKernelSpecsResponse {    kernelspecs: std::collections::HashMap<String, LocalKernelSpec>,}

Lines 1–500 of 624.