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

MostlyK created

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

Change summary

assets/keymaps/default-linux.json         |  15 
assets/keymaps/default-macos.json         |  15 
assets/keymaps/default-windows.json       |  15 
crates/repl/src/kernels/mod.rs            |   6 
crates/repl/src/kernels/native_kernel.rs  |   8 
crates/repl/src/kernels/remote_kernels.rs |   8 
crates/repl/src/notebook/cell.rs          | 923 ++++++++++++++++++++----
crates/repl/src/notebook/notebook_ui.rs   | 937 ++++++++++++++++++++++--
crates/repl/src/outputs.rs                |  37 
crates/repl/src/session.rs                |  98 +-
10 files changed, 1,755 insertions(+), 307 deletions(-)

Detailed changes

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/notebook/cell.rs 🔗

@@ -1,17 +1,20 @@
 #![allow(unused, dead_code)]
 use std::sync::Arc;
+use std::time::{Duration, Instant};
 
 use editor::{Editor, EditorMode, MultiBuffer};
 use futures::future::Shared;
 use gpui::{
-    App, Entity, Hsla, RetainAllImageCache, Task, TextStyleRefinement, image_cache, prelude::*,
+    App, Entity, EventEmitter, Focusable, Hsla, InteractiveElement, RetainAllImageCache,
+    StatefulInteractiveElement, Task, TextStyleRefinement, image_cache, prelude::*,
 };
 use language::{Buffer, Language, LanguageRegistry};
 use markdown_preview::{markdown_parser::parse_markdown, markdown_renderer::render_markdown_block};
 use nbformat::v4::{CellId, CellMetadata, CellType};
+use runtimelib::{JupyterMessage, JupyterMessageContent};
 use settings::Settings as _;
 use theme::ThemeSettings;
