Create RunningKernel trait to allow for native and remote jupyter k…

…ernels (#20842)

Starts setting up a `RunningKernel` trait to make the remote kernel
implementation easy to get started with. No release notes until this is
all hooked up.

Release Notes:

- N/A

Cargo.toml

@@ -368,12 +368,14 @@ indexmap = { version = "1.6.2", features = ["serde"] }
 indoc = "2"
 itertools = "0.13.0"
 jsonwebtoken = "9.3"
+jupyter-protocol = { version = "0.2.0" }
+jupyter-websocket-client = { version = "0.4.1" }
 libc = "0.2"
 linkify = "0.10.0"
 log = { version = "0.4.16", features = ["kv_unstable_serde", "serde"] }
 markup5ever_rcdom = "0.3.0"
 nanoid = "0.4"
-nbformat = "0.5.0"
+nbformat = "0.6.0"
 nix = "0.29"
 num-format = "0.4.4"
 once_cell = "1.19.0"
@@ -407,7 +409,7 @@ reqwest = { git = "https://github.com/zed-industries/reqwest.git", rev = "fd110f
     "stream",
 ] }
 rsa = "0.9.6"
-runtimelib = { version = "0.19.0", default-features = false, features = [
+runtimelib = { version = "0.21.0", default-features = false, features = [
     "async-dispatcher-runtime",
 ] }
 rustc-demangle = "0.1.23"

crates/quick_action_bar/src/repl_menu.rs

@@ -402,7 +402,7 @@ fn session_state(session: View<Session>, cx: &WindowContext) -> ReplMenuState {
             status: session.kernel.status(),
             ..fill_fields()
         },
-        Kernel::RunningKernel(kernel) => match &kernel.execution_state {
+        Kernel::RunningKernel(kernel) => match &kernel.execution_state() {
             ExecutionState::Idle => ReplMenuState {
                 tooltip: format!("Run code on {} ({})", kernel_name, kernel_language).into(),
                 indicator: Some(Indicator::dot().color(Color::Success)),

crates/repl/Cargo.toml

@@ -25,6 +25,8 @@ feature_flags.workspace = true
 futures.workspace = true
 gpui.workspace = true
 image.workspace = true
+jupyter-websocket-client.workspace = true
+jupyter-protocol.workspace = true
 language.workspace = true
 log.workspace = true
 markdown_preview.workspace = true

crates/repl/src/kernels/mod.rs

@@ -0,0 +1,227 @@
+mod native_kernel;
+use std::{fmt::Debug, future::Future, path::PathBuf};
+
+use futures::{
+    channel::mpsc::{self, Receiver},
+    future::Shared,
+    stream,
+};
+use gpui::{AppContext, Model, Task};
+use language::LanguageName;
+pub use native_kernel::*;
+
+mod remote_kernels;
+use project::{Project, WorktreeId};
+pub use remote_kernels::*;
+
+use anyhow::Result;
+use runtimelib::{ExecutionState, JupyterKernelspec, JupyterMessage, KernelInfoReply};
+use smol::process::Command;
+use ui::SharedString;
+
+pub type JupyterMessageChannel = stream::SelectAll<Receiver<JupyterMessage>>;
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum KernelSpecification {
+    Remote(RemoteKernelSpecification),
+    Jupyter(LocalKernelSpecification),
+    PythonEnv(LocalKernelSpecification),
+}
+
+impl KernelSpecification {
+    pub fn name(&self) -> SharedString {
+        match self {
+            Self::Jupyter(spec) => spec.name.clone().into(),
+            Self::PythonEnv(spec) => spec.name.clone().into(),
+            Self::Remote(spec) => spec.name.clone().into(),
+        }
+    }
+
+    pub fn type_name(&self) -> SharedString {
+        match self {
+            Self::Jupyter(_) => "Jupyter".into(),
+            Self::PythonEnv(_) => "Python Environment".into(),
+            Self::Remote(_) => "Remote".into(),
+        }
+    }
+
+    pub fn path(&self) -> SharedString {
+        SharedString::from(match self {
+            Self::Jupyter(spec) => spec.path.to_string_lossy().to_string(),
+            Self::PythonEnv(spec) => spec.path.to_string_lossy().to_string(),
+            Self::Remote(spec) => spec.url.to_string(),
+        })
+    }
+
+    pub fn language(&self) -> SharedString {
+        SharedString::from(match self {
+            Self::Jupyter(spec) => spec.kernelspec.language.clone(),
+            Self::PythonEnv(spec) => spec.kernelspec.language.clone(),
+            Self::Remote(spec) => spec.kernelspec.language.clone(),
+        })
+    }
+}
+
+pub fn python_env_kernel_specifications(
+    project: &Model<Project>,
+    worktree_id: WorktreeId,
+    cx: &mut AppContext,
+) -> impl Future<Output = Result<Vec<KernelSpecification>>> {
+    let python_language = LanguageName::new("Python");
+    let toolchains = project
+        .read(cx)
+        .available_toolchains(worktree_id, python_language, cx);
+    let background_executor = cx.background_executor().clone();
+
+    async move {
+        let toolchains = if let Some(toolchains) = toolchains.await {
+            toolchains
+        } else {
+            return Ok(Vec::new());
+        };
+
+        let kernelspecs = toolchains.toolchains.into_iter().map(|toolchain| {
+            background_executor.spawn(async move {
+                let python_path = toolchain.path.to_string();
+
+                // Check if ipykernel is installed
+                let ipykernel_check = Command::new(&python_path)
+                    .args(&["-c", "import ipykernel"])
+                    .output()
+                    .await;
+
+                if ipykernel_check.is_ok() && ipykernel_check.unwrap().status.success() {
+                    // Create a default kernelspec for this environment
+                    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: toolchain.name.to_string(),
+                        language: "python".to_string(),
+                        interrupt_mode: None,
+                        metadata: None,
+                        env: None,
+                    };
+
+                    Some(KernelSpecification::PythonEnv(LocalKernelSpecification {
+                        name: toolchain.name.to_string(),
+                        path: PathBuf::from(&python_path),
+                        kernelspec: default_kernelspec,
+                    }))
+                } else {
+                    None
+                }
+            })
+        });
+
+        let kernel_specs = futures::future::join_all(kernelspecs)
+            .await
+            .into_iter()
+            .flatten()
+            .collect();
+
+        anyhow::Ok(kernel_specs)
+    }
+}
+
+pub trait RunningKernel: Send + Debug {
+    fn request_tx(&self) -> mpsc::Sender<JupyterMessage>;
+    fn working_directory(&self) -> &PathBuf;
+    fn execution_state(&self) -> &ExecutionState;
+    fn set_execution_state(&mut self, state: ExecutionState);
+    fn kernel_info(&self) -> Option<&KernelInfoReply>;
+    fn set_kernel_info(&mut self, info: KernelInfoReply);
+    fn force_shutdown(&mut self) -> anyhow::Result<()>;
+}
+
+#[derive(Debug, Clone)]
+pub enum KernelStatus {
+    Idle,
+    Busy,
+    Starting,
+    Error,
+    ShuttingDown,
+    Shutdown,
+    Restarting,
+}
+
+impl KernelStatus {
+    pub fn is_connected(&self) -> bool {
+        match self {
+            KernelStatus::Idle | KernelStatus::Busy => true,
+            _ => false,
+        }
+    }
+}
+
+impl ToString for KernelStatus {
+    fn to_string(&self) -> String {
+        match self {
+            KernelStatus::Idle => "Idle".to_string(),
+            KernelStatus::Busy => "Busy".to_string(),
+            KernelStatus::Starting => "Starting".to_string(),
+            KernelStatus::Error => "Error".to_string(),
+            KernelStatus::ShuttingDown => "Shutting Down".to_string(),
+            KernelStatus::Shutdown => "Shutdown".to_string(),
+            KernelStatus::Restarting => "Restarting".to_string(),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub enum Kernel {
+    RunningKernel(Box<dyn RunningKernel>),
+    StartingKernel(Shared<Task<()>>),
+    ErroredLaunch(String),
+    ShuttingDown,
+    Shutdown,
+    Restarting,
+}
+
+impl From<&Kernel> for KernelStatus {
+    fn from(kernel: &Kernel) -> Self {
+        match kernel {
+            Kernel::RunningKernel(kernel) => match kernel.execution_state() {
+                ExecutionState::Idle => KernelStatus::Idle,
+                ExecutionState::Busy => KernelStatus::Busy,
+            },
+            Kernel::StartingKernel(_) => KernelStatus::Starting,
+            Kernel::ErroredLaunch(_) => KernelStatus::Error,
+            Kernel::ShuttingDown => KernelStatus::ShuttingDown,
+            Kernel::Shutdown => KernelStatus::Shutdown,
+            Kernel::Restarting => KernelStatus::Restarting,
+        }
+    }
+}
+
+impl Kernel {
+    pub fn status(&self) -> KernelStatus {
+        self.into()
+    }
+
+    pub fn set_execution_state(&mut self, status: &ExecutionState) {
+        if let Kernel::RunningKernel(running_kernel) = self {
+            running_kernel.set_execution_state(status.clone());
+        }
+    }
+
+    pub fn set_kernel_info(&mut self, kernel_info: &KernelInfoReply) {
+        if let Kernel::RunningKernel(running_kernel) = self {
+            running_kernel.set_kernel_info(kernel_info.clone());
+        }
+    }
+
+    pub fn is_shutting_down(&self) -> bool {
+        match self {
+            Kernel::Restarting | Kernel::ShuttingDown => true,
+            Kernel::RunningKernel(_)
+            | Kernel::StartingKernel(_)
+            | Kernel::ErroredLaunch(_)
+            | Kernel::Shutdown => false,
+        }
+    }
+}