Create RunningKernel trait to allow for native and remote jupyter kernels (#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,
+        }
+    }
+}

crates/repl/src/kernels/native_kernel.rs

@@ -1,69 +1,24 @@
 use anyhow::{Context as _, Result};
 use futures::{
-    channel::mpsc::{self, Receiver},
-    future::Shared,
-    stream::{self, SelectAll, StreamExt},
+    channel::mpsc::{self},
+    stream::{SelectAll, StreamExt},
     SinkExt as _,
 };
-use gpui::{AppContext, EntityId, Model, Task};
-use language::LanguageName;
-use project::{Fs, Project, WorktreeId};
-use runtimelib::{
-    dirs, ConnectionInfo, ExecutionState, JupyterKernelspec, JupyterMessage, JupyterMessageContent,
-    KernelInfoReply,
-};
+use gpui::{AppContext, EntityId, Task};
+use jupyter_protocol::{JupyterMessage, JupyterMessageContent, KernelInfoReply};
+use project::Fs;
+use runtimelib::{dirs, ConnectionInfo, ExecutionState, JupyterKernelspec};
 use smol::{net::TcpListener, process::Command};
 use std::{
     env,
     fmt::Debug,
-    future::Future,
     net::{IpAddr, Ipv4Addr, SocketAddr},
     path::PathBuf,
     sync::Arc,
 };
-use ui::SharedString;
 use uuid::Uuid;
 