-use ui::{IconButtonShape, prelude::*};
+use ui::{CommonAnimationExt, IconButtonShape, prelude::*};
 use util::ResultExt;
 
 use crate::{
@@ -35,6 +38,16 @@ pub enum CellControlType {
     ExpandCell,
 }
 
+pub enum CellEvent {
+    Run(CellId),
+    FocusedIn(CellId),
+}
+
+pub enum MarkdownCellEvent {
+    FinishedEditing,
+    Run(CellId),
+}
+
 impl CellControlType {
     fn icon_name(&self) -> IconName {
         match self {
@@ -113,6 +126,38 @@ fn convert_outputs(
 }
 
 impl Cell {
+    pub fn id(&self, cx: &App) -> CellId {
+        match self {
+            Cell::Code(code_cell) => code_cell.read(cx).id().clone(),
+            Cell::Markdown(markdown_cell) => markdown_cell.read(cx).id().clone(),
+            Cell::Raw(raw_cell) => raw_cell.read(cx).id().clone(),
+        }
+    }
+
+    pub fn current_source(&self, cx: &App) -> String {
+        match self {
+            Cell::Code(code_cell) => code_cell.read(cx).current_source(cx),
+            Cell::Markdown(markdown_cell) => markdown_cell.read(cx).current_source(cx),
+            Cell::Raw(raw_cell) => raw_cell.read(cx).source.clone(),
+        }
+    }
+
+    pub fn to_nbformat_cell(&self, cx: &App) -> nbformat::v4::Cell {
+        match self {
+            Cell::Code(code_cell) => code_cell.read(cx).to_nbformat_cell(cx),
+            Cell::Markdown(markdown_cell) => markdown_cell.read(cx).to_nbformat_cell(cx),
+            Cell::Raw(raw_cell) => raw_cell.read(cx).to_nbformat_cell(),
+        }
+    }
+
+    pub fn is_dirty(&self, cx: &App) -> bool {
+        match self {
+            Cell::Code(code_cell) => code_cell.read(cx).is_dirty(cx),
+            Cell::Markdown(markdown_cell) => markdown_cell.read(cx).is_dirty(cx),
+            Cell::Raw(_) => false,
+        }
+    }
+
     pub fn load(
         cell: &nbformat::v4::Cell,
         languages: &Arc<LanguageRegistry>,
@@ -130,35 +175,14 @@ impl Cell {
                 let source = source.join("");
 
                 let entity = cx.new(|cx| {
-                    let markdown_parsing_task = {
-                        let languages = languages.clone();
-                        let source = source.clone();
-
-                        cx.spawn_in(window, async move |this, cx| {
-                            let parsed_markdown = cx
-                                .background_spawn(async move {
-                                    parse_markdown(&source, None, Some(languages)).await
-                                })
-                                .await;
-
-                            this.update(cx, |cell: &mut MarkdownCell, _| {
-                                cell.parsed_markdown = Some(parsed_markdown);
-                            })
-                            .log_err();
-                        })
-                    };
-
-                    MarkdownCell {
-                        markdown_parsing_task,
-                        image_cache: RetainAllImageCache::new(cx),
-                        languages: languages.clone(),
-                        id: id.clone(),
-                        metadata: metadata.clone(),
-                        source: source.clone(),
-                        parsed_markdown: None,
-                        selected: false,
-                        cell_position: None,
-                    }
+                    MarkdownCell::new(
+                        id.clone(),
+                        metadata.clone(),
+                        source,
+                        languages.clone(),
+                        window,
+                        cx,
+                    )
                 });
 
                 Cell::Markdown(entity)
@@ -169,63 +193,23 @@ impl Cell {
                 execution_count,
                 source,
                 outputs,
-            } => Cell::Code(cx.new(|cx| {
+            } => {
                 let text = source.join("");
-
-                let buffer = cx.new(|cx| Buffer::local(text.clone(), cx));
-                let multi_buffer = cx.new(|cx| MultiBuffer::singleton(buffer.clone(), cx));
-
-                let editor_view = cx.new(|cx| {
-                    let mut editor = Editor::new(
-                        EditorMode::AutoHeight {
-                            min_lines: 1,
-                            max_lines: Some(1024),
-                        },
-                        multi_buffer,
-                        None,
+                let outputs = convert_outputs(outputs, window, cx);
+
+                Cell::Code(cx.new(|cx| {
+                    CodeCell::load(
+                        id.clone(),
+                        metadata.clone(),
+                        *execution_count,
+                        text,
+                        outputs,
+                        notebook_language,
                         window,
                         cx,
-                    );
-
-                    let theme = ThemeSettings::get_global(cx);
-
-                    let refinement = TextStyleRefinement {
-                        font_family: Some(theme.buffer_font.family.clone()),
-                        font_size: Some(theme.buffer_font_size(cx).into()),
-                        color: Some(cx.theme().colors().editor_foreground),
-                        background_color: Some(gpui::transparent_black()),
-                        ..Default::default()
-                    };
-
-                    editor.set_text(text, window, cx);
-                    editor.set_show_gutter(false, cx);
-                    editor.set_text_style_refinement(refinement);
-
-                    // editor.set_read_only(true);
-                    editor
-                });
-
-                let buffer = buffer.clone();
-                let language_task = cx.spawn_in(window, async move |this, cx| {
-                    let language = notebook_language.await;
-
-                    buffer.update(cx, |buffer, cx| {
-                        buffer.set_language(language.clone(), cx);
-                    });
-                });
-
-                CodeCell {
-                    id: id.clone(),
-                    metadata: metadata.clone(),
-                    execution_count: *execution_count,
-                    source: source.join(""),
-                    editor: editor_view,
-                    outputs: convert_outputs(outputs, window, cx),
-                    selected: false,
-                    language_task,
-                    cell_position: None,
-                }
-            })),
+                    )
+                }))
+            }
             nbformat::v4::Cell::Raw {
                 id,
                 metadata,
@@ -252,12 +236,12 @@ pub trait RenderableCell: Render {
     fn set_selected(&mut self, selected: bool) -> &mut Self;
     fn selected_bg_color(&self, window: &mut Window, cx: &mut Context<Self>) -> Hsla {
         if self.selected() {
-            let mut color = cx.theme().colors().icon_accent;
-            color.fade_out(0.9);
+            let mut color = cx.theme().colors().element_hover;
+            color.fade_out(0.5);
             color
         } else {
-            // TODO: this is wrong
-            cx.theme().colors().tab_bar_background
+            // Not sure if this is correct, previous was TODO: this is wrong
+            gpui::transparent_black()
         }
     }
     fn control(&self, _window: &mut Window, _cx: &mut Context<Self>) -> Option<CellControl> {
@@ -337,11 +321,187 @@ pub struct MarkdownCell {
     metadata: CellMetadata,
     image_cache: Entity<RetainAllImageCache>,
     source: String,
+    editor: Entity<Editor>,
     parsed_markdown: Option<markdown_preview::markdown_elements::ParsedMarkdown>,
     markdown_parsing_task: Task<()>,
+    editing: bool,
     selected: bool,
     cell_position: Option<CellPosition>,
     languages: Arc<LanguageRegistry>,
+    _editor_subscription: gpui::Subscription,
+}
+
+impl EventEmitter<MarkdownCellEvent> for MarkdownCell {}
+
+impl MarkdownCell {
+    pub fn new(
+        id: CellId,
+        metadata: CellMetadata,
+        source: String,
+        languages: Arc<LanguageRegistry>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> Self {
+        let buffer = cx.new(|cx| Buffer::local(source.clone(), cx));
+        let multi_buffer = cx.new(|cx| MultiBuffer::singleton(buffer.clone(), cx));
+
+        let markdown_language = languages.language_for_name("Markdown");
+        cx.spawn_in(window, async move |_this, cx| {
+            if let Some(markdown) = markdown_language.await.log_err() {
+                buffer.update(cx, |buffer, cx| {
+                    buffer.set_language(Some(markdown), cx);
+                });
+            }
+        })
+        .detach();
+
+        let editor = cx.new(|cx| {
+            let mut editor = Editor::new(
+                EditorMode::AutoHeight {
+                    min_lines: 1,
+                    max_lines: Some(1024),
+                },
+                multi_buffer,
+                None,
+                window,
+                cx,
+            );
+
+            let theme = ThemeSettings::get_global(cx);
+            let refinement = TextStyleRefinement {
+                font_family: Some(theme.buffer_font.family.clone()),
+                font_size: Some(theme.buffer_font_size(cx).into()),
+                color: Some(cx.theme().colors().editor_foreground),
+                background_color: Some(gpui::transparent_black()),
+                ..Default::default()
+            };
+
+            editor.set_show_gutter(false, cx);
+            editor.set_text_style_refinement(refinement);
+            editor
+        });
+
+        let markdown_parsing_task = {
+            let languages = languages.clone();
+            let source = source.clone();
+
+            cx.spawn_in(window, async move |this, cx| {
+                let parsed_markdown = cx
+                    .background_spawn(async move {
+                        parse_markdown(&source, None, Some(languages)).await
+                    })
+                    .await;
+
+                this.update(cx, |cell: &mut MarkdownCell, _| {
+                    cell.parsed_markdown = Some(parsed_markdown);
+                })
+                .log_err();
+            })
+        };
+
+        let cell_id = id.clone();
+        let editor_subscription =
+            cx.subscribe(&editor, move |this, _editor, event, cx| match event {
+                editor::EditorEvent::Blurred => {
+                    if this.editing {
+                        this.editing = false;
+                        cx.emit(MarkdownCellEvent::FinishedEditing);
+                        cx.notify();
+                    }
+                }
+                _ => {}
+            });
+
+        let start_editing = source.is_empty();
+        Self {
+            id,
+            metadata,
+            image_cache: RetainAllImageCache::new(cx),
+            source,
+            editor,
+            parsed_markdown: None,
+            markdown_parsing_task,
+            editing: start_editing, // Start in edit mode if empty
+            selected: false,
+            cell_position: None,
+            languages,
+            _editor_subscription: editor_subscription,
+        }
+    }
+
+    pub fn editor(&self) -> &Entity<Editor> {
+        &self.editor
+    }
+
+    pub fn current_source(&self, cx: &App) -> String {
+        let editor = self.editor.read(cx);
+        let buffer = editor.buffer().read(cx);
+        buffer
+            .as_singleton()
+            .map(|b| b.read(cx).text())
+            .unwrap_or_default()
+    }
+
+    pub fn is_dirty(&self, cx: &App) -> bool {
+        self.editor.read(cx).buffer().read(cx).is_dirty(cx)
+    }
+
+    pub fn to_nbformat_cell(&self, cx: &App) -> nbformat::v4::Cell {
+        let source = self.current_source(cx);
+        let source_lines: Vec<String> = source.lines().map(|l| format!("{}\n", l)).collect();
+
+        nbformat::v4::Cell::Markdown {
+            id: self.id.clone(),
+            metadata: self.metadata.clone(),
+            source: source_lines,
+            attachments: None,
+        }
+    }
+
+    pub fn is_editing(&self) -> bool {
+        self.editing
+    }
+
+    pub fn set_editing(&mut self, editing: bool) {
+        self.editing = editing;
+    }
+
+    pub fn reparse_markdown(&mut self, cx: &mut Context<Self>) {
+        let editor = self.editor.read(cx);
+        let buffer = editor.buffer().read(cx);
+        let source = buffer
+            .as_singleton()
+            .map(|b| b.read(cx).text())
+            .unwrap_or_default();
+
+        self.source = source.clone();
+        let languages = self.languages.clone();
+
+        self.markdown_parsing_task = cx.spawn(async move |this, cx| {
+            let parsed_markdown = cx
+                .background_spawn(
+                    async move { parse_markdown(&source, None, Some(languages)).await },
+                )
+                .await;
+
+            this.update(cx, |cell: &mut MarkdownCell, cx| {
+                cell.parsed_markdown = Some(parsed_markdown);
+                cx.notify();
+            })
+            .log_err();
+        });
+    }
+
+    /// Called when user presses Shift+Enter or Ctrl+Enter while editing.
+    /// Finishes editing and signals to move to the next cell.
+    pub fn run(&mut self, cx: &mut Context<Self>) {
+        if self.editing {
+            self.editing = false;
+            cx.emit(MarkdownCellEvent::FinishedEditing);
+            cx.emit(MarkdownCellEvent::Run(self.id.clone()));
+            cx.notify();
+        }
+    }
 }
 
 impl RenderableCell for MarkdownCell {
@@ -388,8 +548,71 @@ impl RenderableCell for MarkdownCell {
 
 impl Render for MarkdownCell {
     fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
+        // If editing, show the editor
+        if self.editing {
+            return v_flex()
+                .size_full()
+                .children(self.cell_position_spacer(true, window, cx))
+                .child(
+                    h_flex()
+                        .w_full()
+                        .pr_6()
+                        .rounded_xs()
+                        .items_start()
+                        .gap(DynamicSpacing::Base08.rems(cx))
+                        .bg(self.selected_bg_color(window, cx))
+                        .child(self.gutter(window, cx))
+                        .child(
+                            div()
+                                .flex_1()
+                                .p_3()
+                                .bg(cx.theme().colors().editor_background)
+                                .rounded_sm()
+                                .child(self.editor.clone())
+                                .on_mouse_down(
+                                    gpui::MouseButton::Left,
+                                    cx.listener(|_this, _event, _window, _cx| {
+                                        // Prevent the click from propagating
+                                    }),
+                                ),
+                        ),
+                )
+                .children(self.cell_position_spacer(false, window, cx));
+        }
+
+        // Preview mode - show rendered markdown
         let Some(parsed) = self.parsed_markdown.as_ref() else {
-            return div();
+            // No parsed content yet, show placeholder that can be clicked to edit
+            let focus_handle = self.editor.focus_handle(cx);
+            return v_flex()
+                .size_full()
+                .children(self.cell_position_spacer(true, window, cx))
+                .child(
+                    h_flex()
+                        .w_full()
+                        .pr_6()
+                        .rounded_xs()
+                        .items_start()
+                        .gap(DynamicSpacing::Base08.rems(cx))
+                        .bg(self.selected_bg_color(window, cx))
+                        .child(self.gutter(window, cx))
+                        .child(
+                            div()
+                                .id("markdown-placeholder")
+                                .flex_1()
+                                .p_3()
+                                .italic()
+                                .text_color(cx.theme().colors().text_muted)
+                                .child("Click to edit markdown...")
+                                .cursor_pointer()
+                                .on_click(cx.listener(move |this, _event, window, cx| {
+                                    this.editing = true;
+                                    window.focus(&this.editor.focus_handle(cx), cx);
+                                    cx.notify();
+                                })),
+                        ),
+                )
+                .children(self.cell_position_spacer(false, window, cx));
         };
 
         let mut markdown_render_context =
@@ -397,7 +620,6 @@ impl Render for MarkdownCell {
 
         v_flex()
             .size_full()
-            // TODO: Move base cell render into trait impl so we don't have to repeat this
             .children(self.cell_position_spacer(true, window, cx))
             .child(
                 h_flex()
@@ -411,11 +633,18 @@ impl Render for MarkdownCell {
                     .child(
                         v_flex()
                             .image_cache(self.image_cache.clone())
+                            .id("markdown-content")
                             .size_full()
                             .flex_1()
                             .p_3()
                             .font_ui(cx)
                             .text_size(TextSize::Default.rems(cx))
+                            .cursor_pointer()
+                            .on_click(cx.listener(|this, _event, window, cx| {
+                                this.editing = true;
+                                window.focus(&this.editor.focus_handle(cx), cx);
+                                cx.notify();
+                            }))
                             .children(parsed.children.iter().map(|child| {
                                 div().relative().child(div().relative().child(
                                     render_markdown_block(child, &mut markdown_render_context),
@@ -423,7 +652,6 @@ impl Render for MarkdownCell {
                             })),
                     ),
             )
-            // TODO: Move base cell render into trait impl so we don't have to repeat this
             .children(self.cell_position_spacer(false, window, cx))
     }
 }
@@ -438,18 +666,260 @@ pub struct CodeCell {
     selected: bool,
     cell_position: Option<CellPosition>,
     language_task: Task<()>,
+    execution_start_time: Option<Instant>,
+    execution_duration: Option<Duration>,
+    is_executing: bool,
 }
 
+impl EventEmitter<CellEvent> for CodeCell {}
+
 impl CodeCell {
+    pub fn new(
+        id: CellId,
+        metadata: CellMetadata,
+        source: String,
+        notebook_language: Shared<Task<Option<Arc<Language>>>>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> Self {
+        let buffer = cx.new(|cx| Buffer::local(source.clone(), cx));
+        let multi_buffer = cx.new(|cx| MultiBuffer::singleton(buffer.clone(), cx));
+
+        let editor_view = cx.new(|cx| {
+            let mut editor = Editor::new(
+                EditorMode::AutoHeight {
+                    min_lines: 1,
+                    max_lines: Some(1024),
+                },
+                multi_buffer,
+                None,
+                window,
+                cx,
+            );
+
+            let theme = ThemeSettings::get_global(cx);
+            let refinement = TextStyleRefinement {
+                font_family: Some(theme.buffer_font.family.clone()),
+                font_size: Some(theme.buffer_font_size(cx).into()),
+                color: Some(cx.theme().colors().editor_foreground),
+                background_color: Some(gpui::transparent_black()),
+                ..Default::default()
+            };
+
+            editor.set_show_gutter(false, cx);
+            editor.set_text_style_refinement(refinement);
+            editor
+        });
+
+        let language_task = cx.spawn_in(window, async move |_this, cx| {
+            let language = notebook_language.await;
+            buffer.update(cx, |buffer, cx| {
+                buffer.set_language(language.clone(), cx);
+            });
+        });
+
+        Self {
+            id,
+            metadata,
+            execution_count: None,
+            source,
+            editor: editor_view,
+            outputs: Vec::new(),
+            selected: false,
+            cell_position: None,
+            language_task,
+            execution_start_time: None,
+            execution_duration: None,
+            is_executing: false,
+        }
+    }
+
+    /// Load a code cell from notebook file data, including existing outputs and execution count
+    pub fn load(
+        id: CellId,
+        metadata: CellMetadata,
+        execution_count: Option<i32>,
+        source: String,
+        outputs: Vec<Output>,
+        notebook_language: Shared<Task<Option<Arc<Language>>>>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> Self {
+        let buffer = cx.new(|cx| Buffer::local(source.clone(), cx));
+        let multi_buffer = cx.new(|cx| MultiBuffer::singleton(buffer.clone(), cx));
+
+        let editor_view = cx.new(|cx| {
+            let mut editor = Editor::new(
+                EditorMode::AutoHeight {
+                    min_lines: 1,
+                    max_lines: Some(1024),
+                },
+                multi_buffer,
+                None,
+                window,
+                cx,
+            );
+
+            let theme = ThemeSettings::get_global(cx);
+            let refinement = TextStyleRefinement {
+                font_family: Some(theme.buffer_font.family.clone()),
+                font_size: Some(theme.buffer_font_size(cx).into()),
+                color: Some(cx.theme().colors().editor_foreground),
+                background_color: Some(gpui::transparent_black()),
+                ..Default::default()
+            };
+
+            editor.set_text(source.clone(), window, cx);
+            editor.set_show_gutter(false, cx);
+            editor.set_text_style_refinement(refinement);
+            editor
+        });
+
+        let language_task = cx.spawn_in(window, async move |_this, cx| {
+            let language = notebook_language.await;
+            buffer.update(cx, |buffer, cx| {
+                buffer.set_language(language.clone(), cx);
+            });
+        });
+
+        Self {
+            id,
+            metadata,
+            execution_count,
+            source,
+            editor: editor_view,
+            outputs,
+            selected: false,
+            cell_position: None,
+            language_task,
+            execution_start_time: None,
+            execution_duration: None,
+            is_executing: false,
+        }
+    }
+
+    pub fn editor(&self) -> &Entity<editor::Editor> {
+        &self.editor
+    }
+
+    pub fn current_source(&self, cx: &App) -> String {
+        let editor = self.editor.read(cx);
+        let buffer = editor.buffer().read(cx);
+        buffer
+            .as_singleton()
+            .map(|b| b.read(cx).text())
+            .unwrap_or_default()
+    }
+
     pub fn is_dirty(&self, cx: &App) -> bool {
         self.editor.read(cx).buffer().read(cx).is_dirty(cx)
     }
+
+    pub fn to_nbformat_cell(&self, cx: &App) -> nbformat::v4::Cell {
+        let source = self.current_source(cx);
+        let source_lines: Vec<String> = source.lines().map(|l| format!("{}\n", l)).collect();
+
+        let outputs = self.outputs_to_nbformat(cx);
+
+        nbformat::v4::Cell::Code {
+            id: self.id.clone(),
+            metadata: self.metadata.clone(),
+            execution_count: self.execution_count,
+            source: source_lines,
+            outputs,
+        }
+    }
+
+    fn outputs_to_nbformat(&self, cx: &App) -> Vec<nbformat::v4::Output> {
+        self.outputs
+            .iter()
+            .filter_map(|output| output.to_nbformat(cx))
+            .collect()
+    }
+
     pub fn has_outputs(&self) -> bool {
         !self.outputs.is_empty()
     }
 
     pub fn clear_outputs(&mut self) {
         self.outputs.clear();
+        self.execution_duration = None;
+    }
+
+    pub fn start_execution(&mut self) {
+        self.execution_start_time = Some(Instant::now());
+        self.execution_duration = None;
+        self.is_executing = true;
+    }
+
+    pub fn finish_execution(&mut self) {
+        if let Some(start_time) = self.execution_start_time.take() {
+            self.execution_duration = Some(start_time.elapsed());
+        }
+        self.is_executing = false;
+    }
+
+    pub fn is_executing(&self) -> bool {
+        self.is_executing
+    }
+
+    pub fn execution_duration(&self) -> Option<Duration> {
+        self.execution_duration
+    }
+
+    fn format_duration(duration: Duration) -> String {
+        let total_secs = duration.as_secs_f64();
+        if total_secs < 1.0 {
+            format!("{:.0}ms", duration.as_millis())
+        } else if total_secs < 60.0 {
+            format!("{:.1}s", total_secs)
+        } else {
+            let minutes = (total_secs / 60.0).floor() as u64;
+            let secs = total_secs % 60.0;
+            format!("{}m {:.1}s", minutes, secs)
+        }
+    }
+
+    pub fn handle_message(
+        &mut self,
+        message: &JupyterMessage,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        match &message.content {
+            JupyterMessageContent::StreamContent(stream) => {
+                self.outputs.push(Output::Stream {
+                    content: cx.new(|cx| TerminalOutput::from(&stream.text, window, cx)),
+                });
+            }
+            JupyterMessageContent::DisplayData(display_data) => {
+                self.outputs
+                    .push(Output::new(&display_data.data, None, window, cx));
+            }
+            JupyterMessageContent::ExecuteResult(execute_result) => {
+                self.outputs
+                    .push(Output::new(&execute_result.data, None, window, cx));
+            }
+            JupyterMessageContent::ExecuteInput(input) => {
+                self.execution_count = serde_json::to_value(&input.execution_count)
+                    .ok()
+                    .and_then(|v| v.as_i64())
+                    .map(|v| v as i32);
+            }
+            JupyterMessageContent::ExecuteReply(_) => {
+                self.finish_execution();
+            }
+            JupyterMessageContent::ErrorOutput(error) => {
+                self.outputs.push(Output::ErrorOutput(ErrorView {
+                    ename: error.ename.clone(),
+                    evalue: error.evalue.clone(),
+                    traceback: cx
+                        .new(|cx| TerminalOutput::from(&error.traceback.join("\n"), window, cx)),
+                }));
+            }
+            _ => {}
+        }
+        cx.notify();
     }
 
     fn output_control(&self) -> Option<CellControlType> {
@@ -522,13 +992,22 @@ impl RenderableCell for CodeCell {
     }
 
     fn control(&self, window: &mut Window, cx: &mut Context<Self>) -> Option<CellControl> {
-        let cell_control = if self.has_outputs() {
-            CellControl::new("rerun-cell", CellControlType::RerunCell)
+        let control_type = if self.has_outputs() {
+            CellControlType::RerunCell
         } else {
-            CellControl::new("run-cell", CellControlType::RunCell)
-                .on_click(cx.listener(move |this, _, window, cx| this.run(window, cx)))
+            CellControlType::RunCell
         };
 
+        let cell_control = CellControl::new(
+            if self.has_outputs() {
+                "rerun-cell"
+            } else {
+                "run-cell"
+            },
+            control_type,
+        )
+        .on_click(cx.listener(move |this, _, window, cx| this.run(window, cx)));
+
         Some(cell_control)
     }
 
@@ -549,11 +1028,62 @@ impl RenderableCell for CodeCell {
         self.cell_position = Some(cell_position);
         self
     }
+
+    fn gutter(&self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
+        let is_selected = self.selected();
+        let execution_count = self.execution_count;
+
+        div()
+            .relative()
+            .h_full()
+            .w(px(GUTTER_WIDTH))
+            .child(
+                div()
+                    .w(px(GUTTER_WIDTH))
+                    .flex()
+                    .flex_none()
+                    .justify_center()
+                    .h_full()
+                    .child(
+                        div()
+                            .flex_none()
+                            .w(px(1.))
+                            .h_full()
+                            .when(is_selected, |this| this.bg(cx.theme().colors().icon_accent))
+                            .when(!is_selected, |this| this.bg(cx.theme().colors().border)),
+                    ),
+            )
+            .when_some(self.control(window, cx), |this, control| {
+                this.child(
+                    div()
+                        .absolute()
+                        .top(px(CODE_BLOCK_INSET - 2.0))
+                        .left_0()
+                        .flex()
+                        .flex_col()
+                        .w(px(GUTTER_WIDTH))
+                        .items_center()
+                        .justify_center()
+                        .bg(cx.theme().colors().tab_bar_background)
+                        .child(control.button)
+                        .when_some(execution_count, |this, count| {
+                            this.child(
+                                div()
+                                    .mt_1()
+                                    .text_xs()
+                                    .text_color(cx.theme().colors().text_muted)
+                                    .child(format!("{}", count)),
+                            )
+                        }),
+                )
+            })
+    }
 }
 
 impl RunnableCell for CodeCell {
     fn run(&mut self, window: &mut Window, cx: &mut Context<Self>) {
         println!("Running code cell: {}", self.id);
+        cx.emit(CellEvent::Run(self.id.clone()));
     }
 
     fn execution_count(&self) -> Option<i32> {
@@ -569,6 +1099,16 @@ impl RunnableCell for CodeCell {
 
 impl Render for CodeCell {
     fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
+        // get the language from the editor's buffer
+        let language_name = self
+            .editor
+            .read(cx)
+            .buffer()
+            .read(cx)
+            .as_singleton()
+            .and_then(|buffer| buffer.read(cx).language())
+            .map(|lang| lang.name().to_string());
+
         v_flex()
             .size_full()
             // TODO: Move base cell render into trait impl so we don't have to repeat this
@@ -586,6 +1126,7 @@ impl Render for CodeCell {
                     .child(
                         div().py_1p5().w_full().child(
                             div()
+                                .relative()
                                 .flex()
                                 .size_full()
                                 .flex_1()
@@ -595,73 +1136,137 @@ impl Render for CodeCell {
                                 .border_1()
                                 .border_color(cx.theme().colors().border)
                                 .bg(cx.theme().colors().editor_background)
-                                .child(div().w_full().child(self.editor.clone())),
+                                .child(div().w_full().child(self.editor.clone()))
+                                // lang badge in top-right corner
+                                .when_some(language_name, |this, name| {
+                                    this.child(
+                                        div()
+                                            .absolute()
+                                            .top_1()
+                                            .right_2()
+                                            .px_2()
+                                            .py_0p5()
+                                            .rounded_md()
+                                            .bg(cx.theme().colors().element_background.opacity(0.7))
+                                            .text_xs()
+                                            .text_color(cx.theme().colors().text_muted)
+                                            .child(name),
+                                    )
+                                }),
                         ),
                     ),
             )
             // Output portion
-            .child(
-                h_flex()
-                    .w_full()
-                    .pr_6()
-                    .rounded_xs()
-                    .items_start()
-                    .gap(DynamicSpacing::Base08.rems(cx))
-                    .bg(self.selected_bg_color(window, cx))
-                    .child(self.gutter_output(window, cx))
-                    .child(
-                        div().py_1p5().w_full().child(
-                            div()
-                                .flex()
-                                .size_full()
-                                .flex_1()
-                                .py_3()
-                                .px_5()
-                                .rounded_lg()
-                                .border_1()
-                                // .border_color(cx.theme().colors().border)
-                                // .bg(cx.theme().colors().editor_background)
-                                .child(div().w_full().children(self.outputs.iter().map(
-                                    |output| {
-                                        let content = match output {
-                                            Output::Plain { content, .. } => {
-                                                Some(content.clone().into_any_element())
-                                            }
-                                            Output::Markdown { content, .. } => {
-                                                Some(content.clone().into_any_element())
-                                            }
-                                            Output::Stream { content, .. } => {
-                                                Some(content.clone().into_any_element())
-                                            }
-                                            Output::Image { content, .. } => {
-                                                Some(content.clone().into_any_element())
-                                            }
-                                            Output::Message(message) => Some(
-                                                div().child(message.clone()).into_any_element(),
-                                            ),
-                                            Output::Table { content, .. } => {
-                                                Some(content.clone().into_any_element())
-                                            }
-                                            Output::ErrorOutput(error_view) => {
-                                                error_view.render(window, cx)
-                                            }
-                                            Output::ClearOutputWaitMarker => None,
-                                        };
-
-                                        div()
-                                            // .w_full()
-                                            // .mt_3()
-                                            // .p_3()
-                                            // .rounded_sm()
-                                            // .bg(cx.theme().colors().editor_background)
-                                            // .border(px(1.))
-                                            // .border_color(cx.theme().colors().border)
-                                            // .shadow_xs()
-                                            .children(content)
-                                    },
-                                ))),
-                        ),
-                    ),
+            .when(
+                self.has_outputs() || self.execution_duration.is_some() || self.is_executing,
+                |this| {
+                    let execution_time_label = self.execution_duration.map(Self::format_duration);
+                    let is_executing = self.is_executing;
+                    this.child(
+                        h_flex()
+                            .w_full()
+                            .pr_6()
+                            .rounded_xs()
+                            .items_start()
+                            .gap(DynamicSpacing::Base08.rems(cx))
+                            .bg(self.selected_bg_color(window, cx))
+                            .child(self.gutter_output(window, cx))
+                            .child(
+                                div().py_1p5().w_full().child(
+                                    v_flex()
+                                        .size_full()
+                                        .flex_1()
+                                        .py_3()
+                                        .px_5()
+                                        .rounded_lg()
+                                        .border_1()
+                                        // execution status/time at the TOP
+                                        .when(
+                                            is_executing || execution_time_label.is_some(),
+                                            |this| {
+                                                let time_element = if is_executing {
+                                                    h_flex()
+                                                        .gap_1()
+                                                        .items_center()
+                                                        .child(
+                                                            Icon::new(IconName::ArrowCircle)
+                                                                .size(IconSize::XSmall)
+                                                                .color(Color::Warning)
+                                                                .with_rotate_animation(2)
+                                                                .into_any_element(),
+                                                        )
+                                                        .child(
+                                                            div()
+                                                                .text_xs()
+                                                                .text_color(
+                                                                    cx.theme().colors().text_muted,
+                                                                )
+                                                                .child("Running..."),
+                                                        )
+                                                        .into_any_element()
+                                                } else if let Some(duration_text) =
+                                                    execution_time_label.clone()
+                                                {
+                                                    h_flex()
+                                                        .gap_1()
+                                                        .items_center()
+                                                        .child(
+                                                            Icon::new(IconName::Check)
+                                                                .size(IconSize::XSmall)
+                                                                .color(Color::Success),
+                                                        )
+                                                        .child(
+                                                            div()
+                                                                .text_xs()
+                                                                .text_color(
+                                                                    cx.theme().colors().text_muted,
+                                                                )
+                                                                .child(duration_text),
+                                                        )
+                                                        .into_any_element()
+                                                } else {
+                                                    div().into_any_element()
+                                                };
+                                                this.child(div().mb_2().child(time_element))
+                                            },
+                                        )
+                                        // output at bottom
+                                        .child(div().w_full().children(self.outputs.iter().map(
+                                            |output| {
+                                                let content = match output {
+                                                    Output::Plain { content, .. } => {
+                                                        Some(content.clone().into_any_element())
+                                                    }
+                                                    Output::Markdown { content, .. } => {
+                                                        Some(content.clone().into_any_element())
+                                                    }
+                                                    Output::Stream { content, .. } => {
+                                                        Some(content.clone().into_any_element())
+                                                    }
+                                                    Output::Image { content, .. } => {
+                                                        Some(content.clone().into_any_element())
+                                                    }
+                                                    Output::Message(message) => Some(
+                                                        div()
+                                                            .child(message.clone())
+                                                            .into_any_element(),
+                                                    ),
+                                                    Output::Table { content, .. } => {
+                                                        Some(content.clone().into_any_element())
+                                                    }
+                                                    Output::ErrorOutput(error_view) => {
+                                                        error_view.render(window, cx)
+                                                    }
+                                                    Output::ClearOutputWaitMarker => None,
+                                                };
+
+                                                div().children(content)
+                                            },
+                                        ))),
+                                ),
+                            ),
+                    )
+                },
             )
             // TODO: Move base cell render into trait impl so we don't have to repeat this
             .children(self.cell_position_spacer(false, window, cx))

crates/repl/src/notebook/notebook_ui.rs 🔗

@@ -12,18 +12,31 @@ use gpui::{
     AnyElement, App, Entity, EventEmitter, FocusHandle, Focusable, ListScrollEvent, ListState,
     Point, Task, actions, list, prelude::*,
 };
+use jupyter_protocol::JupyterKernelspec;
 use language::{Language, LanguageRegistry};
 use project::{Project, ProjectEntryId, ProjectPath};
-use ui::{Tooltip, prelude::*};
+use settings::Settings as _;
+use ui::{CommonAnimationExt, Tooltip, prelude::*};
 use workspace::item::{ItemEvent, SaveOptions, TabContentParams};
 use workspace::searchable::SearchableItemHandle;
 use workspace::{Item, ItemHandle, Pane, ProjectItem, ToolbarItemLocation};
-use workspace::{ToolbarItemEvent, ToolbarItemView};
 
-use super::{Cell, CellPosition, RenderableCell};
+use super::{Cell, CellEvent, CellPosition, MarkdownCellEvent, RenderableCell};
 
 use nbformat::v4::CellId;
 use nbformat::v4::Metadata as NotebookMetadata;
+use serde_json;
+use uuid::Uuid;
+
+use crate::components::{KernelPickerDelegate, KernelSelector};
+use crate::kernels::{
+    Kernel, KernelSession, KernelSpecification, KernelStatus, LocalKernelSpecification,
+    NativeRunningKernel, RemoteRunningKernel,
+};
+use crate::repl_store::ReplStore;
+use picker::Picker;
+use runtimelib::{ExecuteRequest, JupyterMessage, JupyterMessageContent};
+use ui::PopoverMenuHandle;
 
 actions!(
     notebook,
@@ -32,6 +45,8 @@ actions!(
         OpenNotebook,
         /// Runs all cells in the notebook.
         RunAll,
+        /// Runs the current cell.
+        Run,
         /// Clears all cell outputs.
         ClearOutputs,
         /// Moves the current cell up.
@@ -42,6 +57,10 @@ actions!(
         AddMarkdownBlock,
         /// Adds a new code cell.
         AddCodeBlock,
+        /// Restarts the kernel.
+        RestartKernel,
+        /// Interrupts the current execution.
+        InterruptKernel,
     ]
 );
 
@@ -74,16 +93,23 @@ pub fn init(cx: &mut App) {
 pub struct NotebookEditor {
     languages: Arc<LanguageRegistry>,
     project: Entity<Project>,
+    worktree_id: project::WorktreeId,
 
     focus_handle: FocusHandle,
     notebook_item: Entity<NotebookItem>,
+    notebook_language: Shared<Task<Option<Arc<Language>>>>,
 
     remote_id: Option<ViewId>,
     cell_list: ListState,
 
     selected_cell_index: usize,
     cell_order: Vec<CellId>,
+    original_cell_order: Vec<CellId>,
     cell_map: HashMap<CellId, Cell>,
+    kernel: Kernel,
+    kernel_specification: Option<KernelSpecification>,
+    execution_requests: HashMap<String, CellId>,
+    kernel_picker_handle: PopoverMenuHandle<Picker<KernelPickerDelegate>>,
 }
 
 impl NotebookEditor {
@@ -97,6 +123,7 @@ impl NotebookEditor {
 
         let languages = project.read(cx).languages().clone();
         let language_name = notebook_item.read(cx).language_name();
+        let worktree_id = notebook_item.read(cx).project_path.worktree_id;
 
         let notebook_language = notebook_item.read(cx).notebook_language();
         let notebook_language = cx
@@ -116,10 +143,85 @@ impl NotebookEditor {
         {
             let cell_id = cell.id();
             cell_order.push(cell_id.clone());
-            cell_map.insert(
-                cell_id.clone(),
-                Cell::load(cell, &languages, notebook_language.clone(), window, cx),
-            );
+            let cell_entity = Cell::load(cell, &languages, notebook_language.clone(), window, cx);
+
+            match &cell_entity {
+                Cell::Code(code_cell) => {
+                    let cell_id_for_focus = cell_id.clone();
+                    cx.subscribe(code_cell, move |this, cell, event, cx| match event {
+                        CellEvent::Run(cell_id) => this.execute_cell(cell_id.clone(), cx),
+                        CellEvent::FocusedIn(_) => {
+                            if let Some(index) = this
+                                .cell_order
+                                .iter()
+                                .position(|id| id == &cell_id_for_focus)
+                            {
+                                this.selected_cell_index = index;
+                                cx.notify();
+                            }
+                        }
+                    })
+                    .detach();
+
+                    let cell_id_for_editor = cell_id.clone();
+                    let editor = code_cell.read(cx).editor().clone();
+                    cx.subscribe(&editor, move |this, _editor, event, cx| {
+                        if let editor::EditorEvent::Focused = event {
+                            if let Some(index) = this
+                                .cell_order
+                                .iter()
+                                .position(|id| id == &cell_id_for_editor)
+                            {
+                                this.selected_cell_index = index;
+                                cx.notify();
+                            }
+                        }
+                    })
+                    .detach();
+                }
+                Cell::Markdown(markdown_cell) => {
+                    let cell_id_for_focus = cell_id.clone();
+                    cx.subscribe(
+                        markdown_cell,
+                        move |_this, cell, event: &MarkdownCellEvent, cx| {
+                            match event {
+                                MarkdownCellEvent::FinishedEditing => {
+                                    cell.update(cx, |cell, cx| {
+                                        cell.reparse_markdown(cx);
+                                    });
+                                }
+                                MarkdownCellEvent::Run(_cell_id) => {
+                                    // run is handled separately by move_to_next_cell
+                                    // Just reparse here
+                                    cell.update(cx, |cell, cx| {
+                                        cell.reparse_markdown(cx);
+                                    });
+                                }
+                            }
+                        },
+                    )
+                    .detach();
+
+                    let cell_id_for_editor = cell_id.clone();
+                    let editor = markdown_cell.read(cx).editor().clone();
+                    cx.subscribe(&editor, move |this, _editor, event, cx| {
+                        if let editor::EditorEvent::Focused = event {
+                            if let Some(index) = this
+                                .cell_order
+                                .iter()
+                                .position(|id| id == &cell_id_for_editor)
+                            {
+                                this.selected_cell_index = index;
+                                cx.notify();
+                            }
+                        }
+                    })
+                    .detach();
+                }
+                Cell::Raw(_) => {}
+            }
+
+            cell_map.insert(cell_id.clone(), cell_entity);
         }
 
         let notebook_handle = cx.entity().downgrade();
@@ -128,16 +230,266 @@ impl NotebookEditor {
         let this = cx.entity();
         let cell_list = ListState::new(cell_count, gpui::ListAlignment::Top, px(1000.));
 
-        Self {
+        let mut editor = Self {
             project,
             languages: languages.clone(),
+            worktree_id,
             focus_handle,
             notebook_item,
+            notebook_language,
             remote_id: None,
             cell_list,
             selected_cell_index: 0,
             cell_order: cell_order.clone(),
+            original_cell_order: cell_order.clone(),
             cell_map: cell_map.clone(),
+            kernel: Kernel::StartingKernel(Task::ready(()).shared()),
+            kernel_specification: None,
+            execution_requests: HashMap::default(),
+            kernel_picker_handle: PopoverMenuHandle::default(),
+        };
+        editor.launch_kernel(window, cx);
+        editor
+    }
+
+    fn has_structural_changes(&self) -> bool {
+        self.cell_order != self.original_cell_order
+    }
+
+    fn has_content_changes(&self, cx: &App) -> bool {
+        self.cell_map.values().any(|cell| cell.is_dirty(cx))
+    }
+
+    pub fn to_notebook(&self, cx: &App) -> nbformat::v4::Notebook {
+        let cells: Vec<nbformat::v4::Cell> = self
+            .cell_order
+            .iter()
+            .filter_map(|cell_id| {
+                self.cell_map
+                    .get(cell_id)
+                    .map(|cell| cell.to_nbformat_cell(cx))
+            })
+            .collect();
+
+        let metadata = self.notebook_item.read(cx).notebook.metadata.clone();
+
+        nbformat::v4::Notebook {
+            metadata,
+            nbformat: 4,
+            nbformat_minor: 5,
+            cells,
+        }
+    }
+
+    pub fn mark_as_saved(&mut self, cx: &mut Context<Self>) {
+        self.original_cell_order = self.cell_order.clone();
+
+        for cell in self.cell_map.values() {
+            match cell {
+                Cell::Code(code_cell) => {
+                    code_cell.update(cx, |code_cell, cx| {
+                        let editor = code_cell.editor();
+                        editor.update(cx, |editor, cx| {
+                            editor.buffer().update(cx, |buffer, cx| {
+                                if let Some(buf) = buffer.as_singleton() {
+                                    buf.update(cx, |b, cx| {
+                                        let version = b.version();
+                                        b.did_save(version, None, cx);
+                                    });
+                                }
+                            });
+                        });
+                    });
+                }
+                Cell::Markdown(markdown_cell) => {
+                    markdown_cell.update(cx, |markdown_cell, cx| {
+                        let editor = markdown_cell.editor();
+                        editor.update(cx, |editor, cx| {
+                            editor.buffer().update(cx, |buffer, cx| {
+                                if let Some(buf) = buffer.as_singleton() {
+                                    buf.update(cx, |b, cx| {
+                                        let version = b.version();
+                                        b.did_save(version, None, cx);
+                                    });
+                                }
+                            });
+                        });
+                    });
+                }
+                Cell::Raw(_) => {}
+            }
+        }
+        cx.notify();
+    }
+
+    fn launch_kernel(&mut self, window: &mut Window, cx: &mut Context<Self>) {
+        // use default Python kernel if no specification is set
+        let spec = self.kernel_specification.clone().unwrap_or_else(|| {
+            KernelSpecification::Jupyter(LocalKernelSpecification {
+                name: "python3".to_string(),
+                path: PathBuf::from("python3"),
+                kernelspec: JupyterKernelspec {
+                    argv: vec![
+                        "python3".to_string(),
+                        "-m".to_string(),
+                        "ipykernel_launcher".to_string(),
+                        "-f".to_string(),
+                        "{connection_file}".to_string(),
+                    ],
+                    display_name: "Python 3".to_string(),
+                    language: "python".to_string(),
+                    interrupt_mode: None,
+                    metadata: None,
+                    env: None,
+                },
+            })
+        });
+
+        self.launch_kernel_with_spec(spec, window, cx);
+    }
+
+    fn launch_kernel_with_spec(
+        &mut self,
+        spec: KernelSpecification,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        let entity_id = cx.entity_id();
+        let working_directory = self
+            .project
+            .read(cx)
+            .worktrees(cx)
+            .next()
+            .map(|worktree| worktree.read(cx).abs_path().to_path_buf())
+            .unwrap_or_else(std::env::temp_dir);
+        let fs = self.project.read(cx).fs().clone();
+        let view = cx.entity();
+
+        self.kernel_specification = Some(spec.clone());
+
+        let kernel_task = match spec {
+            KernelSpecification::Jupyter(local_spec)
+            | KernelSpecification::PythonEnv(local_spec) => NativeRunningKernel::new(
+                local_spec,
+                entity_id,
+                working_directory,
+                fs,
+                view,
+                window,
+                cx,
+            ),
+            KernelSpecification::Remote(remote_spec) => {
+                RemoteRunningKernel::new(remote_spec, working_directory, view, window, cx)
+            }
+        };
+
+        let pending_kernel = cx
+            .spawn(async move |this, cx| {
+                let kernel = kernel_task.await;
+
+                match kernel {
+                    Ok(kernel) => {
+                        this.update(cx, |editor, cx| {
+                            editor.kernel = Kernel::RunningKernel(kernel);
+                            cx.notify();
+                        })
+                        .ok();
+                    }
+                    Err(err) => {
+                        this.update(cx, |editor, cx| {
+                            editor.kernel = Kernel::ErroredLaunch(err.to_string());
+                            cx.notify();
+                        })
+                        .ok();
+                    }
+                }
+            })
+            .shared();
+
+        self.kernel = Kernel::StartingKernel(pending_kernel);
+        cx.notify();
+    }
+
+    // Note: Python environments are only detected as kernels if ipykernel is installed.
+    // Users need to run `pip install ipykernel` (or `uv pip install ipykernel`) in their
+    // virtual environment for it to appear in the kernel selector.
+    // This happens because we have an ipykernel check inside the function python_env_kernel_specification in mod.rs L:121
+
+    fn change_kernel(
+        &mut self,
+        spec: KernelSpecification,
+        window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        if let Kernel::RunningKernel(kernel) = &mut self.kernel {
+            kernel.force_shutdown(window, cx).detach();
+        }
+
+        self.execution_requests.clear();
+
+        self.launch_kernel_with_spec(spec, window, cx);
+    }
+
+    fn restart_kernel(&mut self, _: &RestartKernel, window: &mut Window, cx: &mut Context<Self>) {
+        if let Some(spec) = self.kernel_specification.clone() {
+            if let Kernel::RunningKernel(kernel) = &mut self.kernel {
+                kernel.force_shutdown(window, cx).detach();
+            }
+
+            self.kernel = Kernel::Restarting;
+            cx.notify();
+
+            self.launch_kernel_with_spec(spec, window, cx);
+        }
+    }
+
+    fn interrupt_kernel(
+        &mut self,
+        _: &InterruptKernel,
+        _window: &mut Window,
+        cx: &mut Context<Self>,
+    ) {
+        if let Kernel::RunningKernel(kernel) = &self.kernel {
+            let interrupt_request = runtimelib::InterruptRequest {};
+            let message: JupyterMessage = interrupt_request.into();
+            kernel.request_tx().try_send(message).ok();
+            cx.notify();
+        }
+    }
+
+    fn execute_cell(&mut self, cell_id: CellId, cx: &mut Context<Self>) {
+        let code = if let Some(Cell::Code(cell)) = self.cell_map.get(&cell_id) {
+            let editor = cell.read(cx).editor().clone();
+            let buffer = editor.read(cx).buffer().read(cx);
+            buffer
+                .as_singleton()
+                .map(|b| b.read(cx).text())
+                .unwrap_or_default()
+        } else {
+            return;
+        };
+
+        if let Some(Cell::Code(cell)) = self.cell_map.get(&cell_id) {
+            cell.update(cx, |cell, cx| {
+                if cell.has_outputs() {
+                    cell.clear_outputs();
+                }
+                cell.start_execution();
+                cx.notify();
+            });
+        }
+
+        let request = ExecuteRequest {
+            code,
+            ..Default::default()
+        };
+        let message: JupyterMessage = request.into();
+        let msg_id = message.header.msg_id.clone();
+
+        self.execution_requests.insert(msg_id, cell_id.clone());
+
+        if let Kernel::RunningKernel(kernel) = &mut self.kernel {
+            kernel.request_tx().try_send(message).ok();
         }
     }
 
@@ -154,15 +506,73 @@ impl NotebookEditor {
     fn clear_outputs(&mut self, window: &mut Window, cx: &mut Context<Self>) {
         for cell in self.cell_map.values() {
             if let Cell::Code(code_cell) = cell {
-                code_cell.update(cx, |cell, _cx| {
+                code_cell.update(cx, |cell, cx| {
                     cell.clear_outputs();
+                    cx.notify();
                 });
             }
         }
+        cx.notify();
     }
 
     fn run_cells(&mut self, window: &mut Window, cx: &mut Context<Self>) {
-        println!("Cells would all run here, if that was implemented!");
+        println!("Cells would run here!");
+        for cell_id in self.cell_order.clone() {
+            self.execute_cell(cell_id, cx);
+        }
+    }
+
+    fn run_current_cell(&mut self, _: &Run, window: &mut Window, cx: &mut Context<Self>) {
+        if let Some(cell_id) = self.cell_order.get(self.selected_cell_index).cloned() {
+            if let Some(cell) = self.cell_map.get(&cell_id) {
+                match cell {
+                    Cell::Code(_) => {
+                        self.execute_cell(cell_id, cx);
+                    }
+                    Cell::Markdown(markdown_cell) => {
+                        // for markdown, finish editing and move to next cell
+                        let is_editing = markdown_cell.read(cx).is_editing();
+                        if is_editing {
+                            markdown_cell.update(cx, |cell, cx| {
+                                cell.run(cx);
+                            });
+                            // move to the next cell
+                            // Discussion can be done on this default implementation
+                            self.move_to_next_cell(window, cx);
+                        }
+                    }
+                    Cell::Raw(_) => {}
+                }
+            }
+        }
+    }
+
+    // Discussion can be done on this default implementation
+    /// Moves focus to the next cell, or creates a new code cell if at the end
+    fn move_to_next_cell(&mut self, window: &mut Window, cx: &mut Context<Self>) {
+        if self.selected_cell_index < self.cell_order.len() - 1 {
+            self.selected_cell_index += 1;
+            // focus the new cell's editor
+            if let Some(cell_id) = self.cell_order.get(self.selected_cell_index) {
+                if let Some(cell) = self.cell_map.get(cell_id) {
+                    match cell {
+                        Cell::Code(code_cell) => {
+                            let editor = code_cell.read(cx).editor();
+                            window.focus(&editor.focus_handle(cx), cx);
+                        }
+                        Cell::Markdown(markdown_cell) => {
+                            // Don't auto-enter edit mode for next markdown cell
+                            // Just select it
+                        }
+                        Cell::Raw(_) => {}
+                    }
+                }
+            }
+            cx.notify();
+        } else {
+            // in the end, could optionally create a new cell
+            // For now, just stay on the current cell
+        }
     }
 
     fn open_notebook(&mut self, _: &OpenNotebook, _window: &mut Window, _cx: &mut Context<Self>) {
@@ -171,18 +581,132 @@ impl NotebookEditor {
 
     fn move_cell_up(&mut self, window: &mut Window, cx: &mut Context<Self>) {
         println!("Move cell up triggered");
+        if self.selected_cell_index > 0 {
+            self.cell_order
+                .swap(self.selected_cell_index, self.selected_cell_index - 1);
+            self.selected_cell_index -= 1;
+            cx.notify();
+        }
     }
 
     fn move_cell_down(&mut self, window: &mut Window, cx: &mut Context<Self>) {
         println!("Move cell down triggered");
+        if self.selected_cell_index < self.cell_order.len() - 1 {
+            self.cell_order
+                .swap(self.selected_cell_index, self.selected_cell_index + 1);
+            self.selected_cell_index += 1;
+            cx.notify();
+        }
     }
 
     fn add_markdown_block(&mut self, window: &mut Window, cx: &mut Context<Self>) {
-        println!("Add markdown block triggered");
+        let new_cell_id: CellId = Uuid::new_v4().into();
+        let languages = self.languages.clone();
+        let metadata: nbformat::v4::CellMetadata =
+            serde_json::from_str("{}").expect("empty object should parse");
+
+        let markdown_cell = cx.new(|cx| {
+            super::MarkdownCell::new(
+                new_cell_id.clone(),
+                metadata,
+                String::new(),
+                languages,
+                window,
+                cx,
+            )
+        });
+
+        let insert_index = self.selected_cell_index + 1;
+        self.cell_order.insert(insert_index, new_cell_id.clone());
+        self.cell_map
+            .insert(new_cell_id.clone(), Cell::Markdown(markdown_cell.clone()));
+        self.selected_cell_index = insert_index;
+
+        cx.subscribe(
+            &markdown_cell,
+            move |_this, cell, event: &MarkdownCellEvent, cx| match event {
+                MarkdownCellEvent::FinishedEditing | MarkdownCellEvent::Run(_) => {
+                    cell.update(cx, |cell, cx| {
+                        cell.reparse_markdown(cx);
+                    });
+                }
+            },
+        )
+        .detach();
+
+        let cell_id_for_editor = new_cell_id.clone();
+        let editor = markdown_cell.read(cx).editor().clone();
+        cx.subscribe(&editor, move |this, _editor, event, cx| {
+            if let editor::EditorEvent::Focused = event {
+                if let Some(index) = this
+                    .cell_order
+                    .iter()
+                    .position(|id| id == &cell_id_for_editor)
+                {
+                    this.selected_cell_index = index;
+                    cx.notify();
+                }
+            }
+        })
+        .detach();
+
+        self.cell_list.reset(self.cell_order.len());
+        cx.notify();
     }
 
     fn add_code_block(&mut self, window: &mut Window, cx: &mut Context<Self>) {
-        println!("Add code block triggered");
+        let new_cell_id: CellId = Uuid::new_v4().into();
+        let notebook_language = self.notebook_language.clone();
+        let metadata: nbformat::v4::CellMetadata =
+            serde_json::from_str("{}").expect("empty object should parse");
+
+        let code_cell = cx.new(|cx| {
+            super::CodeCell::new(
+                new_cell_id.clone(),
+                metadata,
+                String::new(),
+                notebook_language,
+                window,
+                cx,
+            )
+        });
+
+        let insert_index = self.selected_cell_index + 1;
+        self.cell_order.insert(insert_index, new_cell_id.clone());
+        self.cell_map
+            .insert(new_cell_id.clone(), Cell::Code(code_cell.clone()));
+        self.selected_cell_index = insert_index;
+
+        let cell_id_for_run = new_cell_id.clone();
+        cx.subscribe(&code_cell, move |this, _cell, event, cx| match event {
+            CellEvent::Run(cell_id) => this.execute_cell(cell_id.clone(), cx),
+            CellEvent::FocusedIn(_) => {
+                if let Some(index) = this.cell_order.iter().position(|id| id == &cell_id_for_run) {
+                    this.selected_cell_index = index;
+                    cx.notify();
+                }
+            }
+        })
+        .detach();
+
+        let cell_id_for_editor = new_cell_id.clone();
+        let editor = code_cell.read(cx).editor().clone();
+        cx.subscribe(&editor, move |this, _editor, event, cx| {
+            if let editor::EditorEvent::Focused = event {
+                if let Some(index) = this
+                    .cell_order
+                    .iter()
+                    .position(|id| id == &cell_id_for_editor)
+                {
+                    this.selected_cell_index = index;
+                    cx.notify();
+                }
+            }
+        })
+        .detach();
+
+        self.cell_list.reset(self.cell_order.len());
+        cx.notify();
     }
 
     fn cell_count(&self) -> usize {
@@ -415,19 +939,160 @@ impl NotebookEditor {
                 v_flex()
                     .gap(DynamicSpacing::Base08.rems(cx))
                     .items_center()
-                    .child(Self::render_notebook_control(
-                        "more-menu",
-                        IconName::Ellipsis,
-                        window,
-                        cx,
-                    ))
                     .child(
-                        Self::button_group(window, cx)
-                            .child(IconButton::new("repl", IconName::ReplNeutral)),
+                        Self::render_notebook_control("more-menu", IconName::Ellipsis, window, cx)
+                            .tooltip(move |window, cx| (Tooltip::text("More options"))(window, cx)),
+                    )
+                    .child(Self::button_group(window, cx).child({
+                        let kernel_status = self.kernel.status();
+                        let (icon, icon_color) = match &kernel_status {
+                            KernelStatus::Idle => (IconName::ReplNeutral, Color::Success),
+                            KernelStatus::Busy => (IconName::ReplNeutral, Color::Warning),
+                            KernelStatus::Starting => (IconName::ReplNeutral, Color::Muted),
+                            KernelStatus::Error => (IconName::ReplNeutral, Color::Error),
+                            KernelStatus::ShuttingDown => (IconName::ReplNeutral, Color::Muted),
+                            KernelStatus::Shutdown => (IconName::ReplNeutral, Color::Disabled),
+                            KernelStatus::Restarting => (IconName::ReplNeutral, Color::Warning),
+                        };
+                        let kernel_name = self
+                            .kernel_specification
+                            .as_ref()
+                            .map(|spec| spec.name().to_string())
+                            .unwrap_or_else(|| "Select Kernel".to_string());
+                        IconButton::new("repl", icon)
+                            .icon_color(icon_color)
+                            .tooltip(move |window, cx| {
+                                Tooltip::text(format!(
+                                    "{} ({}). Click to change kernel.",
+                                    kernel_name,
+                                    kernel_status.to_string()
+                                ))(window, cx)
+                            })
+                            .on_click(cx.listener(|this, _, window, cx| {
+                                this.kernel_picker_handle.toggle(window, cx);
+                            }))
+                    })),
+            )
+    }
+
+    fn render_kernel_status_bar(
+        &self,
+        _window: &mut Window,
+        cx: &mut Context<Self>,
+    ) -> impl IntoElement {
+        let kernel_status = self.kernel.status();
+        let kernel_name = self
+            .kernel_specification
+            .as_ref()
+            .map(|spec| spec.name().to_string())
+            .unwrap_or_else(|| "Select Kernel".to_string());
+
+        let (status_icon, status_color) = match &kernel_status {
+            KernelStatus::Idle => (IconName::Circle, Color::Success),
+            KernelStatus::Busy => (IconName::ArrowCircle, Color::Warning),
+            KernelStatus::Starting => (IconName::ArrowCircle, Color::Muted),
+            KernelStatus::Error => (IconName::XCircle, Color::Error),
+            KernelStatus::ShuttingDown => (IconName::ArrowCircle, Color::Muted),
+            KernelStatus::Shutdown => (IconName::Circle, Color::Muted),
+            KernelStatus::Restarting => (IconName::ArrowCircle, Color::Warning),
+        };
+
+        let is_spinning = matches!(
+            kernel_status,
+            KernelStatus::Busy
+                | KernelStatus::Starting
+                | KernelStatus::ShuttingDown
+                | KernelStatus::Restarting
+        );
+
+        let status_icon_element = if is_spinning {
+            Icon::new(status_icon)
+                .size(IconSize::Small)
+                .color(status_color)
+                .with_rotate_animation(2)
+                .into_any_element()
+        } else {
+            Icon::new(status_icon)
+                .size(IconSize::Small)
+                .color(status_color)
+                .into_any_element()
+        };
+
+        let worktree_id = self.worktree_id;
+        let kernel_picker_handle = self.kernel_picker_handle.clone();
+        let view = cx.entity().downgrade();
+
+        h_flex()
+            .w_full()
+            .px_3()
+            .py_1()
+            .gap_2()
+            .items_center()
+            .justify_between()
+            .bg(cx.theme().colors().status_bar_background)
+            .child(
+                KernelSelector::new(
+                    Box::new(move |spec: KernelSpecification, window, cx| {
+                        if let Some(view) = view.upgrade() {
+                            view.update(cx, |this, cx| {
+                                this.change_kernel(spec, window, cx);
+                            });
+                        }
+                    }),
+                    worktree_id,
+                    Button::new("kernel-selector", kernel_name.clone())
+                        .label_size(LabelSize::Small)
+                        .icon(status_icon)
+                        .icon_size(IconSize::Small)
+                        .icon_color(status_color)
+                        .icon_position(IconPosition::Start),
+                    Tooltip::text(format!(
+                        "Kernel: {} ({}). Click to change.",
+                        kernel_name,
+                        kernel_status.to_string()
+                    )),
+                )
+                .with_handle(kernel_picker_handle),
+            )
+            .child(
+                h_flex()
+                    .gap_1()
+                    .child(
+                        IconButton::new("restart-kernel", IconName::RotateCw)
+                            .icon_size(IconSize::Small)
+                            .tooltip(|window, cx| {
+                                Tooltip::for_action("Restart Kernel", &RestartKernel, cx)
+                            })
+                            .on_click(cx.listener(|this, _, window, cx| {
+                                this.restart_kernel(&RestartKernel, window, cx);
+                            })),
+                    )
+                    .child(
+                        IconButton::new("interrupt-kernel", IconName::Stop)
+                            .icon_size(IconSize::Small)
+                            .disabled(!matches!(kernel_status, KernelStatus::Busy))
+                            .tooltip(|window, cx| {
+                                Tooltip::for_action("Interrupt Kernel", &InterruptKernel, cx)
+                            })
+                            .on_click(cx.listener(|this, _, window, cx| {
+                                this.interrupt_kernel(&InterruptKernel, window, cx);
+                            })),
                     ),
             )
     }
 
+    fn cell_list(&self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
+        let view = cx.entity();
+        list(self.cell_list.clone(), move |index, window, cx| {
+            view.update(cx, |this, cx| {
+                let cell_id = &this.cell_order[index];
+                let cell = this.cell_map.get(cell_id).unwrap();
+                this.render_cell(index, cell, window, cx).into_any_element()
+            })
+        })
+        .size_full()
+    }
+
     fn cell_position(&self, index: usize) -> CellPosition {
         match index {
             0 => CellPosition::First,
@@ -475,8 +1140,9 @@ impl NotebookEditor {
 
 impl Render for NotebookEditor {
     fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
-        div()
-            .key_context("notebook")
+        v_flex()
+            .size_full()
+            .key_context("NotebookEditor")
             .track_focus(&self.focus_handle)
             .on_action(cx.listener(|this, &OpenNotebook, window, cx| {
                 this.open_notebook(&OpenNotebook, window, cx)
@@ -484,6 +1150,9 @@ impl Render for NotebookEditor {
             .on_action(
                 cx.listener(|this, &ClearOutputs, window, cx| this.clear_outputs(window, cx)),
             )
+            .on_action(
+                cx.listener(|this, &Run, window, cx| this.run_current_cell(&Run, window, cx)),
+            )
             .on_action(cx.listener(|this, &RunAll, window, cx| this.run_cells(window, cx)))
             .on_action(cx.listener(|this, &MoveCellUp, window, cx| this.move_cell_up(window, cx)))
             .on_action(
@@ -495,38 +1164,22 @@ impl Render for NotebookEditor {
             .on_action(
                 cx.listener(|this, &AddCodeBlock, window, cx| this.add_code_block(window, cx)),
             )
-            .on_action(cx.listener(Self::select_next))
-            .on_action(cx.listener(Self::select_previous))
-            .on_action(cx.listener(Self::select_first))
-            .on_action(cx.listener(Self::select_last))
-            .flex()
-            .items_start()
-            .size_full()
-            .overflow_hidden()
-            .px(DynamicSpacing::Base12.px(cx))
-            .gap(DynamicSpacing::Base12.px(cx))
-            .bg(cx.theme().colors().tab_bar_background)
+            .on_action(
+                cx.listener(|this, action, window, cx| this.restart_kernel(action, window, cx)),
+            )
+            .on_action(
+                cx.listener(|this, action, window, cx| this.interrupt_kernel(action, window, cx)),
+            )
             .child(
-                v_flex()
-                    .id("notebook-cells")
+                h_flex()
                     .flex_1()
-                    .size_full()
-                    .overflow_y_scroll()
-                    .child(list(
-                        self.cell_list.clone(),
-                        cx.processor(|this, ix, window, cx| {
-                            this.cell_order
-                                .get(ix)
-                                .and_then(|cell_id| this.cell_map.get(cell_id))
-                                .map(|cell| {
-                                    this.render_cell(ix, cell, window, cx).into_any_element()
-                                })
-                                .unwrap_or_else(|| div().into_any())
-                        }),
-                    ))
-                    .size_full(),
+                    .w_full()
+                    .h_full()
+                    .gap_2()
+                    .child(div().flex_1().h_full().child(self.cell_list(window, cx)))
+                    .child(self.render_notebook_controls(window, cx)),
             )
-            .child(self.render_notebook_controls(window, cx))
+            .child(self.render_kernel_status_bar(window, cx))
     }
 }
 
@@ -566,6 +1219,18 @@ impl project::ProjectItem for NotebookItem {
 
                 // todo: watch for changes to the file
                 let file_content = fs.load(abs_path.as_path()).await?;
+
+                // Pre-process to ensure IDs exist
+                let mut json: serde_json::Value = serde_json::from_str(&file_content)?;
+                if let Some(cells) = json.get_mut("cells").and_then(|c| c.as_array_mut()) {
+                    for cell in cells {
+                        if cell.get("id").is_none() {
+                            cell["id"] = serde_json::Value::String(Uuid::new_v4().to_string());
+                        }
+                    }
+                }
+                let file_content = serde_json::to_string(&json)?;
+
                 let notebook = nbformat::parse_notebook(&file_content);
 
                 let notebook = match notebook {
@@ -611,6 +1276,7 @@ impl project::ProjectItem for NotebookItem {
     }
 
     fn is_dirty(&self) -> bool {
+        // TODO: Track if notebook metadata or structure has changed
         false
     }
 }
@@ -724,6 +1390,17 @@ impl Item for NotebookEditor {
         f(self.notebook_item.entity_id(), self.notebook_item.read(cx))
     }
 
+    fn tab_content_text(&self, _detail: usize, cx: &App) -> SharedString {
+        self.notebook_item
+            .read(cx)
+            .project_path
+            .path
+            .file_name()
+            .map(|s| s.to_string())
+            .unwrap_or_default()
+            .into()
+    }
+
     fn tab_content(&self, params: TabContentParams, window: &Window, cx: &App) -> AnyElement {
         Label::new(self.tab_content_text(params.detail.unwrap_or(0), cx))
             .single_line()
@@ -732,16 +1409,6 @@ impl Item for NotebookEditor {
             .into_any_element()
     }
 
-    fn tab_content_text(&self, _detail: usize, cx: &App) -> SharedString {
-        let path = &self.notebook_item.read(cx).path;
-        let title = path
-            .file_name()
-            .unwrap_or_else(|| path.as_os_str())
-            .to_string_lossy()
-            .to_string();
-        title.into()
-    }
-
     fn tab_icon(&self, _window: &Window, _cx: &App) -> Option<Icon> {
         Some(IconName::Book.into())
     }
@@ -769,68 +1436,154 @@ impl Item for NotebookEditor {
         // TODO
     }
 
-    // TODO
     fn can_save(&self, _cx: &App) -> bool {
-        false
+        true
     }
-    // TODO
+
     fn save(
         &mut self,
         _options: SaveOptions,
-        _project: Entity<Project>,
+        project: Entity<Project>,
         _window: &mut Window,
-        _cx: &mut Context<Self>,
+        cx: &mut Context<Self>,
     ) -> Task<Result<()>> {
-        unimplemented!("save() must be implemented if can_save() returns true")
+        let notebook = self.to_notebook(cx);
+        let path = self.notebook_item.read(cx).path.clone();
+        let fs = project.read(cx).fs().clone();
+
+        self.mark_as_saved(cx);
+
+        cx.spawn(async move |_this, _cx| {
+            let json =
+                serde_json::to_string_pretty(&notebook).context("Failed to serialize notebook")?;
+            fs.atomic_write(path, json).await?;
+            Ok(())
+        })
     }
 
-    // TODO
     fn save_as(
         &mut self,
-        _project: Entity<Project>,
-        _path: ProjectPath,
+        project: Entity<Project>,
+        path: ProjectPath,
         _window: &mut Window,
-        _cx: &mut Context<Self>,
+        cx: &mut Context<Self>,
     ) -> Task<Result<()>> {
-        unimplemented!("save_as() must be implemented if can_save() returns true")
+        let notebook = self.to_notebook(cx);
+        let fs = project.read(cx).fs().clone();
+
+        let abs_path = project.read(cx).absolute_path(&path, cx);
+
+        self.mark_as_saved(cx);
+
+        cx.spawn(async move |_this, _cx| {
+            let abs_path = abs_path.context("Failed to get absolute path")?;
+            let json =
+                serde_json::to_string_pretty(&notebook).context("Failed to serialize notebook")?;
+            fs.atomic_write(abs_path, json).await?;
+            Ok(())
+        })
     }
-    // TODO
+
     fn reload(
         &mut self,
-        _project: Entity<Project>,
-        _window: &mut Window,
-        _cx: &mut Context<Self>,
+        project: Entity<Project>,
+        window: &mut Window,
+        cx: &mut Context<Self>,
     ) -> Task<Result<()>> {
-        unimplemented!("reload() must be implemented if can_save() returns true")
-    }
+        let path = self.notebook_item.read(cx).path.clone();
+        let fs = project.read(cx).fs().clone();
+        let languages = self.languages.clone();
+        let notebook_language = self.notebook_language.clone();
 
-    fn is_dirty(&self, cx: &App) -> bool {
-        self.cell_map.values().any(|cell| {
-            if let Cell::Code(code_cell) = cell {
-                code_cell.read(cx).is_dirty(cx)
-            } else {
-                false
+        cx.spawn_in(window, async move |this, cx| {
+            let file_content = fs.load(&path).await?;
+
+            let mut json: serde_json::Value = serde_json::from_str(&file_content)?;
+            if let Some(cells) = json.get_mut("cells").and_then(|c| c.as_array_mut()) {
+                for cell in cells {
+                    if cell.get("id").is_none() {
+                        cell["id"] = serde_json::Value::String(Uuid::new_v4().to_string());
+                    }
+                }
             }
+            let file_content = serde_json::to_string(&json)?;
+
+            let notebook = nbformat::parse_notebook(&file_content);
+            let notebook = match notebook {
+                Ok(nbformat::Notebook::V4(notebook)) => notebook,
+                Ok(nbformat::Notebook::Legacy(legacy_notebook)) => {
+                    nbformat::upgrade_legacy_notebook(legacy_notebook)?
+                }
+                Err(e) => {
+                    anyhow::bail!("Failed to parse notebook: {:?}", e);
+                }
+            };
+
+            this.update_in(cx, |this, window, cx| {
+                let mut cell_order = vec![];
+                let mut cell_map = HashMap::default();
+
+                for cell in notebook.cells.iter() {
+                    let cell_id = cell.id();
+                    cell_order.push(cell_id.clone());
+                    let cell_entity =
+                        Cell::load(cell, &languages, notebook_language.clone(), window, cx);
+                    cell_map.insert(cell_id.clone(), cell_entity);
+                }
+
+                this.cell_order = cell_order.clone();
+                this.original_cell_order = cell_order;
+                this.cell_map = cell_map;
+                this.cell_list =
+                    ListState::new(this.cell_order.len(), gpui::ListAlignment::Top, px(1000.));
+                cx.notify();
+            })?;
+
+            Ok(())
         })
     }
-}
 
-// TODO: Implement this to allow us to persist to the database, etc:
-// impl SerializableItem for NotebookEditor {}
+    fn is_dirty(&self, cx: &App) -> bool {
+        self.has_structural_changes() || self.has_content_changes(cx)
+    }
+}
 
 impl ProjectItem for NotebookEditor {
     type Item = NotebookItem;
 
     fn for_project_item(
         project: Entity<Project>,
-        _: Option<&Pane>,
+        _pane: Option<&Pane>,
         item: Entity<Self::Item>,
         window: &mut Window,
         cx: &mut Context<Self>,
-    ) -> Self
-    where
-        Self: Sized,
-    {
+    ) -> Self {
         Self::new(project, item, window, cx)
     }
 }
+
+impl KernelSession for NotebookEditor {
+    fn route(&mut self, message: &JupyterMessage, window: &mut Window, cx: &mut Context<Self>) {
+        // Handle kernel status updates (these are broadcast to all)
+        if let JupyterMessageContent::Status(status) = &message.content {
+            self.kernel.set_execution_state(&status.execution_state);
+            cx.notify();
+        }
+
+        // Handle cell-specific messages
+        if let Some(parent_header) = &message.parent_header {
+            if let Some(cell_id) = self.execution_requests.get(&parent_header.msg_id) {
+                if let Some(Cell::Code(cell)) = self.cell_map.get(cell_id) {
+                    cell.update(cx, |cell, cx| {
+                        cell.handle_message(message, window, cx);
+                    });
+                }
+            }
+        }
+    }
+
+    fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>) {
+        self.kernel = Kernel::ErroredLaunch(error_message);
+        cx.notify();
+    }
+}

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);
+    }
+}