Implement RunningKernel trait for native and remote kernels (#20934)

This PR introduces a unified interface for both native and remote
kernels through the `RunningKernel` trait. When either the native kernel
or the remote kernels are started, they return a `Box<dyn
RunningKernel>` to make it easier to work with the session. As a bonus
of this refactor, I've dropped some of the mpsc channels to instead opt
for passing messages directly to `session.route(message)`. There was a
lot of simplification of `Session` by moving responsibilities to
`NativeRunningKernel`.

No release notes yet until this is finalized.

* [x] Detect remote kernelspecs from configured remote servers
* [x] Launch kernel on demand

For now, this allows you to set env vars `JUPYTER_SERVER` and
`JUPYTER_TOKEN` to access a remote server. `JUPYTER_SERVER` should be a
base path like `http://localhost:8888` or
`https://notebooks.gesis.org/binder/jupyter/user/rubydata-binder-w6igpy4l/`

Release Notes:

- N/A

crates/repl/src/kernels/mod.rs

@@ -6,7 +6,7 @@ use futures::{
     future::Shared,
     stream,
 };
-use gpui::{AppContext, Model, Task};
+use gpui::{AppContext, Model, Task, WindowContext};
 use language::LanguageName;
 pub use native_kernel::*;
 
@@ -16,7 +16,7 @@ pub use remote_kernels::*;
 
 use anyhow::Result;
 use runtimelib::{ExecutionState, JupyterKernelspec, JupyterMessage, KernelInfoReply};
-use ui::SharedString;
+use ui::{Icon, IconName, SharedString};
 
 pub type JupyterMessageChannel = stream::SelectAll<Receiver<JupyterMessage>>;
 
@@ -59,6 +59,19 @@ impl KernelSpecification {
             Self::Remote(spec) => spec.kernelspec.language.clone(),
         })
     }
