repl: Add initial notebook execution + KernelSession abstraction (#43553)

Coming from discussion on #25936 and
https://github.com/zed-industries/zed/pull/19756 .
This PR introduces the KernelSession abstraction and adds basic notebook
cell execution inside NotebookEditor.

The following provides a base starter for the development on Notebooks.

What this PR includes today:

Release Notes:

- KernelSession trait extracted. Both REPL and NotebookEditor now share
the same routing mechanism.
- NotebookEditor can launch kernels, execute code cells, and display
outputs.
- Basic cell operations: run current cell, run all, move up/down, add
code/markdown blocks. Keybindings follow Jupyter defaults.


Next Steps:

- [x] Editing support for markdown and code blocks.
- Buffer integration instead of temporary cell text.
- [x] Proper focus behavior when executing or adding cells.
- Kernel control UI.

A little far fetched steps:

- Vim Support
- Cell Handling Improvement and other convenient features and design
from other editors
- Ability to have better parsing for AI Support. 


I have attached a video of showcasing some of the features:


https://github.com/user-attachments/assets/37e6f3e5-2022-45f0-a73d-2dd01ebc2932

assets/keymaps/default-linux.json

@@ -1334,6 +1334,21 @@
       "alt-right": "dev::EditPredictionContextGoForward",
     },
   },
+  {
+    "context": "NotebookEditor > Editor",
+    "bindings": {
+      "enter": "editor::Newline",
+      "shift-enter": "notebook::Run",
+      "ctrl-enter": "notebook::Run",
+      "ctrl-shift-enter": "notebook::RunAll",
+      "alt-up": "notebook::MoveCellUp",
+      "alt-down": "notebook::MoveCellDown",
+      "ctrl-m": "notebook::AddCodeBlock",
+      "ctrl-shift-m": "notebook::AddMarkdownBlock",
+      "ctrl-shift-r": "notebook::RestartKernel",
+      "ctrl-c": "notebook::InterruptKernel"
+    }
+  },
   {
     "context": "GitBranchSelector || (GitBranchSelector > Picker > Editor)",
     "use_key_equivalents": true,

assets/keymaps/default-macos.json

@@ -1479,4 +1479,19 @@
       "shift-backspace": "project_dropdown::RemoveSelectedFolder",
     },
   },
+  {
+    "context": "NotebookEditor > Editor",
+    "bindings": {
+      "enter": "editor::Newline",
+      "shift-enter": "notebook::Run",
+      "cmd-enter": "notebook::Run",
+      "cmd-shift-enter": "notebook::RunAll",
+      "alt-up": "notebook::MoveCellUp",
+      "alt-down": "notebook::MoveCellDown",
+      "cmd-m": "notebook::AddCodeBlock",
+      "cmd-shift-m": "notebook::AddMarkdownBlock",
+      "cmd-shift-r": "notebook::RestartKernel",
+      "cmd-c": "notebook::InterruptKernel",
+    },
+  },
 ]

assets/keymaps/default-windows.json

@@ -1397,4 +1397,19 @@
       "shift-backspace": "project_dropdown::RemoveSelectedFolder",
     },
   },
+  {
+    "context": "NotebookEditor > Editor",
+    "bindings": {
+      "enter": "editor::Newline",
+      "shift-enter": "notebook::Run",
+      "ctrl-enter": "notebook::Run",
+      "ctrl-shift-enter": "notebook::RunAll",
+      "alt-up": "notebook::MoveCellUp",
+      "alt-down": "notebook::MoveCellDown",
+      "ctrl-m": "notebook::AddCodeBlock",
+      "ctrl-shift-m": "notebook::AddMarkdownBlock",
+      "ctrl-shift-r": "notebook::RestartKernel",
+      "ctrl-c": "notebook::InterruptKernel"
+    }
+  },
 ]

crates/repl/src/kernels/mod.rs

@@ -15,11 +15,17 @@ use project::{Project, ProjectPath, Toolchains, WorktreeId};
 pub use remote_kernels::*;
 
 use anyhow::Result;
