From e9aadaf0afc2563ce6b6dde8c93e626e9965bfb0 Mon Sep 17 00:00:00 2001 From: MostlyK <135974627+MostlyKIGuess@users.noreply.github.com> Date: Thu, 22 Jan 2026 20:30:29 +0530 Subject: [PATCH] repl: Add initial notebook execution + KernelSession abstraction (#43553) Coming from discussion on #25936 and https://github.com/zed-industries/zed/pull/19756 . This PR introduces the KernelSession abstraction and adds basic notebook cell execution inside NotebookEditor. The following provides a base starter for the development on Notebooks. What this PR includes today: Release Notes: - KernelSession trait extracted. Both REPL and NotebookEditor now share the same routing mechanism. - NotebookEditor can launch kernels, execute code cells, and display outputs. - Basic cell operations: run current cell, run all, move up/down, add code/markdown blocks. Keybindings follow Jupyter defaults. Next Steps: - [x] Editing support for markdown and code blocks. - Buffer integration instead of temporary cell text. - [x] Proper focus behavior when executing or adding cells. - Kernel control UI. A little far fetched steps: - Vim Support - Cell Handling Improvement and other convenient features and design from other editors - Ability to have better parsing for AI Support. I have attached a video of showcasing some of the features: https://github.com/user-attachments/assets/37e6f3e5-2022-45f0-a73d-2dd01ebc2932 --- assets/keymaps/default-linux.json | 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 | 935 +++++++++++++++++---- crates/repl/src/notebook/notebook_ui.rs | 937 +++++++++++++++++++--- crates/repl/src/outputs.rs | 37 + crates/repl/src/session.rs | 98 +-- 10 files changed, 1767 insertions(+), 307 deletions(-) diff --git a/assets/keymaps/default-linux.json b/assets/keymaps/default-linux.json index ad23d7051e92751bb6dc1b7df0dace076ee324bd..920bd24da6f2c9431bc162deb5b2f2df97bc3a28 100644 --- a/assets/keymaps/default-linux.json +++ b/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, diff --git a/assets/keymaps/default-macos.json b/assets/keymaps/default-macos.json index 8ce8102297d6664fe8ffa29d115baadf003358a9..3c4e3e010f5dcdd078eec974e7cab75e99c92780 100644 --- a/assets/keymaps/default-macos.json +++ b/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", + }, + }, ] diff --git a/assets/keymaps/default-windows.json b/assets/keymaps/default-windows.json index eed4da1c79c046c65736f649ed79b0740cf67f67..bfb6a5b7ddd7a6bc0995f060671ec87e2370c7ca 100644 --- a/assets/keymaps/default-windows.json +++ b/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" + } + }, ] diff --git a/crates/repl/src/kernels/mod.rs b/crates/repl/src/kernels/mod.rs index ab8f27121e91cca444d27fed5c5d728426166610..250628e4f3a219240cc7b68c6c2e3e445896ab40 100644 --- a/crates/repl/src/kernels/mod.rs +++ b/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); + fn kernel_errored(&mut self, error_message: String, cx: &mut Context); +} + pub type JupyterMessageChannel = stream::SelectAll>; #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/crates/repl/src/kernels/native_kernel.rs b/crates/repl/src/kernels/native_kernel.rs index 6152958925fa023efe8f4e9a0c816b2793024281..e5eb32d3499a9f7e35fc0104162be3e053d51483 100644 --- a/crates/repl/src/kernels/native_kernel.rs +++ b/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( kernel_specification: LocalKernelSpecification, entity_id: EntityId, working_directory: PathBuf, fs: Arc, // todo: convert to weak view - session: Entity, + session: Entity, window: &mut Window, cx: &mut App, ) -> Task>> { diff --git a/crates/repl/src/kernels/remote_kernels.rs b/crates/repl/src/kernels/remote_kernels.rs index 6bc8b0d1b1c8d1894b61153814dda0307d0e08f5..0be657bb8d32954cede08020b7133baf0581689a 100644 --- a/crates/repl/src/kernels/remote_kernels.rs +++ b/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( kernelspec: RemoteKernelSpecification, working_directory: std::path::PathBuf, - session: Entity, + session: Entity, window: &mut Window, cx: &mut App, ) -> Task>> { diff --git a/crates/repl/src/notebook/cell.rs b/crates/repl/src/notebook/cell.rs index 87b8e1d55ae85e09c0398848a989b7764e0d3b04..495285f86e44266d22455d56ff202e69f5f4c2fe 100644 --- a/crates/repl/src/notebook/cell.rs +++ b/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, @@ -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) -> 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) -> Option { @@ -337,11 +321,187 @@ pub struct MarkdownCell { metadata: CellMetadata, image_cache: Entity, source: String, + editor: Entity, parsed_markdown: Option, markdown_parsing_task: Task<()>, + editing: bool, selected: bool, cell_position: Option, languages: Arc, + _editor_subscription: gpui::Subscription, +} + +impl EventEmitter for MarkdownCell {} + +impl MarkdownCell { + pub fn new( + id: CellId, + metadata: CellMetadata, + source: String, + languages: Arc, + window: &mut Window, + cx: &mut Context, + ) -> 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 { + &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 = 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) { + 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) { + 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) -> 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, language_task: Task<()>, + execution_start_time: Option, + execution_duration: Option, + is_executing: bool, } +impl EventEmitter for CodeCell {} + impl CodeCell { + pub fn new( + id: CellId, + metadata: CellMetadata, + source: String, + notebook_language: Shared>>>, + window: &mut Window, + cx: &mut Context, + ) -> 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, + source: String, + outputs: Vec, + notebook_language: Shared>>>, + window: &mut Window, + cx: &mut Context, + ) -> 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 { + &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 = 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 { + 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 { + 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, + ) { + 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 { @@ -522,13 +992,22 @@ impl RenderableCell for CodeCell { } fn control(&self, window: &mut Window, cx: &mut Context) -> Option { - 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) -> 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) { println!("Running code cell: {}", self.id); + cx.emit(CellEvent::Run(self.id.clone())); } fn execution_count(&self) -> Option { @@ -569,6 +1099,16 @@ impl RunnableCell for CodeCell { impl Render for CodeCell { fn render(&mut self, window: &mut Window, cx: &mut Context) -> 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)) @@ -676,6 +1281,18 @@ pub struct RawCell { cell_position: Option, } +impl RawCell { + pub fn to_nbformat_cell(&self) -> nbformat::v4::Cell { + let source_lines: Vec = self.source.lines().map(|l| format!("{}\n", l)).collect(); + + nbformat::v4::Cell::Raw { + id: self.id.clone(), + metadata: self.metadata.clone(), + source: source_lines, + } + } +} + impl RenderableCell for RawCell { const CELL_TYPE: CellType = CellType::Raw; diff --git a/crates/repl/src/notebook/notebook_ui.rs b/crates/repl/src/notebook/notebook_ui.rs index 3dbbed46a4f5169d7fbee9a2782269e00ba41975..35a285d213e63086a9c1c2d5a6f565dd5b7a4d01 100644 --- a/crates/repl/src/notebook/notebook_ui.rs +++ b/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, project: Entity, + worktree_id: project::WorktreeId, focus_handle: FocusHandle, notebook_item: Entity, + notebook_language: Shared>>>, remote_id: Option, cell_list: ListState, selected_cell_index: usize, cell_order: Vec, + original_cell_order: Vec, cell_map: HashMap, + kernel: Kernel, + kernel_specification: Option, + execution_requests: HashMap, + kernel_picker_handle: PopoverMenuHandle>, } 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 = 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.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) { + // 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, + ) { + 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, + ) { + 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) { + 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, + ) { + 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) { + 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) { 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) { - 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) { + 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) { + 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) { @@ -171,18 +581,132 @@ impl NotebookEditor { fn move_cell_up(&mut self, window: &mut Window, cx: &mut Context) { 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) { 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) { - 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) { - 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, + ) -> 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) -> 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) -> 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 { 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: Entity, _window: &mut Window, - _cx: &mut Context, + cx: &mut Context, ) -> Task> { - 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(¬ebook).context("Failed to serialize notebook")?; + fs.atomic_write(path, json).await?; + Ok(()) + }) } - // TODO fn save_as( &mut self, - _project: Entity, - _path: ProjectPath, + project: Entity, + path: ProjectPath, _window: &mut Window, - _cx: &mut Context, + cx: &mut Context, ) -> Task> { - 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(¬ebook).context("Failed to serialize notebook")?; + fs.atomic_write(abs_path, json).await?; + Ok(()) + }) } - // TODO + fn reload( &mut self, - _project: Entity, - _window: &mut Window, - _cx: &mut Context, + project: Entity, + window: &mut Window, + cx: &mut Context, ) -> Task> { - 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, - _: Option<&Pane>, + _pane: Option<&Pane>, item: Entity, window: &mut Window, cx: &mut Context, - ) -> 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) { + // 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.kernel = Kernel::ErroredLaunch(error_message); + cx.notify(); + } +} diff --git a/crates/repl/src/outputs.rs b/crates/repl/src/outputs.rs index ae2ba8539ad8d4ca1bca43244705360aac0f8b91..8a48868ab31a20795b166ab63aae014fa8e14ac0 100644 --- a/crates/repl/src/outputs.rs +++ b/crates/repl/src/outputs.rs @@ -127,6 +127,43 @@ pub enum Output { ClearOutputWaitMarker, } +impl Output { + pub fn to_nbformat(&self, cx: &App) -> Option { + 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 = + 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: Entity, diff --git a/crates/repl/src/session.rs b/crates/repl/src/session.rs index 337b3714362914d00bc5d1ddb771d15019e1665d..ef6757df07be310913b83eda2cc9c4a3c8d0ff09 100644 --- a/crates/repl/src/session.rs +++ b/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) { - 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) { 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) { + 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.kernel_errored(error_message, cx); + } +}