+
+    pub fn icon(&self, cx: &AppContext) -> Icon {
+        let lang_name = match self {
+            Self::Jupyter(spec) => spec.kernelspec.language.clone(),
+            Self::PythonEnv(spec) => spec.kernelspec.language.clone(),
+            Self::Remote(spec) => spec.kernelspec.language.clone(),
+        };
+
+        file_icons::FileIcons::get(cx)
+            .get_type_icon(&lang_name.to_lowercase())
+            .map(Icon::from_path)
+            .unwrap_or(Icon::new(IconName::ReplNeutral))
+    }
 }
 
 pub fn python_env_kernel_specifications(
@@ -134,7 +147,7 @@ pub trait RunningKernel: Send + Debug {
     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<()>;
+    fn force_shutdown(&mut self, cx: &mut WindowContext) -> Task<anyhow::Result<()>>;
 }
 
 #[derive(Debug, Clone)]

crates/repl/src/kernels/native_kernel.rs

@@ -1,10 +1,11 @@
 use anyhow::{Context as _, Result};
 use futures::{
     channel::mpsc::{self},
+    io::BufReader,
     stream::{SelectAll, StreamExt},
-    SinkExt as _,
+    AsyncBufReadExt as _, SinkExt as _,
 };
-use gpui::{AppContext, EntityId, Task};
+use gpui::{EntityId, Task, View, WindowContext};
 use jupyter_protocol::{JupyterMessage, JupyterMessageContent, KernelInfoReply};
 use project::Fs;
 use runtimelib::{dirs, ConnectionInfo, ExecutionState, JupyterKernelspec};
@@ -18,7 +19,9 @@ use std::{
 };
 use uuid::Uuid;
 
-use super::{JupyterMessageChannel, RunningKernel};
+use crate::Session;
+
+use super::RunningKernel;
 
 #[derive(Debug, Clone)]
 pub struct LocalKernelSpecification {
@@ -83,10 +86,10 @@ async fn peek_ports(ip: IpAddr) -> Result<[u16; 5]> {
 pub struct NativeRunningKernel {
     pub process: smol::process::Child,
     _shell_task: Task<Result<()>>,
-    _iopub_task: Task<Result<()>>,
     _control_task: Task<Result<()>>,
     _routing_task: Task<Result<()>>,
     connection_path: PathBuf,
+    _process_status_task: Option<Task<()>>,
     pub working_directory: PathBuf,
     pub request_tx: mpsc::Sender<JupyterMessage>,
     pub execution_state: ExecutionState,
@@ -107,8 +110,10 @@ impl NativeRunningKernel {
         entity_id: EntityId,
         working_directory: PathBuf,
         fs: Arc<dyn Fs>,
-        cx: &mut AppContext,
-    ) -> Task<Result<(Self, JupyterMessageChannel)>> {
+        // todo: convert to weak view
+        session: View<Session>,
+        cx: &mut WindowContext,
+    ) -> Task<Result<Box<dyn RunningKernel>>> {
         cx.spawn(|cx| async move {
             let ip = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
             let ports = peek_ports(ip).await?;
@@ -136,7 +141,7 @@ impl NativeRunningKernel {
 
             let mut cmd = kernel_specification.command(&connection_path)?;
 
-            let process = cmd
+            let mut process = cmd
                 .current_dir(&working_directory)
                 .stdout(std::process::Stdio::piped())
                 .stderr(std::process::Stdio::piped())
@@ -155,8 +160,6 @@ impl NativeRunningKernel {
             let mut control_socket =
                 runtimelib::create_client_control_connection(&connection_info, &session_id).await?;
 
-            let (mut iopub, iosub) = futures::channel::mpsc::channel(100);
-
             let (request_tx, mut request_rx) =
                 futures::channel::mpsc::channel::<JupyterMessage>(100);
 
@@ -164,18 +167,41 @@ impl NativeRunningKernel {
             let (mut shell_reply_tx, shell_reply_rx) = futures::channel::mpsc::channel(100);
 
             let mut messages_rx = SelectAll::new();
-            messages_rx.push(iosub);
             messages_rx.push(control_reply_rx);
             messages_rx.push(shell_reply_rx);
 
-            let iopub_task = cx.background_executor().spawn({
-                async move {
-                    while let Ok(message) = iopub_socket.read().await {
-                        iopub.send(message).await?;
+            cx.spawn({
+                let session = session.clone();
+
+                |mut cx| async move {
+                    while let Some(message) = messages_rx.next().await {
+                        session
+                            .update(&mut cx, |session, cx| {
+                                session.route(&message, cx);
+                            })
+                            .ok();
                     }
                     anyhow::Ok(())
                 }
-            });
+            })
+            .detach();
+
+            // iopub task
+            cx.spawn({
+                let session = session.clone();
+
+                |mut cx| async move {
+                    while let Ok(message) = iopub_socket.read().await {
+                        session
+                            .update(&mut cx, |session, cx| {
+                                session.route(&message, cx);
+                            })
+                            .ok();
+                    }
+                    anyhow::Ok(())
+                }
+            })
+            .detach();
 
             let (mut control_request_tx, mut control_request_rx) =
                 futures::channel::mpsc::channel(100);
@@ -221,21 +247,74 @@ impl NativeRunningKernel {
                 }
             });
 
-            anyhow::Ok((
-                Self {
-                    process,
-                    request_tx,
-                    working_directory,
-                    _shell_task: shell_task,
-                    _iopub_task: iopub_task,
-                    _control_task: control_task,
-                    _routing_task: routing_task,
-                    connection_path,
-                    execution_state: ExecutionState::Idle,
-                    kernel_info: None,
-                },
-                messages_rx,
-            ))
+            let stderr = process.stderr.take();
+
+            cx.spawn(|mut _cx| async move {
+                if stderr.is_none() {
+                    return;
+                }
+                let reader = BufReader::new(stderr.unwrap());
+                let mut lines = reader.lines();
+                while let Some(Ok(line)) = lines.next().await {
+                    log::error!("kernel: {}", line);
+                }
+            })
+            .detach();
+
+            let stdout = process.stdout.take();
+
+            cx.spawn(|mut _cx| async move {
+                if stdout.is_none() {
+                    return;
+                }
+                let reader = BufReader::new(stdout.unwrap());
+                let mut lines = reader.lines();
+                while let Some(Ok(line)) = lines.next().await {
+                    log::info!("kernel: {}", line);
+                }
+            })
+            .detach();
+
+            let status = process.status();
+
+            let process_status_task = cx.spawn(|mut cx| async move {
+                let error_message = match status.await {
+                    Ok(status) => {
+                        if status.success() {
+                            log::info!("kernel process exited successfully");
+                            return;
+                        }
+
+                        format!("kernel process exited with status: {:?}", status)
+                    }
+                    Err(err) => {
+                        format!("kernel process exited with error: {:?}", err)
+                    }
+                };
+
+                log::error!("{}", error_message);
+
+                session
+                    .update(&mut cx, |session, cx| {
+                        session.kernel_errored(error_message, cx);
+
+                        cx.notify();
+                    })
+                    .ok();
+            });
+
+            anyhow::Ok(Box::new(Self {
+                process,
+                request_tx,
+                working_directory,
+                _process_status_task: Some(process_status_task),
+                _shell_task: shell_task,
+                _control_task: control_task,
+                _routing_task: routing_task,
+                connection_path,
+                execution_state: ExecutionState::Idle,
+                kernel_info: None,
+            }) as Box<dyn RunningKernel>)
         })
     }
 }
@@ -265,14 +344,17 @@ impl RunningKernel for NativeRunningKernel {
         self.kernel_info = Some(info);
     }
 
-    fn force_shutdown(&mut self) -> anyhow::Result<()> {
-        match self.process.kill() {
+    fn force_shutdown(&mut self, _cx: &mut WindowContext) -> Task<anyhow::Result<()>> {
+        self._process_status_task.take();
+        self.request_tx.close_channel();
+
+        Task::ready(match self.process.kill() {
             Ok(_) => Ok(()),
             Err(error) => Err(anyhow::anyhow!(
                 "Failed to kill the kernel process: {}",
                 error
             )),
-        }
+        })
     }
 }
 

crates/repl/src/kernels/remote_kernels.rs

@@ -1,12 +1,21 @@
-use futures::{channel::mpsc, StreamExt as _};
-use gpui::AppContext;
+use futures::{channel::mpsc, SinkExt as _};
+use gpui::{Task, View, WindowContext};
+use http_client::{AsyncBody, HttpClient, Request};
 use jupyter_protocol::{ExecutionState, JupyterMessage, KernelInfoReply};
-// todo(kyle): figure out if this needs to be different
 use runtimelib::JupyterKernelspec;
 
+use futures::StreamExt;
+use smol::io::AsyncReadExt as _;
+
+use crate::Session;
+
 use super::RunningKernel;
-use jupyter_websocket_client::RemoteServer;
-use std::fmt::Debug;
+use anyhow::Result;
+use jupyter_websocket_client::{
+    JupyterWebSocketReader, JupyterWebSocketWriter, KernelLaunchRequest, KernelSpecsResponse,
+    RemoteServer,
+};
+use std::{fmt::Debug, sync::Arc};
 
 #[derive(Debug, Clone)]
 pub struct RemoteKernelSpecification {
@@ -16,6 +25,101 @@ pub struct RemoteKernelSpecification {
     pub kernelspec: JupyterKernelspec,
 }
 
+pub async fn launch_remote_kernel(
+    remote_server: &RemoteServer,
+    http_client: Arc<dyn HttpClient>,
+    kernel_name: &str,
+    _path: &str,
+) -> Result<String> {
+    //
+    let kernel_launch_request = KernelLaunchRequest {
+        name: kernel_name.to_string(),
+        // todo: add path to runtimelib
+        // path,
+    };
+
+    let kernel_launch_request = serde_json::to_string(&kernel_launch_request)?;
+
+    let request = Request::builder()
+        .method("POST")
+        .uri(&remote_server.api_url("/kernels"))
+        .header("Authorization", format!("token {}", remote_server.token))
+        .body(AsyncBody::from(kernel_launch_request))?;
+
+    let response = http_client.send(request).await?;
+
+    if !response.status().is_success() {
+        let mut body = String::new();
+        response.into_body().read_to_string(&mut body).await?;
+        return Err(anyhow::anyhow!("Failed to launch kernel: {}", body));
+    }
+
+    let mut body = String::new();
+    response.into_body().read_to_string(&mut body).await?;
+
+    let response: jupyter_websocket_client::Kernel = serde_json::from_str(&body)?;
+
+    Ok(response.id)
+}
+
+pub async fn list_remote_kernelspecs(
+    remote_server: RemoteServer,
+    http_client: Arc<dyn HttpClient>,
+) -> Result<Vec<RemoteKernelSpecification>> {
+    let url = remote_server.api_url("/kernelspecs");
+
+    let request = Request::builder()
+        .method("GET")
+        .uri(&url)
+        .header("Authorization", format!("token {}", remote_server.token))
+        .body(AsyncBody::default())?;
+
+    let response = http_client.send(request).await?;
+
+    if response.status().is_success() {
+        let mut body = response.into_body();
+
+        let mut body_bytes = Vec::new();
+        body.read_to_end(&mut body_bytes).await?;
+
+        let kernel_specs: KernelSpecsResponse = serde_json::from_slice(&body_bytes)?;
+
+        let remote_kernelspecs = kernel_specs
+            .kernelspecs
+            .into_iter()
+            .map(|(name, spec)| RemoteKernelSpecification {
+                name: name.clone(),
+                url: remote_server.base_url.clone(),
+                token: remote_server.token.clone(),
+                // todo: line up the jupyter kernelspec from runtimelib with
+                //       the kernelspec pulled from the API
+                //
+                //        There are _small_ differences, so we may just want a impl `From`
+                kernelspec: JupyterKernelspec {
+                    argv: spec.spec.argv,
+                    display_name: spec.spec.display_name,
+                    language: spec.spec.language,
+                    // todo: fix up mismatch in types here
+                    metadata: None,
+                    interrupt_mode: None,
+                    env: None,
+                },
+            })
+            .collect::<Vec<RemoteKernelSpecification>>();
+
+        if remote_kernelspecs.is_empty() {
+            Err(anyhow::anyhow!("No kernel specs found"))
+        } else {
+            Ok(remote_kernelspecs.clone())
+        }
+    } else {
+        Err(anyhow::anyhow!(
+            "Failed to fetch kernel specs: {}",
+            response.status()
+        ))
+    }
+}
+
 impl PartialEq for RemoteKernelSpecification {
     fn eq(&self, other: &Self) -> bool {
         self.name == other.name && self.url == other.url
@@ -26,55 +130,91 @@ impl Eq for RemoteKernelSpecification {}
 
 pub struct RemoteRunningKernel {
     remote_server: RemoteServer,
+    _receiving_task: Task<Result<()>>,
+    _routing_task: Task<Result<()>>,
+    http_client: Arc<dyn HttpClient>,
     pub working_directory: std::path::PathBuf,
     pub request_tx: mpsc::Sender<JupyterMessage>,
     pub execution_state: ExecutionState,
     pub kernel_info: Option<KernelInfoReply>,
+    pub kernel_id: String,
 }
 
 impl RemoteRunningKernel {
-    pub async fn new(
+    pub fn new(
         kernelspec: RemoteKernelSpecification,
         working_directory: std::path::PathBuf,
-        request_tx: mpsc::Sender<JupyterMessage>,
-        _cx: &mut AppContext,
-    ) -> anyhow::Result<(
-        Self,
-        (), // Stream<Item=JupyterMessage>
-    )> {
+        session: View<Session>,
+        cx: &mut WindowContext,
+    ) -> Task<Result<Box<dyn RunningKernel>>> {
         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 http_client = cx.http_client();
 
-        let kernel_socket = remote_server.connect_to_kernel(kernel_id).await?;
+        cx.spawn(|cx| async move {
+            let kernel_id = launch_remote_kernel(
+                &remote_server,
+                http_client.clone(),
+                &kernelspec.name,
+                working_directory.to_str().unwrap_or_default(),
+            )
+            .await?;
 
-        let (mut _w, mut _r) = kernel_socket.split();
+            let kernel_socket = remote_server.connect_to_kernel(&kernel_id).await?;
 
-        let (_messages_tx, _messages_rx) = mpsc::channel::<JupyterMessage>(100);
+            let (mut w, mut r): (JupyterWebSocketWriter, JupyterWebSocketReader) =
+                kernel_socket.split();
 
-        // 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();
+            let (request_tx, mut request_rx) =
+                futures::channel::mpsc::channel::<JupyterMessage>(100);
 
-        anyhow::Ok((
-            Self {
+            let routing_task = cx.background_executor().spawn({
+                async move {
+                    while let Some(message) = request_rx.next().await {
+                        w.send(message).await.ok();
+                    }
+                    Ok(())
+                }
+            });
+
+            let receiving_task = cx.spawn({
+                let session = session.clone();
+
+                |mut cx| async move {
+                    while let Some(message) = r.next().await {
+                        match message {
+                            Ok(message) => {
+                                session
+                                    .update(&mut cx, |session, cx| {
+                                        session.route(&message, cx);
+                                    })
+                                    .ok();
+                            }
+                            Err(e) => {
+                                log::error!("Error receiving message: {:?}", e);
+                            }
+                        }
+                    }
+                    Ok(())
+                }
+            });
+
+            anyhow::Ok(Box::new(Self {
+                _routing_task: routing_task,
+                _receiving_task: receiving_task,
                 remote_server,
                 working_directory,
                 request_tx,
+                // todo(kyle): pull this from the kernel API to start with
                 execution_state: ExecutionState::Idle,
                 kernel_info: None,
-            },
-            (),
-        ))
+                kernel_id,
+                http_client: http_client.clone(),
+            }) as Box<dyn RunningKernel>)
+        })
     }
 }
 
@@ -116,7 +256,30 @@ impl RunningKernel for RemoteRunningKernel {
         self.kernel_info = Some(info);
     }
 
-    fn force_shutdown(&mut self) -> anyhow::Result<()> {
-        unimplemented!("force_shutdown")
+    fn force_shutdown(&mut self, cx: &mut WindowContext) -> Task<anyhow::Result<()>> {
+        let url = self
+            .remote_server
+            .api_url(&format!("/kernels/{}", self.kernel_id));
+        let token = self.remote_server.token.clone();
+        let http_client = self.http_client.clone();
+
+        cx.spawn(|_| async move {
+            let request = Request::builder()
+                .method("DELETE")
+                .uri(&url)
+                .header("Authorization", format!("token {}", token))
+                .body(AsyncBody::default())?;
+
+            let response = http_client.send(request).await?;
+
+            if response.status().is_success() {
+                Ok(())
+            } else {
+                Err(anyhow::anyhow!(
+                    "Failed to shutdown kernel: {}",
+                    response.status()
+                ))
+            }
+        })
     }
 }