+use gpui::Context;
 use jupyter_protocol::JupyterKernelspec;
 use runtimelib::{ExecutionState, JupyterMessage, KernelInfoReply};
 use ui::{Icon, IconName, SharedString};
 use util::rel_path::RelPath;
 
+pub trait KernelSession: Sized {
+    fn route(&mut self, message: &JupyterMessage, window: &mut Window, cx: &mut Context<Self>);
+    fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>);
+}
+
 pub type JupyterMessageChannel = stream::SelectAll<Receiver<JupyterMessage>>;
 
 #[derive(Debug, Clone, PartialEq, Eq)]

crates/repl/src/kernels/native_kernel.rs

@@ -22,9 +22,7 @@ use std::{
 };
 use uuid::Uuid;
 
-use crate::Session;
-
-use super::RunningKernel;
+use super::{KernelSession, RunningKernel};
 
 #[derive(Debug, Clone)]
 pub struct LocalKernelSpecification {
@@ -105,13 +103,13 @@ impl Debug for NativeRunningKernel {
 }
 
 impl NativeRunningKernel {
-    pub fn new(
+    pub fn new<S: KernelSession + 'static>(
         kernel_specification: LocalKernelSpecification,
         entity_id: EntityId,
         working_directory: PathBuf,
         fs: Arc<dyn Fs>,
         // todo: convert to weak view
-        session: Entity<Session>,
+        session: Entity<S>,
         window: &mut Window,
         cx: &mut App,
     ) -> Task<Result<Box<dyn RunningKernel>>> {

crates/repl/src/kernels/remote_kernels.rs

@@ -9,9 +9,7 @@ use async_tungstenite::tungstenite::{client::IntoClientRequest, http::HeaderValu
 use futures::StreamExt;
 use smol::io::AsyncReadExt as _;
 
-use crate::Session;
-
-use super::RunningKernel;
+use super::{KernelSession, RunningKernel};
 use anyhow::Result;
 use jupyter_websocket_client::{
     JupyterWebSocket, JupyterWebSocketReader, JupyterWebSocketWriter, KernelLaunchRequest,
@@ -127,10 +125,10 @@ pub struct RemoteRunningKernel {
 }
 
 impl RemoteRunningKernel {
-    pub fn new(
+    pub fn new<S: KernelSession + 'static>(
         kernelspec: RemoteKernelSpecification,
         working_directory: std::path::PathBuf,
-        session: Entity<Session>,
+        session: Entity<S>,
         window: &mut Window,
         cx: &mut App,
     ) -> Task<Result<Box<dyn RunningKernel>>> {

crates/repl/src/outputs.rs

@@ -127,6 +127,43 @@ pub enum Output {
     ClearOutputWaitMarker,
 }
 
+impl Output {
+    pub fn to_nbformat(&self, cx: &App) -> Option<nbformat::v4::Output> {
+        match self {
+            Output::Stream { content } => {
+                let text = content.read(cx).full_text();
+                Some(nbformat::v4::Output::Stream {
+                    name: "stdout".to_string(),
+                    text: nbformat::v4::MultilineString(text),
+                })
+            }
+            Output::Plain { content, .. } => {
+                let text = content.read(cx).full_text();
+                let mut data = jupyter_protocol::media::Media::default();
+                data.content.push(jupyter_protocol::MediaType::Plain(text));
+                Some(nbformat::v4::Output::DisplayData(
+                    nbformat::v4::DisplayData {
+                        data,
+                        metadata: serde_json::Map::new(),
+                    },
+                ))
+            }
+            Output::ErrorOutput(error_view) => {
+                let traceback_text = error_view.traceback.read(cx).full_text();
+                let traceback_lines: Vec<String> =
+                    traceback_text.lines().map(|s| s.to_string()).collect();
+                Some(nbformat::v4::Output::Error(nbformat::v4::ErrorOutput {
+                    ename: error_view.ename.clone(),
+                    evalue: error_view.evalue.clone(),
+                    traceback: traceback_lines,
+                }))
+            }
+            Output::Message(_) | Output::ClearOutputWaitMarker => None,
+            Output::Image { .. } | Output::Table { .. } | Output::Markdown { .. } => None,
+        }
+    }
+}
+
 impl Output {
     fn render_output_controls<V: OutputContent + 'static>(
         v: Entity<V>,

crates/repl/src/session.rs

@@ -3,7 +3,7 @@ use crate::kernels::RemoteRunningKernel;
 use crate::setup_editor_session_actions;
 use crate::{
     KernelStatus,
-    kernels::{Kernel, KernelSpecification, NativeRunningKernel},
+    kernels::{Kernel, KernelSession, KernelSpecification, NativeRunningKernel},
     outputs::{
         ExecutionStatus, ExecutionView, ExecutionViewFinishedEmpty, ExecutionViewFinishedSmall,
     },
@@ -648,51 +648,6 @@ impl Session {
         }
     }
 
-    pub fn route(&mut self, message: &JupyterMessage, window: &mut Window, cx: &mut Context<Self>) {
-        let parent_message_id = match message.parent_header.as_ref() {
-            Some(header) => &header.msg_id,
-            None => return,
-        };
-
-        match &message.content {
-            JupyterMessageContent::Status(status) => {
-                self.kernel.set_execution_state(&status.execution_state);
-
-                telemetry::event!(
-                    "Kernel Status Changed",
-                    kernel_language = self.kernel_specification.language(),
-                    kernel_status = KernelStatus::from(&self.kernel).to_string(),
-                    repl_session_id = cx.entity_id().to_string(),
-                );
-
-                cx.notify();
-            }
-            JupyterMessageContent::KernelInfoReply(reply) => {
-                self.kernel.set_kernel_info(reply);
-                cx.notify();
-            }
-            JupyterMessageContent::UpdateDisplayData(update) => {
-                let display_id = if let Some(display_id) = update.transient.display_id.clone() {
-                    display_id
-                } else {
-                    return;
-                };
-
-                self.blocks.iter_mut().for_each(|(_, block)| {
-                    block.execution_view.update(cx, |execution_view, cx| {
-                        execution_view.update_display_data(&update.data, &display_id, window, cx);
-                    });
-                });
-                return;
-            }
-            _ => {}
-        }
-
-        if let Some(block) = self.blocks.get_mut(parent_message_id) {
-            block.handle_message(message, window, cx);
-        }
-    }
-
     pub fn interrupt(&mut self, cx: &mut Context<Self>) {
         match &mut self.kernel {
             Kernel::RunningKernel(_kernel) => {
@@ -861,3 +816,54 @@ impl Render for Session {
             .buttons(interrupt_button)
     }
 }
+
+impl KernelSession for Session {
+    fn route(&mut self, message: &JupyterMessage, window: &mut Window, cx: &mut Context<Self>) {
+        let parent_message_id = match message.parent_header.as_ref() {
+            Some(header) => &header.msg_id,
+            None => return,
+        };
+
+        match &message.content {
+            JupyterMessageContent::Status(status) => {
+                self.kernel.set_execution_state(&status.execution_state);
+
+                telemetry::event!(
+                    "Kernel Status Changed",
+                    kernel_language = self.kernel_specification.language(),
+                    kernel_status = KernelStatus::from(&self.kernel).to_string(),
+                    repl_session_id = cx.entity_id().to_string(),
+                );
+
+                cx.notify();
+            }
+            JupyterMessageContent::KernelInfoReply(reply) => {
+                self.kernel.set_kernel_info(reply);
+                cx.notify();
+            }
+            JupyterMessageContent::UpdateDisplayData(update) => {
+                let display_id = if let Some(display_id) = update.transient.display_id.clone() {
+                    display_id
+                } else {
+                    return;
+                };
+
+                self.blocks.iter_mut().for_each(|(_, block)| {
+                    block.execution_view.update(cx, |execution_view, cx| {
+                        execution_view.update_display_data(&update.data, &display_id, window, cx);
+                    });
+                });
+                return;
+            }
+            _ => {}
+        }
+
+        if let Some(block) = self.blocks.get_mut(parent_message_id) {
+            block.handle_message(message, window, cx);
+        }
+    }
+
+    fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>) {
+        self.kernel_errored(error_message, cx);
+    }
+}