-#[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(),
-        })
-    }
-}
+use super::{JupyterMessageChannel, RunningKernel};
 
 #[derive(Debug, Clone)]
 pub struct LocalKernelSpecification {
@@ -80,22 +35,6 @@ impl PartialEq for LocalKernelSpecification {
 
 impl Eq for LocalKernelSpecification {}
 
-#[derive(Debug, Clone)]
-pub struct RemoteKernelSpecification {
-    pub name: String,
-    pub url: String,
-    pub token: String,
-    pub kernelspec: JupyterKernelspec,
-}
-
-impl PartialEq for RemoteKernelSpecification {
-    fn eq(&self, other: &Self) -> bool {
-        self.name == other.name && self.url == other.url
-    }
-}
-
-impl Eq for RemoteKernelSpecification {}
-
 impl LocalKernelSpecification {
     #[must_use]
     fn command(&self, connection_path: &PathBuf) -> Result<Command> {
@@ -147,95 +86,7 @@ async fn peek_ports(ip: IpAddr) -> Result<[u16; 5]> {
     Ok(ports)
 }
 
-#[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(),
-        }
-    }
-}
-
-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,
-        }
-    }
-}
-
-#[derive(Debug)]
-pub enum Kernel {
-    RunningKernel(RunningKernel),
-    StartingKernel(Shared<Task<()>>),
-    ErroredLaunch(String),
-    ShuttingDown,
-    Shutdown,
-    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.execution_state = status.clone();
-        }
-    }
-
-    pub fn set_kernel_info(&mut self, kernel_info: &KernelInfoReply) {
-        if let Kernel::RunningKernel(running_kernel) = self {
-            running_kernel.kernel_info = Some(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,
-        }
-    }
-}
-
-pub struct RunningKernel {
+pub struct NativeRunningKernel {
     pub process: smol::process::Child,
     _shell_task: Task<Result<()>>,
     _iopub_task: Task<Result<()>>,
@@ -248,9 +99,7 @@ pub struct RunningKernel {
     pub kernel_info: Option<KernelInfoReply>,
 }
 
-type JupyterMessageChannel = stream::SelectAll<Receiver<JupyterMessage>>;
-
-impl Debug for RunningKernel {
+impl Debug for NativeRunningKernel {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("RunningKernel")
             .field("process", &self.process)
@@ -258,25 +107,14 @@ impl Debug for RunningKernel {
     }
 }
 
-impl RunningKernel {
+impl NativeRunningKernel {
     pub fn new(
-        kernel_specification: KernelSpecification,
+        kernel_specification: LocalKernelSpecification,
         entity_id: EntityId,
         working_directory: PathBuf,
         fs: Arc<dyn Fs>,
         cx: &mut AppContext,
     ) -> Task<Result<(Self, JupyterMessageChannel)>> {
-        let kernel_specification = match kernel_specification {
-            KernelSpecification::Jupyter(spec) => spec,
-            KernelSpecification::PythonEnv(spec) => spec,
-            KernelSpecification::Remote(_spec) => {
-                // todo!(): Implement remote kernel specification
-                return Task::ready(Err(anyhow::anyhow!(
-                    "Running remote kernels is not supported"
-                )));
-            }
-        };
-
         cx.spawn(|cx| async move {
             let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
             let ports = peek_ports(ip).await?;
@@ -315,15 +153,13 @@ impl RunningKernel {
 
             let session_id = Uuid::new_v4().to_string();
 
-            let mut iopub_socket = connection_info
-                .create_client_iopub_connection("", &session_id)
-                .await?;
-            let mut shell_socket = connection_info
-                .create_client_shell_connection(&session_id)
-                .await?;
-            let mut control_socket = connection_info
-                .create_client_control_connection(&session_id)
-                .await?;
+            let mut iopub_socket =
+                runtimelib::create_client_iopub_connection(&connection_info, "", &session_id)
+                    .await?;
+            let mut shell_socket =
+                runtimelib::create_client_shell_connection(&connection_info, &session_id).await?;
+            let mut control_socket =
+                runtimelib::create_client_control_connection(&connection_info, &session_id).await?;
 
             let (mut iopub, iosub) = futures::channel::mpsc::channel(100);
 
@@ -410,7 +246,43 @@ impl RunningKernel {
     }
 }
 
-impl Drop for RunningKernel {
+impl RunningKernel for NativeRunningKernel {
+    fn request_tx(&self) -> mpsc::Sender<JupyterMessage> {
+        self.request_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) -> anyhow::Result<()> {
+        match self.process.kill() {
+            Ok(_) => Ok(()),
+            Err(error) => Err(anyhow::anyhow!(
+                "Failed to kill the kernel process: {}",
+                error
+            )),
+        }
+    }
+}
+
+impl Drop for NativeRunningKernel {
     fn drop(&mut self) {
         std::fs::remove_file(&self.connection_path).ok();
         self.request_tx.close_channel();
@@ -467,72 +339,6 @@ async fn read_kernels_dir(path: PathBuf, fs: &dyn Fs) -> Result<Vec<LocalKernelS
     Ok(valid_kernelspecs)
 }
 
-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 async fn local_kernel_specifications(fs: Arc<dyn Fs>) -> Result<Vec<LocalKernelSpecification>> {
     let mut data_dirs = dirs::data_dirs();
 

crates/repl/src/kernels/remote_kernels.rs

@@ -0,0 +1,122 @@
+use futures::{channel::mpsc, StreamExt as _};
+use gpui::AppContext;
+use jupyter_protocol::{ExecutionState, JupyterMessage, KernelInfoReply};
+// todo(kyle): figure out if this needs to be different
+use runtimelib::JupyterKernelspec;
+
+use super::RunningKernel;
+use jupyter_websocket_client::RemoteServer;
+use std::fmt::Debug;
+
+#[derive(Debug, Clone)]
+pub struct RemoteKernelSpecification {
+    pub name: String,
+    pub url: String,
+    pub token: String,
+    pub kernelspec: JupyterKernelspec,
+}
+
+impl PartialEq for RemoteKernelSpecification {
+    fn eq(&self, other: &Self) -> bool {
+        self.name == other.name && self.url == other.url
+    }
+}
+
+impl Eq for RemoteKernelSpecification {}
+
+pub struct RemoteRunningKernel {
+    remote_server: RemoteServer,
+    pub working_directory: std::path::PathBuf,
+    pub request_tx: mpsc::Sender<JupyterMessage>,
+    pub execution_state: ExecutionState,
+    pub kernel_info: Option<KernelInfoReply>,
+}
+
+impl RemoteRunningKernel {
+    pub async fn new(
+        kernelspec: RemoteKernelSpecification,
+        working_directory: std::path::PathBuf,
+        request_tx: mpsc::Sender<JupyterMessage>,
+        _cx: &mut AppContext,
+    ) -> anyhow::Result<(
+        Self,
+        (), // Stream<Item=JupyterMessage>
+    )> {
+        let remote_server = RemoteServer {
+            base_url: kernelspec.url,
+            token: kernelspec.token,
+        };
+
+        // todo: launch a kernel to get a kernel ID
+        let kernel_id = "not-implemented";
+
+        let kernel_socket = remote_server.connect_to_kernel(kernel_id).await?;
+
+        let (mut _w, mut _r) = kernel_socket.split();
+
+        let (_messages_tx, _messages_rx) = mpsc::channel::<JupyterMessage>(100);
+
+        // let routing_task = cx.background_executor().spawn({
+        //     async move {
+        //         while let Some(message) = request_rx.next().await {
+        //             w.send(message).await;
+        //         }
+        //     }
+        // });
+        // let messages_rx = r.into();
+
+        anyhow::Ok((
+            Self {
+                remote_server,
+                working_directory,
+                request_tx,
+                execution_state: ExecutionState::Idle,
+                kernel_info: None,
+            },
+            (),
+        ))
+    }
+}
+
+impl Debug for RemoteRunningKernel {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("RemoteRunningKernel")
+            // custom debug that keeps tokens out of logs
+            .field("remote_server url", &self.remote_server.base_url)
+            .field("working_directory", &self.working_directory)
+            .field("request_tx", &self.request_tx)
+            .field("execution_state", &self.execution_state)
+            .field("kernel_info", &self.kernel_info)
+            .finish()
+    }
+}
+
+impl RunningKernel for RemoteRunningKernel {
+    fn request_tx(&self) -> futures::channel::mpsc::Sender<runtimelib::JupyterMessage> {
+        self.request_tx.clone()
+    }
+
+    fn working_directory(&self) -> &std::path::PathBuf {
+        &self.working_directory
+    }
+
+    fn execution_state(&self) -> &runtimelib::ExecutionState {
+        &self.execution_state
+    }
+
+    fn set_execution_state(&mut self, state: runtimelib::ExecutionState) {
+        self.execution_state = state;
+    }
+
+    fn kernel_info(&self) -> Option<&runtimelib::KernelInfoReply> {
+        self.kernel_info.as_ref()
+    }
+
+    fn set_kernel_info(&mut self, info: runtimelib::KernelInfoReply) {
+        self.kernel_info = Some(info);
+    }
+
+    fn force_shutdown(&mut self) -> anyhow::Result<()> {
+        unimplemented!("force_shutdown")
+    }
+}

crates/repl/src/repl.rs

@@ -1,6 +1,6 @@
 pub mod components;
 mod jupyter_settings;
-mod kernels;
+pub mod kernels;
 pub mod notebook;
 mod outputs;
 mod repl_editor;

crates/repl/src/session.rs

@@ -1,7 +1,7 @@
 use crate::components::KernelListItem;
 use crate::setup_editor_session_actions;
 use crate::{
-    kernels::{Kernel, KernelSpecification, RunningKernel},
+    kernels::{Kernel, KernelSpecification, NativeRunningKernel},
     outputs::{ExecutionStatus, ExecutionView},
     KernelStatus,
 };
@@ -246,13 +246,19 @@ impl Session {
             cx.entity_id().to_string(),
         );
 
-        let kernel = RunningKernel::new(
-            self.kernel_specification.clone(),
-            entity_id,
-            working_directory,
-            self.fs.clone(),
-            cx,
-        );
+        let kernel = match self.kernel_specification.clone() {
+            KernelSpecification::Jupyter(kernel_specification)
+            | KernelSpecification::PythonEnv(kernel_specification) => NativeRunningKernel::new(
+                kernel_specification,
+                entity_id,
+                working_directory,
+                self.fs.clone(),
+                cx,
+            ),
+            KernelSpecification::Remote(_remote_kernel_specification) => {
+                unimplemented!()
+            }
+        };
 
         let pending_kernel = cx
             .spawn(|this, mut cx| async move {
@@ -291,7 +297,7 @@ impl Session {
                             .detach();
 
                             let status = kernel.process.status();
-                            session.kernel(Kernel::RunningKernel(kernel), cx);
+                            session.kernel(Kernel::RunningKernel(Box::new(kernel)), cx);
 
                             let process_status_task = cx.spawn(|session, mut cx| async move {
                                 let error_message = match status.await {
@@ -416,7 +422,7 @@ impl Session {
 
     fn send(&mut self, message: JupyterMessage, _cx: &mut ViewContext<Self>) -> anyhow::Result<()> {
         if let Kernel::RunningKernel(kernel) = &mut self.kernel {
-            kernel.request_tx.try_send(message).ok();
+            kernel.request_tx().try_send(message).ok();
         }
 
         anyhow::Ok(())
@@ -631,7 +637,7 @@ impl Session {
 
         match kernel {
             Kernel::RunningKernel(mut kernel) => {
-                let mut request_tx = kernel.request_tx.clone();
+                let mut request_tx = kernel.request_tx().clone();
 
                 cx.spawn(|this, mut cx| async move {
                     let message: JupyterMessage = ShutdownRequest { restart: false }.into();
@@ -646,7 +652,7 @@ impl Session {
                     })
                     .ok();
 
-                    kernel.process.kill().ok();
+                    kernel.force_shutdown().ok();
 
                     this.update(&mut cx, |session, cx| {
                         session.clear_outputs(cx);
@@ -674,7 +680,7 @@ impl Session {
                 // Do nothing if already restarting
             }
             Kernel::RunningKernel(mut kernel) => {
-                let mut request_tx = kernel.request_tx.clone();
+                let mut request_tx = kernel.request_tx().clone();
 
                 cx.spawn(|this, mut cx| async move {
                     // Send shutdown request with restart flag
@@ -692,7 +698,7 @@ impl Session {
                     cx.background_executor().timer(Duration::from_secs(1)).await;
 
                     // Force kill the kernel if it hasn't shut down
-                    kernel.process.kill().ok();
+                    kernel.force_shutdown().ok();
 
                     // Start a new kernel
                     this.update(&mut cx, |session, cx| {
@@ -727,7 +733,7 @@ impl Render for Session {
         let (status_text, interrupt_button) = match &self.kernel {
             Kernel::RunningKernel(kernel) => (
                 kernel
-                    .kernel_info
+                    .kernel_info()
                     .as_ref()
                     .map(|info| info.language_info.name.clone()),
                 Some(
@@ -747,7 +753,7 @@ impl Render for Session {
 
         KernelListItem::new(self.kernel_specification.clone())
             .status_color(match &self.kernel {
-                Kernel::RunningKernel(kernel) => match kernel.execution_state {
+                Kernel::RunningKernel(kernel) => match kernel.execution_state() {
                     ExecutionState::Idle => Color::Success,
                     ExecutionState::Busy => Color::Modified,
                 },