notebook_ui.rs

   1#![allow(unused, dead_code)]
   2use std::future::Future;
   3use std::{path::PathBuf, sync::Arc};
   4
   5use anyhow::{Context as _, Result};
   6use client::proto::ViewId;
   7use collections::HashMap;
   8use editor::DisplayPoint;
   9use feature_flags::{FeatureFlagAppExt as _, NotebookFeatureFlag};
  10use futures::FutureExt;
  11use futures::future::Shared;
  12use gpui::{
  13    AnyElement, App, Entity, EventEmitter, FocusHandle, Focusable, ListScrollEvent, ListState,
  14    Point, Task, actions, list, prelude::*,
  15};
  16use jupyter_protocol::JupyterKernelspec;
  17use language::{Language, LanguageRegistry};
  18use log;
  19use project::{Project, ProjectEntryId, ProjectPath};
  20use settings::Settings as _;
  21use ui::{CommonAnimationExt, Tooltip, prelude::*};
  22use workspace::item::{ItemEvent, SaveOptions, TabContentParams};
  23use workspace::searchable::SearchableItemHandle;
  24use workspace::{Item, ItemHandle, Pane, ProjectItem, ToolbarItemLocation};
  25
  26use super::{Cell, CellEvent, CellPosition, MarkdownCellEvent, RenderableCell};
  27
  28use nbformat::v4::CellId;
  29use nbformat::v4::Metadata as NotebookMetadata;
  30use serde_json;
  31use uuid::Uuid;
  32
  33use crate::components::{KernelPickerDelegate, KernelSelector};
  34use crate::kernels::{
  35    Kernel, KernelSession, KernelSpecification, KernelStatus, LocalKernelSpecification,
  36    NativeRunningKernel, RemoteRunningKernel, SshRunningKernel, WslRunningKernel,
  37};
  38use crate::repl_store::ReplStore;
  39
  40use picker::Picker;
  41use runtimelib::{ExecuteRequest, JupyterMessage, JupyterMessageContent};
  42use ui::PopoverMenuHandle;
  43use zed_actions::editor::{MoveDown, MoveUp};
  44use zed_actions::notebook::{NotebookMoveDown, NotebookMoveUp};
  45
  46actions!(
  47    notebook,
  48    [
  49        /// Opens a Jupyter notebook file.
  50        OpenNotebook,
  51        /// Runs all cells in the notebook.
  52        RunAll,
  53        /// Runs the current cell.
  54        Run,
  55        /// Clears all cell outputs.
  56        ClearOutputs,
  57        /// Moves the current cell up.
  58        MoveCellUp,
  59        /// Moves the current cell down.
  60        MoveCellDown,
  61        /// Adds a new markdown cell.
  62        AddMarkdownBlock,
  63        /// Adds a new code cell.
  64        AddCodeBlock,
  65        /// Restarts the kernel.
  66        RestartKernel,
  67        /// Interrupts the current execution.
  68        InterruptKernel,
  69    ]
  70);
  71
  72pub(crate) const MAX_TEXT_BLOCK_WIDTH: f32 = 9999.0;
  73pub(crate) const SMALL_SPACING_SIZE: f32 = 8.0;
  74pub(crate) const MEDIUM_SPACING_SIZE: f32 = 12.0;
  75pub(crate) const LARGE_SPACING_SIZE: f32 = 16.0;
  76pub(crate) const GUTTER_WIDTH: f32 = 19.0;
  77pub(crate) const CODE_BLOCK_INSET: f32 = MEDIUM_SPACING_SIZE;
  78pub(crate) const CONTROL_SIZE: f32 = 20.0;
  79
  80pub fn init(cx: &mut App) {
  81    if cx.has_flag::<NotebookFeatureFlag>() || std::env::var("LOCAL_NOTEBOOK_DEV").is_ok() {
  82        workspace::register_project_item::<NotebookEditor>(cx);
  83    }
  84
  85    cx.observe_flag::<NotebookFeatureFlag, _>({
  86        move |is_enabled, cx| {
  87            if is_enabled {
  88                workspace::register_project_item::<NotebookEditor>(cx);
  89            } else {
  90                // todo: there is no way to unregister a project item, so if the feature flag
  91                // gets turned off they need to restart Zed.
  92            }
  93        }
  94    })
  95    .detach();
  96}
  97
  98pub struct NotebookEditor {
  99    languages: Arc<LanguageRegistry>,
 100    project: Entity<Project>,
 101    worktree_id: project::WorktreeId,
 102
 103    focus_handle: FocusHandle,
 104    notebook_item: Entity<NotebookItem>,
 105    notebook_language: Shared<Task<Option<Arc<Language>>>>,
 106
 107    remote_id: Option<ViewId>,
 108    cell_list: ListState,
 109
 110    selected_cell_index: usize,
 111    cell_order: Vec<CellId>,
 112    original_cell_order: Vec<CellId>,
 113    cell_map: HashMap<CellId, Cell>,
 114    kernel: Kernel,
 115    kernel_specification: Option<KernelSpecification>,
 116    execution_requests: HashMap<String, CellId>,
 117    kernel_picker_handle: PopoverMenuHandle<Picker<KernelPickerDelegate>>,
 118}
 119
 120impl NotebookEditor {
 121    pub fn new(
 122        project: Entity<Project>,
 123        notebook_item: Entity<NotebookItem>,
 124        window: &mut Window,
 125        cx: &mut Context<Self>,
 126    ) -> Self {
 127        let focus_handle = cx.focus_handle();
 128
 129        let languages = project.read(cx).languages().clone();
 130        let language_name = notebook_item.read(cx).language_name();
 131        let worktree_id = notebook_item.read(cx).project_path.worktree_id;
 132
 133        let notebook_language = notebook_item.read(cx).notebook_language();
 134        let notebook_language = cx
 135            .spawn_in(window, async move |_, _| notebook_language.await)
 136            .shared();
 137
 138        let mut cell_order = vec![]; // Vec<CellId>
 139        let mut cell_map = HashMap::default(); // HashMap<CellId, Cell>
 140
 141        let cell_count = notebook_item.read(cx).notebook.cells.len();
 142        for index in 0..cell_count {
 143            let cell = notebook_item.read(cx).notebook.cells[index].clone();
 144            let cell_id = cell.id();
 145            cell_order.push(cell_id.clone());
 146            let cell_entity = Cell::load(&cell, &languages, notebook_language.clone(), window, cx);
 147
 148            match &cell_entity {
 149                Cell::Code(code_cell) => {
 150                    let cell_id_for_focus = cell_id.clone();
 151                    cx.subscribe(code_cell, move |this, cell, event, cx| match event {
 152                        CellEvent::Run(cell_id) => this.execute_cell(cell_id.clone(), cx),
 153                        CellEvent::FocusedIn(_) => {
 154                            if let Some(index) = this
 155                                .cell_order
 156                                .iter()
 157                                .position(|id| id == &cell_id_for_focus)
 158                            {
 159                                this.selected_cell_index = index;
 160                                cx.notify();
 161                            }
 162                        }
 163                    })
 164                    .detach();
 165
 166                    let cell_id_for_editor = cell_id.clone();
 167                    let editor = code_cell.read(cx).editor().clone();
 168                    cx.subscribe(&editor, move |this, _editor, event, cx| {
 169                        if let editor::EditorEvent::Focused = event {
 170                            if let Some(index) = this
 171                                .cell_order
 172                                .iter()
 173                                .position(|id| id == &cell_id_for_editor)
 174                            {
 175                                this.selected_cell_index = index;
 176                                cx.notify();
 177                            }
 178                        }
 179                    })
 180                    .detach();
 181                }
 182                Cell::Markdown(markdown_cell) => {
 183                    let cell_id_for_focus = cell_id.clone();
 184                    cx.subscribe(
 185                        markdown_cell,
 186                        move |_this, cell, event: &MarkdownCellEvent, cx| {
 187                            match event {
 188                                MarkdownCellEvent::FinishedEditing => {
 189                                    cell.update(cx, |cell, cx| {
 190                                        cell.reparse_markdown(cx);
 191                                    });
 192                                }
 193                                MarkdownCellEvent::Run(_cell_id) => {
 194                                    // run is handled separately by move_to_next_cell
 195                                    // Just reparse here
 196                                    cell.update(cx, |cell, cx| {
 197                                        cell.reparse_markdown(cx);
 198                                    });
 199                                }
 200                            }
 201                        },
 202                    )
 203                    .detach();
 204
 205                    let cell_id_for_editor = cell_id.clone();
 206                    let editor = markdown_cell.read(cx).editor().clone();
 207                    cx.subscribe(&editor, move |this, _editor, event, cx| {
 208                        if let editor::EditorEvent::Focused = event {
 209                            if let Some(index) = this
 210                                .cell_order
 211                                .iter()
 212                                .position(|id| id == &cell_id_for_editor)
 213                            {
 214                                this.selected_cell_index = index;
 215                                cx.notify();
 216                            }
 217                        }
 218                    })
 219                    .detach();
 220                }
 221                Cell::Raw(_) => {}
 222            }
 223
 224            cell_map.insert(cell_id.clone(), cell_entity);
 225        }
 226
 227        let notebook_handle = cx.entity().downgrade();
 228        let cell_count = cell_order.len();
 229
 230        let this = cx.entity();
 231        let cell_list = ListState::new(cell_count, gpui::ListAlignment::Top, px(1000.));
 232
 233        let mut editor = Self {
 234            project,
 235            languages: languages.clone(),
 236            worktree_id,
 237            focus_handle,
 238            notebook_item: notebook_item.clone(),
 239            notebook_language,
 240            remote_id: None,
 241            cell_list,
 242            selected_cell_index: 0,
 243            cell_order: cell_order.clone(),
 244            original_cell_order: cell_order.clone(),
 245            cell_map: cell_map.clone(),
 246            kernel: Kernel::Shutdown, // TODO: use recommended kernel after the implementation is done in repl
 247            kernel_specification: None,
 248            execution_requests: HashMap::default(),
 249            kernel_picker_handle: PopoverMenuHandle::default(),
 250        };
 251        editor.launch_kernel(window, cx);
 252        editor.refresh_language(cx);
 253
 254        cx.subscribe(&notebook_item, |this, _item, _event, cx| {
 255            this.refresh_language(cx);
 256        })
 257        .detach();
 258
 259        editor
 260    }
 261
 262    fn refresh_language(&mut self, cx: &mut Context<Self>) {
 263        let notebook_language = self.notebook_item.read(cx).notebook_language();
 264        let task = cx.spawn(async move |this, cx| {
 265            let language = notebook_language.await;
 266            if let Some(this) = this.upgrade() {
 267                this.update(cx, |this, cx| {
 268                    for cell in this.cell_map.values() {
 269                        if let Cell::Code(code_cell) = cell {
 270                            code_cell.update(cx, |cell, cx| {
 271                                cell.set_language(language.clone(), cx);
 272                            });
 273                        }
 274                    }
 275                });
 276            }
 277            language
 278        });
 279        self.notebook_language = task.shared();
 280    }
 281
 282    fn has_structural_changes(&self) -> bool {
 283        self.cell_order != self.original_cell_order
 284    }
 285
 286    fn has_content_changes(&self, cx: &App) -> bool {
 287        self.cell_map.values().any(|cell| cell.is_dirty(cx))
 288    }
 289
 290    pub fn to_notebook(&self, cx: &App) -> nbformat::v4::Notebook {
 291        let cells: Vec<nbformat::v4::Cell> = self
 292            .cell_order
 293            .iter()
 294            .filter_map(|cell_id| {
 295                self.cell_map
 296                    .get(cell_id)
 297                    .map(|cell| cell.to_nbformat_cell(cx))
 298            })
 299            .collect();
 300
 301        let metadata = self.notebook_item.read(cx).notebook.metadata.clone();
 302
 303        nbformat::v4::Notebook {
 304            metadata,
 305            nbformat: 4,
 306            nbformat_minor: 5,
 307            cells,
 308        }
 309    }
 310
 311    pub fn mark_as_saved(&mut self, cx: &mut Context<Self>) {
 312        self.original_cell_order = self.cell_order.clone();
 313
 314        for cell in self.cell_map.values() {
 315            match cell {
 316                Cell::Code(code_cell) => {
 317                    code_cell.update(cx, |code_cell, cx| {
 318                        let editor = code_cell.editor();
 319                        editor.update(cx, |editor, cx| {
 320                            editor.buffer().update(cx, |buffer, cx| {
 321                                if let Some(buf) = buffer.as_singleton() {
 322                                    buf.update(cx, |b, cx| {
 323                                        let version = b.version();
 324                                        b.did_save(version, None, cx);
 325                                    });
 326                                }
 327                            });
 328                        });
 329                    });
 330                }
 331                Cell::Markdown(markdown_cell) => {
 332                    markdown_cell.update(cx, |markdown_cell, cx| {
 333                        let editor = markdown_cell.editor();
 334                        editor.update(cx, |editor, cx| {
 335                            editor.buffer().update(cx, |buffer, cx| {
 336                                if let Some(buf) = buffer.as_singleton() {
 337                                    buf.update(cx, |b, cx| {
 338                                        let version = b.version();
 339                                        b.did_save(version, None, cx);
 340                                    });
 341                                }
 342                            });
 343                        });
 344                    });
 345                }
 346                Cell::Raw(_) => {}
 347            }
 348        }
 349        cx.notify();
 350    }
 351
 352    fn launch_kernel(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 353        // use default Python kernel if no specification is set
 354        let spec = self.kernel_specification.clone().unwrap_or_else(|| {
 355            KernelSpecification::Jupyter(LocalKernelSpecification {
 356                name: "python3".to_string(),
 357                path: PathBuf::from("python3"),
 358                kernelspec: JupyterKernelspec {
 359                    argv: vec![
 360                        "python3".to_string(),
 361                        "-m".to_string(),
 362                        "ipykernel_launcher".to_string(),
 363                        "-f".to_string(),
 364                        "{connection_file}".to_string(),
 365                    ],
 366                    display_name: "Python 3".to_string(),
 367                    language: "python".to_string(),
 368                    interrupt_mode: None,
 369                    metadata: None,
 370                    env: None,
 371                },
 372            })
 373        });
 374
 375        self.launch_kernel_with_spec(spec, window, cx);
 376    }
 377
 378    fn launch_kernel_with_spec(
 379        &mut self,
 380        spec: KernelSpecification,
 381        window: &mut Window,
 382        cx: &mut Context<Self>,
 383    ) {
 384        let entity_id = cx.entity_id();
 385        let working_directory = self
 386            .project
 387            .read(cx)
 388            .worktrees(cx)
 389            .next()
 390            .map(|worktree| worktree.read(cx).abs_path().to_path_buf())
 391            .unwrap_or_else(std::env::temp_dir);
 392        let fs = self.project.read(cx).fs().clone();
 393        let view = cx.entity();
 394
 395        self.kernel_specification = Some(spec.clone());
 396
 397        self.notebook_item.update(cx, |item, cx| {
 398            let kernel_name = spec.name().to_string();
 399            let language = spec.language().to_string();
 400
 401            let display_name = match &spec {
 402                KernelSpecification::Jupyter(s) => s.kernelspec.display_name.clone(),
 403                KernelSpecification::PythonEnv(s) => s.kernelspec.display_name.clone(),
 404                KernelSpecification::JupyterServer(s) => s.kernelspec.display_name.clone(),
 405                KernelSpecification::SshRemote(s) => s.kernelspec.display_name.clone(),
 406                KernelSpecification::WslRemote(s) => s.kernelspec.display_name.clone(),
 407            };
 408
 409            let kernelspec_json = serde_json::json!({
 410                "display_name": display_name,
 411                "name": kernel_name,
 412                "language": language
 413            });
 414
 415            if let Ok(k) = serde_json::from_value(kernelspec_json) {
 416                item.notebook.metadata.kernelspec = Some(k);
 417                cx.emit(());
 418            }
 419        });
 420
 421        let kernel_task = match spec {
 422            KernelSpecification::Jupyter(local_spec) => NativeRunningKernel::new(
 423                local_spec,
 424                entity_id,
 425                working_directory,
 426                fs,
 427                view,
 428                window,
 429                cx,
 430            ),
 431            KernelSpecification::PythonEnv(env_spec) => NativeRunningKernel::new(
 432                env_spec.as_local_spec(),
 433                entity_id,
 434                working_directory,
 435                fs,
 436                view,
 437                window,
 438                cx,
 439            ),
 440            KernelSpecification::JupyterServer(remote_spec) => {
 441                RemoteRunningKernel::new(remote_spec, working_directory, view, window, cx)
 442            }
 443
 444            KernelSpecification::SshRemote(spec) => {
 445                let project = self.project.clone();
 446                SshRunningKernel::new(spec, working_directory, project, view, window, cx)
 447            }
 448            KernelSpecification::WslRemote(spec) => {
 449                WslRunningKernel::new(spec, entity_id, working_directory, fs, view, window, cx)
 450            }
 451        };
 452
 453        let pending_kernel = cx
 454            .spawn(async move |this, cx| {
 455                let kernel = kernel_task.await;
 456
 457                match kernel {
 458                    Ok(kernel) => {
 459                        this.update(cx, |editor, cx| {
 460                            editor.kernel = Kernel::RunningKernel(kernel);
 461                            cx.notify();
 462                        })
 463                        .ok();
 464                    }
 465                    Err(err) => {
 466                        log::error!("Kernel failed to start: {:?}", err);
 467                        this.update(cx, |editor, cx| {
 468                            editor.kernel = Kernel::ErroredLaunch(err.to_string());
 469                            cx.notify();
 470                        })
 471                        .ok();
 472                    }
 473                }
 474            })
 475            .shared();
 476
 477        self.kernel = Kernel::StartingKernel(pending_kernel);
 478        cx.notify();
 479    }
 480
 481    // Note: Python environments are only detected as kernels if ipykernel is installed.
 482    // Users need to run `pip install ipykernel` (or `uv pip install ipykernel`) in their
 483    // virtual environment for it to appear in the kernel selector.
 484    // This happens because we have an ipykernel check inside the function python_env_kernel_specification in mod.rs L:121
 485
 486    fn change_kernel(
 487        &mut self,
 488        spec: KernelSpecification,
 489        window: &mut Window,
 490        cx: &mut Context<Self>,
 491    ) {
 492        if let Kernel::RunningKernel(kernel) = &mut self.kernel {
 493            kernel.force_shutdown(window, cx).detach();
 494        }
 495
 496        self.execution_requests.clear();
 497
 498        self.launch_kernel_with_spec(spec, window, cx);
 499    }
 500
 501    fn restart_kernel(&mut self, _: &RestartKernel, window: &mut Window, cx: &mut Context<Self>) {
 502        if let Some(spec) = self.kernel_specification.clone() {
 503            if let Kernel::RunningKernel(kernel) = &mut self.kernel {
 504                kernel.force_shutdown(window, cx).detach();
 505            }
 506
 507            self.kernel = Kernel::Restarting;
 508            cx.notify();
 509
 510            self.launch_kernel_with_spec(spec, window, cx);
 511        }
 512    }
 513
 514    fn interrupt_kernel(
 515        &mut self,
 516        _: &InterruptKernel,
 517        _window: &mut Window,
 518        cx: &mut Context<Self>,
 519    ) {
 520        if let Kernel::RunningKernel(kernel) = &self.kernel {
 521            let interrupt_request = runtimelib::InterruptRequest {};
 522            let message: JupyterMessage = interrupt_request.into();
 523            kernel.request_tx().try_send(message).ok();
 524            cx.notify();
 525        }
 526    }
 527
 528    fn execute_cell(&mut self, cell_id: CellId, cx: &mut Context<Self>) {
 529        let code = if let Some(Cell::Code(cell)) = self.cell_map.get(&cell_id) {
 530            let editor = cell.read(cx).editor().clone();
 531            let buffer = editor.read(cx).buffer().read(cx);
 532            buffer
 533                .as_singleton()
 534                .map(|b| b.read(cx).text())
 535                .unwrap_or_default()
 536        } else {
 537            return;
 538        };
 539
 540        if let Some(Cell::Code(cell)) = self.cell_map.get(&cell_id) {
 541            cell.update(cx, |cell, cx| {
 542                if cell.has_outputs() {
 543                    cell.clear_outputs();
 544                }
 545                cell.start_execution();
 546                cx.notify();
 547            });
 548        }
 549
 550        let request = ExecuteRequest {
 551            code,
 552            ..Default::default()
 553        };
 554        let message: JupyterMessage = request.into();
 555        let msg_id = message.header.msg_id.clone();
 556
 557        self.execution_requests.insert(msg_id, cell_id.clone());
 558
 559        if let Kernel::RunningKernel(kernel) = &mut self.kernel {
 560            kernel.request_tx().try_send(message).ok();
 561        }
 562    }
 563
 564    fn has_outputs(&self, window: &mut Window, cx: &mut Context<Self>) -> bool {
 565        self.cell_map.values().any(|cell| {
 566            if let Cell::Code(code_cell) = cell {
 567                code_cell.read(cx).has_outputs()
 568            } else {
 569                false
 570            }
 571        })
 572    }
 573
 574    fn clear_outputs(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 575        for cell in self.cell_map.values() {
 576            if let Cell::Code(code_cell) = cell {
 577                code_cell.update(cx, |cell, cx| {
 578                    cell.clear_outputs();
 579                    cx.notify();
 580                });
 581            }
 582        }
 583        cx.notify();
 584    }
 585
 586    fn run_cells(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 587        for cell_id in self.cell_order.clone() {
 588            self.execute_cell(cell_id, cx);
 589        }
 590    }
 591
 592    fn run_current_cell(&mut self, _: &Run, window: &mut Window, cx: &mut Context<Self>) {
 593        if let Some(cell_id) = self.cell_order.get(self.selected_cell_index).cloned() {
 594            if let Some(cell) = self.cell_map.get(&cell_id) {
 595                match cell {
 596                    Cell::Code(_) => {
 597                        self.execute_cell(cell_id, cx);
 598                    }
 599                    Cell::Markdown(markdown_cell) => {
 600                        // for markdown, finish editing and move to next cell
 601                        let is_editing = markdown_cell.read(cx).is_editing();
 602                        if is_editing {
 603                            markdown_cell.update(cx, |cell, cx| {
 604                                cell.run(cx);
 605                            });
 606                            // move to the next cell
 607                            // Discussion can be done on this default implementation
 608                            self.move_to_next_cell(window, cx);
 609                        }
 610                    }
 611                    Cell::Raw(_) => {}
 612                }
 613            }
 614        }
 615    }
 616
 617    // Discussion can be done on this default implementation
 618    /// Moves focus to the next cell, or creates a new code cell if at the end
 619    fn move_to_next_cell(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 620        if !self.cell_order.is_empty() && self.selected_cell_index < self.cell_order.len() - 1 {
 621            self.selected_cell_index += 1;
 622            // focus the new cell's editor
 623            if let Some(cell_id) = self.cell_order.get(self.selected_cell_index) {
 624                if let Some(cell) = self.cell_map.get(cell_id) {
 625                    match cell {
 626                        Cell::Code(code_cell) => {
 627                            let editor = code_cell.read(cx).editor();
 628                            window.focus(&editor.focus_handle(cx), cx);
 629                        }
 630                        Cell::Markdown(markdown_cell) => {
 631                            // Don't auto-enter edit mode for next markdown cell
 632                            // Just select it
 633                        }
 634                        Cell::Raw(_) => {}
 635                    }
 636                }
 637            }
 638            cx.notify();
 639        } else {
 640            // in the end, could optionally create a new cell
 641            // For now, just stay on the current cell
 642        }
 643    }
 644
 645    fn open_notebook(&mut self, _: &OpenNotebook, _window: &mut Window, _cx: &mut Context<Self>) {
 646        println!("Open notebook triggered");
 647    }
 648
 649    fn move_cell_up(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 650        println!("Move cell up triggered");
 651        if self.selected_cell_index > 0 {
 652            self.cell_order
 653                .swap(self.selected_cell_index, self.selected_cell_index - 1);
 654            self.selected_cell_index -= 1;
 655            cx.notify();
 656        }
 657    }
 658
 659    fn move_cell_down(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 660        println!("Move cell down triggered");
 661        if !self.cell_order.is_empty() && self.selected_cell_index < self.cell_order.len() - 1 {
 662            self.cell_order
 663                .swap(self.selected_cell_index, self.selected_cell_index + 1);
 664            self.selected_cell_index += 1;
 665            cx.notify();
 666        }
 667    }
 668
 669    fn add_markdown_block(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 670        let new_cell_id: CellId = Uuid::new_v4().into();
 671        let languages = self.languages.clone();
 672        let metadata: nbformat::v4::CellMetadata =
 673            serde_json::from_str("{}").expect("empty object should parse");
 674
 675        let markdown_cell = cx.new(|cx| {
 676            super::MarkdownCell::new(
 677                new_cell_id.clone(),
 678                metadata,
 679                String::new(),
 680                languages,
 681                window,
 682                cx,
 683            )
 684        });
 685
 686        let insert_index = if self.cell_order.is_empty() {
 687            0
 688        } else {
 689            self.selected_cell_index + 1
 690        };
 691        self.cell_order.insert(insert_index, new_cell_id.clone());
 692        self.cell_map
 693            .insert(new_cell_id.clone(), Cell::Markdown(markdown_cell.clone()));
 694        self.selected_cell_index = insert_index;
 695
 696        cx.subscribe(
 697            &markdown_cell,
 698            move |_this, cell, event: &MarkdownCellEvent, cx| match event {
 699                MarkdownCellEvent::FinishedEditing | MarkdownCellEvent::Run(_) => {
 700                    cell.update(cx, |cell, cx| {
 701                        cell.reparse_markdown(cx);
 702                    });
 703                }
 704            },
 705        )
 706        .detach();
 707
 708        let cell_id_for_editor = new_cell_id.clone();
 709        let editor = markdown_cell.read(cx).editor().clone();
 710        cx.subscribe(&editor, move |this, _editor, event, cx| {
 711            if let editor::EditorEvent::Focused = event {
 712                if let Some(index) = this
 713                    .cell_order
 714                    .iter()
 715                    .position(|id| id == &cell_id_for_editor)
 716                {
 717                    this.selected_cell_index = index;
 718                    cx.notify();
 719                }
 720            }
 721        })
 722        .detach();
 723
 724        self.cell_list.reset(self.cell_order.len());
 725        cx.notify();
 726    }
 727
 728    fn add_code_block(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 729        let new_cell_id: CellId = Uuid::new_v4().into();
 730        let notebook_language = self.notebook_language.clone();
 731        let metadata: nbformat::v4::CellMetadata =
 732            serde_json::from_str("{}").expect("empty object should parse");
 733
 734        let code_cell = cx.new(|cx| {
 735            super::CodeCell::new(
 736                new_cell_id.clone(),
 737                metadata,
 738                String::new(),
 739                notebook_language,
 740                window,
 741                cx,
 742            )
 743        });
 744
 745        let insert_index = if self.cell_order.is_empty() {
 746            0
 747        } else {
 748            self.selected_cell_index + 1
 749        };
 750        self.cell_order.insert(insert_index, new_cell_id.clone());
 751        self.cell_map
 752            .insert(new_cell_id.clone(), Cell::Code(code_cell.clone()));
 753        self.selected_cell_index = insert_index;
 754
 755        let cell_id_for_run = new_cell_id.clone();
 756        cx.subscribe(&code_cell, move |this, _cell, event, cx| match event {
 757            CellEvent::Run(cell_id) => this.execute_cell(cell_id.clone(), cx),
 758            CellEvent::FocusedIn(_) => {
 759                if let Some(index) = this.cell_order.iter().position(|id| id == &cell_id_for_run) {
 760                    this.selected_cell_index = index;
 761                    cx.notify();
 762                }
 763            }
 764        })
 765        .detach();
 766
 767        let cell_id_for_editor = new_cell_id.clone();
 768        let editor = code_cell.read(cx).editor().clone();
 769        cx.subscribe(&editor, move |this, _editor, event, cx| {
 770            if let editor::EditorEvent::Focused = event {
 771                if let Some(index) = this
 772                    .cell_order
 773                    .iter()
 774                    .position(|id| id == &cell_id_for_editor)
 775                {
 776                    this.selected_cell_index = index;
 777                    cx.notify();
 778                }
 779            }
 780        })
 781        .detach();
 782
 783        self.cell_list.reset(self.cell_order.len());
 784        cx.notify();
 785    }
 786
 787    fn cell_count(&self) -> usize {
 788        self.cell_map.len()
 789    }
 790
 791    fn selected_index(&self) -> usize {
 792        self.selected_cell_index
 793    }
 794
 795    pub fn set_selected_index(
 796        &mut self,
 797        index: usize,
 798        jump_to_index: bool,
 799        window: &mut Window,
 800        cx: &mut Context<Self>,
 801    ) {
 802        // let previous_index = self.selected_cell_index;
 803        self.selected_cell_index = index;
 804        let current_index = self.selected_cell_index;
 805
 806        // in the future we may have some `on_cell_change` event that we want to fire here
 807
 808        if jump_to_index {
 809            self.jump_to_cell(current_index, window, cx);
 810        }
 811    }
 812
 813    pub fn select_next(
 814        &mut self,
 815        _: &menu::SelectNext,
 816        window: &mut Window,
 817        cx: &mut Context<Self>,
 818    ) {
 819        let count = self.cell_count();
 820        if count > 0 {
 821            let index = self.selected_index();
 822            let ix = if index == count - 1 {
 823                count - 1
 824            } else {
 825                index + 1
 826            };
 827            self.set_selected_index(ix, true, window, cx);
 828            cx.notify();
 829        }
 830    }
 831
 832    pub fn select_previous(
 833        &mut self,
 834        _: &menu::SelectPrevious,
 835        window: &mut Window,
 836        cx: &mut Context<Self>,
 837    ) {
 838        let count = self.cell_count();
 839        if count > 0 {
 840            let index = self.selected_index();
 841            let ix = if index == 0 { 0 } else { index - 1 };
 842            self.set_selected_index(ix, true, window, cx);
 843            cx.notify();
 844        }
 845    }
 846
 847    pub fn select_first(
 848        &mut self,
 849        _: &menu::SelectFirst,
 850        window: &mut Window,
 851        cx: &mut Context<Self>,
 852    ) {
 853        let count = self.cell_count();
 854        if count > 0 {
 855            self.set_selected_index(0, true, window, cx);
 856            cx.notify();
 857        }
 858    }
 859
 860    pub fn select_last(
 861        &mut self,
 862        _: &menu::SelectLast,
 863        window: &mut Window,
 864        cx: &mut Context<Self>,
 865    ) {
 866        let count = self.cell_count();
 867        if count > 0 {
 868            self.set_selected_index(count - 1, true, window, cx);
 869            cx.notify();
 870        }
 871    }
 872
 873    fn jump_to_cell(&mut self, index: usize, _window: &mut Window, _cx: &mut Context<Self>) {
 874        self.cell_list.scroll_to_reveal_item(index);
 875    }
 876
 877    fn button_group(window: &mut Window, cx: &mut Context<Self>) -> Div {
 878        v_flex()
 879            .gap(DynamicSpacing::Base04.rems(cx))
 880            .items_center()
 881            .w(px(CONTROL_SIZE + 4.0))
 882            .overflow_hidden()
 883            .rounded(px(5.))
 884            .bg(cx.theme().colors().title_bar_background)
 885            .p_px()
 886            .border_1()
 887            .border_color(cx.theme().colors().border)
 888    }
 889
 890    fn render_notebook_control(
 891        id: impl Into<SharedString>,
 892        icon: IconName,
 893        _window: &mut Window,
 894        _cx: &mut Context<Self>,
 895    ) -> IconButton {
 896        let id: ElementId = ElementId::Name(id.into());
 897        IconButton::new(id, icon).width(px(CONTROL_SIZE))
 898    }
 899
 900    fn render_notebook_controls(
 901        &self,
 902        window: &mut Window,
 903        cx: &mut Context<Self>,
 904    ) -> impl IntoElement {
 905        let has_outputs = self.has_outputs(window, cx);
 906
 907        v_flex()
 908            .max_w(px(CONTROL_SIZE + 4.0))
 909            .items_center()
 910            .gap(DynamicSpacing::Base16.rems(cx))
 911            .justify_between()
 912            .flex_none()
 913            .h_full()
 914            .py(DynamicSpacing::Base12.px(cx))
 915            .child(
 916                v_flex()
 917                    .gap(DynamicSpacing::Base08.rems(cx))
 918                    .child(
 919                        Self::button_group(window, cx)
 920                            .child(
 921                                Self::render_notebook_control(
 922                                    "run-all-cells",
 923                                    IconName::PlayFilled,
 924                                    window,
 925                                    cx,
 926                                )
 927                                .tooltip(move |window, cx| {
 928                                    Tooltip::for_action("Execute all cells", &RunAll, cx)
 929                                })
 930                                .on_click(|_, window, cx| {
 931                                    window.dispatch_action(Box::new(RunAll), cx);
 932                                }),
 933                            )
 934                            .child(
 935                                Self::render_notebook_control(
 936                                    "clear-all-outputs",
 937                                    IconName::ListX,
 938                                    window,
 939                                    cx,
 940                                )
 941                                .disabled(!has_outputs)
 942                                .tooltip(move |window, cx| {
 943                                    Tooltip::for_action("Clear all outputs", &ClearOutputs, cx)
 944                                })
 945                                .on_click(|_, window, cx| {
 946                                    window.dispatch_action(Box::new(ClearOutputs), cx);
 947                                }),
 948                            ),
 949                    )
 950                    .child(
 951                        Self::button_group(window, cx)
 952                            .child(
 953                                Self::render_notebook_control(
 954                                    "move-cell-up",
 955                                    IconName::ArrowUp,
 956                                    window,
 957                                    cx,
 958                                )
 959                                .tooltip(move |window, cx| {
 960                                    Tooltip::for_action("Move cell up", &MoveCellUp, cx)
 961                                })
 962                                .on_click(|_, window, cx| {
 963                                    window.dispatch_action(Box::new(MoveCellUp), cx);
 964                                }),
 965                            )
 966                            .child(
 967                                Self::render_notebook_control(
 968                                    "move-cell-down",
 969                                    IconName::ArrowDown,
 970                                    window,
 971                                    cx,
 972                                )
 973                                .tooltip(move |window, cx| {
 974                                    Tooltip::for_action("Move cell down", &MoveCellDown, cx)
 975                                })
 976                                .on_click(|_, window, cx| {
 977                                    window.dispatch_action(Box::new(MoveCellDown), cx);
 978                                }),
 979                            ),
 980                    )
 981                    .child(
 982                        Self::button_group(window, cx)
 983                            .child(
 984                                Self::render_notebook_control(
 985                                    "new-markdown-cell",
 986                                    IconName::Plus,
 987                                    window,
 988                                    cx,
 989                                )
 990                                .tooltip(move |window, cx| {
 991                                    Tooltip::for_action("Add markdown block", &AddMarkdownBlock, cx)
 992                                })
 993                                .on_click(|_, window, cx| {
 994                                    window.dispatch_action(Box::new(AddMarkdownBlock), cx);
 995                                }),
 996                            )
 997                            .child(
 998                                Self::render_notebook_control(
 999                                    "new-code-cell",
1000                                    IconName::Code,
1001                                    window,
1002                                    cx,
1003                                )
1004                                .tooltip(move |window, cx| {
1005                                    Tooltip::for_action("Add code block", &AddCodeBlock, cx)
1006                                })
1007                                .on_click(|_, window, cx| {
1008                                    window.dispatch_action(Box::new(AddCodeBlock), cx);
1009                                }),
1010                            ),
1011                    ),
1012            )
1013            .child(
1014                v_flex()
1015                    .gap(DynamicSpacing::Base08.rems(cx))
1016                    .items_center()
1017                    .child(
1018                        Self::render_notebook_control("more-menu", IconName::Ellipsis, window, cx)
1019                            .tooltip(move |window, cx| (Tooltip::text("More options"))(window, cx)),
1020                    )
1021                    .child(Self::button_group(window, cx).child({
1022                        let kernel_status = self.kernel.status();
1023                        let (icon, icon_color) = match &kernel_status {
1024                            KernelStatus::Idle => (IconName::ReplNeutral, Color::Success),
1025                            KernelStatus::Busy => (IconName::ReplNeutral, Color::Warning),
1026                            KernelStatus::Starting => (IconName::ReplNeutral, Color::Muted),
1027                            KernelStatus::Error => (IconName::ReplNeutral, Color::Error),
1028                            KernelStatus::ShuttingDown => (IconName::ReplNeutral, Color::Muted),
1029                            KernelStatus::Shutdown => (IconName::ReplNeutral, Color::Disabled),
1030                            KernelStatus::Restarting => (IconName::ReplNeutral, Color::Warning),
1031                        };
1032                        let kernel_name = self
1033                            .kernel_specification
1034                            .as_ref()
1035                            .map(|spec| spec.name().to_string())
1036                            .unwrap_or_else(|| "Select Kernel".to_string());
1037                        IconButton::new("repl", icon)
1038                            .icon_color(icon_color)
1039                            .tooltip(move |window, cx| {
1040                                Tooltip::text(format!(
1041                                    "{} ({}). Click to change kernel.",
1042                                    kernel_name,
1043                                    kernel_status.to_string()
1044                                ))(window, cx)
1045                            })
1046                            .on_click(cx.listener(|this, _, window, cx| {
1047                                this.kernel_picker_handle.toggle(window, cx);
1048                            }))
1049                    })),
1050            )
1051    }
1052
1053    fn render_kernel_status_bar(
1054        &self,
1055        _window: &mut Window,
1056        cx: &mut Context<Self>,
1057    ) -> impl IntoElement {
1058        let kernel_status = self.kernel.status();
1059        let kernel_name = self
1060            .kernel_specification
1061            .as_ref()
1062            .map(|spec| spec.name().to_string())
1063            .unwrap_or_else(|| "Select Kernel".to_string());
1064
1065        let (status_icon, status_color) = match &kernel_status {
1066            KernelStatus::Idle => (IconName::Circle, Color::Success),
1067            KernelStatus::Busy => (IconName::ArrowCircle, Color::Warning),
1068            KernelStatus::Starting => (IconName::ArrowCircle, Color::Muted),
1069            KernelStatus::Error => (IconName::XCircle, Color::Error),
1070            KernelStatus::ShuttingDown => (IconName::ArrowCircle, Color::Muted),
1071            KernelStatus::Shutdown => (IconName::Circle, Color::Muted),
1072            KernelStatus::Restarting => (IconName::ArrowCircle, Color::Warning),
1073        };
1074
1075        let is_spinning = matches!(
1076            kernel_status,
1077            KernelStatus::Busy
1078                | KernelStatus::Starting
1079                | KernelStatus::ShuttingDown
1080                | KernelStatus::Restarting
1081        );
1082
1083        let status_icon_element = if is_spinning {
1084            Icon::new(status_icon)
1085                .size(IconSize::Small)
1086                .color(status_color)
1087                .with_rotate_animation(2)
1088                .into_any_element()
1089        } else {
1090            Icon::new(status_icon)
1091                .size(IconSize::Small)
1092                .color(status_color)
1093                .into_any_element()
1094        };
1095
1096        let worktree_id = self.worktree_id;
1097        let kernel_picker_handle = self.kernel_picker_handle.clone();
1098        let view = cx.entity().downgrade();
1099
1100        h_flex()
1101            .w_full()
1102            .px_3()
1103            .py_1()
1104            .gap_2()
1105            .items_center()
1106            .justify_between()
1107            .bg(cx.theme().colors().status_bar_background)
1108            .child(
1109                KernelSelector::new(
1110                    Box::new(move |spec: KernelSpecification, window, cx| {
1111                        if let Some(view) = view.upgrade() {
1112                            view.update(cx, |this, cx| {
1113                                this.change_kernel(spec, window, cx);
1114                            });
1115                        }
1116                    }),
1117                    worktree_id,
1118                    Button::new("kernel-selector", kernel_name.clone())
1119                        .label_size(LabelSize::Small)
1120                        .start_icon(
1121                            Icon::new(status_icon)
1122                                .size(IconSize::Small)
1123                                .color(status_color),
1124                        ),
1125                    Tooltip::text(format!(
1126                        "Kernel: {} ({}). Click to change.",
1127                        kernel_name,
1128                        kernel_status.to_string()
1129                    )),
1130                )
1131                .with_handle(kernel_picker_handle),
1132            )
1133            .child(
1134                h_flex()
1135                    .gap_1()
1136                    .child(
1137                        IconButton::new("restart-kernel", IconName::RotateCw)
1138                            .icon_size(IconSize::Small)
1139                            .tooltip(|window, cx| {
1140                                Tooltip::for_action("Restart Kernel", &RestartKernel, cx)
1141                            })
1142                            .on_click(cx.listener(|this, _, window, cx| {
1143                                this.restart_kernel(&RestartKernel, window, cx);
1144                            })),
1145                    )
1146                    .child(
1147                        IconButton::new("interrupt-kernel", IconName::Stop)
1148                            .icon_size(IconSize::Small)
1149                            .disabled(!matches!(kernel_status, KernelStatus::Busy))
1150                            .tooltip(|window, cx| {
1151                                Tooltip::for_action("Interrupt Kernel", &InterruptKernel, cx)
1152                            })
1153                            .on_click(cx.listener(|this, _, window, cx| {
1154                                this.interrupt_kernel(&InterruptKernel, window, cx);
1155                            })),
1156                    ),
1157            )
1158    }
1159
1160    fn cell_list(&self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1161        let view = cx.entity();
1162        list(self.cell_list.clone(), move |index, window, cx| {
1163            view.update(cx, |this, cx| {
1164                let cell_id = &this.cell_order[index];
1165                let cell = this.cell_map.get(cell_id).unwrap();
1166                this.render_cell(index, cell, window, cx).into_any_element()
1167            })
1168        })
1169        .size_full()
1170    }
1171
1172    fn cell_position(&self, index: usize) -> CellPosition {
1173        match index {
1174            0 => CellPosition::First,
1175            index if index == self.cell_count() - 1 => CellPosition::Last,
1176            _ => CellPosition::Middle,
1177        }
1178    }
1179
1180    fn render_cell(
1181        &self,
1182        index: usize,
1183        cell: &Cell,
1184        window: &mut Window,
1185        cx: &mut Context<Self>,
1186    ) -> impl IntoElement {
1187        let cell_position = self.cell_position(index);
1188
1189        let is_selected = index == self.selected_cell_index;
1190
1191        match cell {
1192            Cell::Code(cell) => {
1193                cell.update(cx, |cell, _cx| {
1194                    cell.set_selected(is_selected)
1195                        .set_cell_position(cell_position);
1196                });
1197                cell.clone().into_any_element()
1198            }
1199            Cell::Markdown(cell) => {
1200                cell.update(cx, |cell, _cx| {
1201                    cell.set_selected(is_selected)
1202                        .set_cell_position(cell_position);
1203                });
1204                cell.clone().into_any_element()
1205            }
1206            Cell::Raw(cell) => {
1207                cell.update(cx, |cell, _cx| {
1208                    cell.set_selected(is_selected)
1209                        .set_cell_position(cell_position);
1210                });
1211                cell.clone().into_any_element()
1212            }
1213        }
1214    }
1215}
1216
1217impl Render for NotebookEditor {
1218    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1219        v_flex()
1220            .size_full()
1221            .key_context("NotebookEditor")
1222            .track_focus(&self.focus_handle)
1223            .on_action(cx.listener(|this, _: &OpenNotebook, window, cx| {
1224                this.open_notebook(&OpenNotebook, window, cx)
1225            }))
1226            .on_action(
1227                cx.listener(|this, _: &ClearOutputs, window, cx| this.clear_outputs(window, cx)),
1228            )
1229            .on_action(
1230                cx.listener(|this, _: &Run, window, cx| this.run_current_cell(&Run, window, cx)),
1231            )
1232            .on_action(cx.listener(|this, _: &RunAll, window, cx| this.run_cells(window, cx)))
1233            .on_action(
1234                cx.listener(|this, _: &MoveCellUp, window, cx| this.move_cell_up(window, cx)),
1235            )
1236            .on_action(
1237                cx.listener(|this, _: &MoveCellDown, window, cx| this.move_cell_down(window, cx)),
1238            )
1239            .on_action(cx.listener(|this, _: &AddMarkdownBlock, window, cx| {
1240                this.add_markdown_block(window, cx)
1241            }))
1242            .on_action(
1243                cx.listener(|this, _: &AddCodeBlock, window, cx| this.add_code_block(window, cx)),
1244            )
1245            .on_action(cx.listener(|this, _: &MoveUp, window, cx| {
1246                this.select_previous(&menu::SelectPrevious, window, cx);
1247                if let Some(cell_id) = this.cell_order.get(this.selected_cell_index) {
1248                    if let Some(cell) = this.cell_map.get(cell_id) {
1249                        match cell {
1250                            Cell::Code(cell) => {
1251                                let editor = cell.read(cx).editor().clone();
1252                                editor.update(cx, |editor, cx| {
1253                                    editor.move_to_end(&Default::default(), window, cx);
1254                                });
1255                                editor.focus_handle(cx).focus(window, cx);
1256                            }
1257                            Cell::Markdown(cell) => {
1258                                cell.update(cx, |cell, cx| {
1259                                    cell.set_editing(true);
1260                                    cx.notify();
1261                                });
1262                                let editor = cell.read(cx).editor().clone();
1263                                editor.update(cx, |editor, cx| {
1264                                    editor.move_to_end(&Default::default(), window, cx);
1265                                });
1266                                editor.focus_handle(cx).focus(window, cx);
1267                            }
1268                            _ => {}
1269                        }
1270                    }
1271                }
1272            }))
1273            .on_action(cx.listener(|this, _: &MoveDown, window, cx| {
1274                this.select_next(&menu::SelectNext, window, cx);
1275                if let Some(cell_id) = this.cell_order.get(this.selected_cell_index) {
1276                    if let Some(cell) = this.cell_map.get(cell_id) {
1277                        match cell {
1278                            Cell::Code(cell) => {
1279                                let editor = cell.read(cx).editor().clone();
1280                                editor.update(cx, |editor, cx| {
1281                                    editor.move_to_beginning(&Default::default(), window, cx);
1282                                });
1283                                editor.focus_handle(cx).focus(window, cx);
1284                            }
1285                            Cell::Markdown(cell) => {
1286                                cell.update(cx, |cell, cx| {
1287                                    cell.set_editing(true);
1288                                    cx.notify();
1289                                });
1290                                let editor = cell.read(cx).editor().clone();
1291                                editor.update(cx, |editor, cx| {
1292                                    editor.move_to_beginning(&Default::default(), window, cx);
1293                                });
1294                                editor.focus_handle(cx).focus(window, cx);
1295                            }
1296                            _ => {}
1297                        }
1298                    }
1299                }
1300            }))
1301            .on_action(cx.listener(|this, _: &NotebookMoveDown, window, cx| {
1302                let Some(cell_id) = this.cell_order.get(this.selected_cell_index) else {
1303                    return;
1304                };
1305                let Some(cell) = this.cell_map.get(cell_id) else {
1306                    return;
1307                };
1308
1309                let editor = match cell {
1310                    Cell::Code(cell) => cell.read(cx).editor().clone(),
1311                    Cell::Markdown(cell) => cell.read(cx).editor().clone(),
1312                    _ => return,
1313                };
1314
1315                let is_at_last_line = editor.update(cx, |editor, cx| {
1316                    let display_snapshot = editor.display_snapshot(cx);
1317                    let selections = editor.selections.all_display(&display_snapshot);
1318                    if let Some(selection) = selections.last() {
1319                        let head = selection.head();
1320                        let cursor_row = head.row();
1321                        let max_row = display_snapshot.max_point().row();
1322
1323                        cursor_row >= max_row
1324                    } else {
1325                        false
1326                    }
1327                });
1328
1329                if is_at_last_line {
1330                    this.select_next(&menu::SelectNext, window, cx);
1331                    if let Some(cell_id) = this.cell_order.get(this.selected_cell_index) {
1332                        if let Some(cell) = this.cell_map.get(cell_id) {
1333                            match cell {
1334                                Cell::Code(cell) => {
1335                                    let editor = cell.read(cx).editor().clone();
1336                                    editor.update(cx, |editor, cx| {
1337                                        editor.move_to_beginning(&Default::default(), window, cx);
1338                                    });
1339                                    editor.focus_handle(cx).focus(window, cx);
1340                                }
1341                                Cell::Markdown(cell) => {
1342                                    cell.update(cx, |cell, cx| {
1343                                        cell.set_editing(true);
1344                                        cx.notify();
1345                                    });
1346                                    let editor = cell.read(cx).editor().clone();
1347                                    editor.update(cx, |editor, cx| {
1348                                        editor.move_to_beginning(&Default::default(), window, cx);
1349                                    });
1350                                    editor.focus_handle(cx).focus(window, cx);
1351                                }
1352                                _ => {}
1353                            }
1354                        }
1355                    }
1356                } else {
1357                    editor.update(cx, |editor, cx| {
1358                        editor.move_down(&Default::default(), window, cx);
1359                    });
1360                }
1361            }))
1362            .on_action(cx.listener(|this, _: &NotebookMoveUp, window, cx| {
1363                let Some(cell_id) = this.cell_order.get(this.selected_cell_index) else {
1364                    return;
1365                };
1366                let Some(cell) = this.cell_map.get(cell_id) else {
1367                    return;
1368                };
1369
1370                let editor = match cell {
1371                    Cell::Code(cell) => cell.read(cx).editor().clone(),
1372                    Cell::Markdown(cell) => cell.read(cx).editor().clone(),
1373                    _ => return,
1374                };
1375
1376                let is_at_first_line = editor.update(cx, |editor, cx| {
1377                    let display_snapshot = editor.display_snapshot(cx);
1378                    let selections = editor.selections.all_display(&display_snapshot);
1379                    if let Some(selection) = selections.first() {
1380                        let head = selection.head();
1381                        let cursor_row = head.row();
1382
1383                        cursor_row.0 == 0
1384                    } else {
1385                        false
1386                    }
1387                });
1388
1389                if is_at_first_line {
1390                    this.select_previous(&menu::SelectPrevious, window, cx);
1391                    if let Some(cell_id) = this.cell_order.get(this.selected_cell_index) {
1392                        if let Some(cell) = this.cell_map.get(cell_id) {
1393                            match cell {
1394                                Cell::Code(cell) => {
1395                                    let editor = cell.read(cx).editor().clone();
1396                                    editor.update(cx, |editor, cx| {
1397                                        editor.move_to_end(&Default::default(), window, cx);
1398                                    });
1399                                    editor.focus_handle(cx).focus(window, cx);
1400                                }
1401                                Cell::Markdown(cell) => {
1402                                    cell.update(cx, |cell, cx| {
1403                                        cell.set_editing(true);
1404                                        cx.notify();
1405                                    });
1406                                    let editor = cell.read(cx).editor().clone();
1407                                    editor.update(cx, |editor, cx| {
1408                                        editor.move_to_end(&Default::default(), window, cx);
1409                                    });
1410                                    editor.focus_handle(cx).focus(window, cx);
1411                                }
1412                                _ => {}
1413                            }
1414                        }
1415                    }
1416                } else {
1417                    editor.update(cx, |editor, cx| {
1418                        editor.move_up(&Default::default(), window, cx);
1419                    });
1420                }
1421            }))
1422            .on_action(
1423                cx.listener(|this, action, window, cx| this.restart_kernel(action, window, cx)),
1424            )
1425            .on_action(
1426                cx.listener(|this, action, window, cx| this.interrupt_kernel(action, window, cx)),
1427            )
1428            .child(
1429                h_flex()
1430                    .flex_1()
1431                    .w_full()
1432                    .h_full()
1433                    .gap_2()
1434                    .child(div().flex_1().h_full().child(self.cell_list(window, cx)))
1435                    .child(self.render_notebook_controls(window, cx)),
1436            )
1437            .child(self.render_kernel_status_bar(window, cx))
1438    }
1439}
1440
1441impl Focusable for NotebookEditor {
1442    fn focus_handle(&self, _: &App) -> FocusHandle {
1443        self.focus_handle.clone()
1444    }
1445}
1446
1447// Intended to be a NotebookBuffer
1448pub struct NotebookItem {
1449    path: PathBuf,
1450    project_path: ProjectPath,
1451    languages: Arc<LanguageRegistry>,
1452    // Raw notebook data
1453    notebook: nbformat::v4::Notebook,
1454    // Store our version of the notebook in memory (cell_order, cell_map)
1455    id: ProjectEntryId,
1456}
1457
1458impl project::ProjectItem for NotebookItem {
1459    fn try_open(
1460        project: &Entity<Project>,
1461        path: &ProjectPath,
1462        cx: &mut App,
1463    ) -> Option<Task<anyhow::Result<Entity<Self>>>> {
1464        let path = path.clone();
1465        let project = project.clone();
1466        let fs = project.read(cx).fs().clone();
1467        let languages = project.read(cx).languages().clone();
1468
1469        if path.path.extension().unwrap_or_default() == "ipynb" {
1470            Some(cx.spawn(async move |cx| {
1471                let abs_path = project
1472                    .read_with(cx, |project, cx| project.absolute_path(&path, cx))
1473                    .with_context(|| format!("finding the absolute path of {path:?}"))?;
1474
1475                // todo: watch for changes to the file
1476                let buffer = project
1477                    .update(cx, |project, cx| project.open_buffer(path.clone(), cx))
1478                    .await?;
1479                let file_content = buffer.read_with(cx, |buffer, _| buffer.text());
1480
1481                let notebook = if file_content.trim().is_empty() {
1482                    nbformat::v4::Notebook {
1483                        nbformat: 4,
1484                        nbformat_minor: 5,
1485                        cells: vec![],
1486                        metadata: serde_json::from_str("{}").unwrap(),
1487                    }
1488                } else {
1489                    let notebook = match nbformat::parse_notebook(&file_content) {
1490                        Ok(nb) => nb,
1491                        Err(_) => {
1492                            // Pre-process to ensure IDs exist
1493                            let mut json: serde_json::Value = serde_json::from_str(&file_content)?;
1494                            if let Some(cells) =
1495                                json.get_mut("cells").and_then(|c| c.as_array_mut())
1496                            {
1497                                for cell in cells {
1498                                    if cell.get("id").is_none() {
1499                                        cell["id"] =
1500                                            serde_json::Value::String(Uuid::new_v4().to_string());
1501                                    }
1502                                }
1503                            }
1504                            let file_content = serde_json::to_string(&json)?;
1505                            nbformat::parse_notebook(&file_content)?
1506                        }
1507                    };
1508
1509                    match notebook {
1510                        nbformat::Notebook::V4(notebook) => notebook,
1511                        // 4.1 - 4.4 are converted to 4.5
1512                        nbformat::Notebook::Legacy(legacy_notebook) => {
1513                            // TODO: Decide if we want to mutate the notebook by including Cell IDs
1514                            // and any other conversions
1515
1516                            nbformat::upgrade_legacy_notebook(legacy_notebook)?
1517                        }
1518                        nbformat::Notebook::V3(v3_notebook) => {
1519                            nbformat::upgrade_v3_notebook(v3_notebook)?
1520                        }
1521                    }
1522                };
1523
1524                let id = project
1525                    .update(cx, |project, cx| {
1526                        project.entry_for_path(&path, cx).map(|entry| entry.id)
1527                    })
1528                    .context("Entry not found")?;
1529
1530                Ok(cx.new(|_| NotebookItem {
1531                    path: abs_path,
1532                    project_path: path,
1533                    languages,
1534                    notebook,
1535                    id,
1536                }))
1537            }))
1538        } else {
1539            None
1540        }
1541    }
1542
1543    fn entry_id(&self, _: &App) -> Option<ProjectEntryId> {
1544        Some(self.id)
1545    }
1546
1547    fn project_path(&self, _: &App) -> Option<ProjectPath> {
1548        Some(self.project_path.clone())
1549    }
1550
1551    fn is_dirty(&self) -> bool {
1552        // TODO: Track if notebook metadata or structure has changed
1553        false
1554    }
1555}
1556
1557impl NotebookItem {
1558    pub fn language_name(&self) -> Option<String> {
1559        self.notebook
1560            .metadata
1561            .language_info
1562            .as_ref()
1563            .map(|l| l.name.clone())
1564            .or(self
1565                .notebook
1566                .metadata
1567                .kernelspec
1568                .as_ref()
1569                .and_then(|spec| spec.language.clone()))
1570    }
1571
1572    pub fn notebook_language(&self) -> impl Future<Output = Option<Arc<Language>>> + use<> {
1573        let language_name = self.language_name();
1574        let languages = self.languages.clone();
1575
1576        async move {
1577            if let Some(language_name) = language_name {
1578                languages.language_for_name(&language_name).await.ok()
1579            } else {
1580                None
1581            }
1582        }
1583    }
1584}
1585
1586impl EventEmitter<()> for NotebookItem {}
1587
1588impl EventEmitter<()> for NotebookEditor {}
1589
1590// pub struct NotebookControls {
1591//     pane_focused: bool,
1592//     active_item: Option<Box<dyn ItemHandle>>,
1593//     // subscription: Option<Subscription>,
1594// }
1595
1596// impl NotebookControls {
1597//     pub fn new() -> Self {
1598//         Self {
1599//             pane_focused: false,
1600//             active_item: Default::default(),
1601//             // subscription: Default::default(),
1602//         }
1603//     }
1604// }
1605
1606// impl EventEmitter<ToolbarItemEvent> for NotebookControls {}
1607
1608// impl Render for NotebookControls {
1609//     fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
1610//         div().child("notebook controls")
1611//     }
1612// }
1613
1614// impl ToolbarItemView for NotebookControls {
1615//     fn set_active_pane_item(
1616//         &mut self,
1617//         active_pane_item: Option<&dyn workspace::ItemHandle>,
1618//         window: &mut Window, cx: &mut Context<Self>,
1619//     ) -> workspace::ToolbarItemLocation {
1620//         cx.notify();
1621//         self.active_item = None;
1622
1623//         let Some(item) = active_pane_item else {
1624//             return ToolbarItemLocation::Hidden;
1625//         };
1626
1627//         ToolbarItemLocation::PrimaryLeft
1628//     }
1629
1630//     fn pane_focus_update(&mut self, pane_focused: bool, _window: &mut Window, _cx: &mut Context<Self>) {
1631//         self.pane_focused = pane_focused;
1632//     }
1633// }
1634
1635impl Item for NotebookEditor {
1636    type Event = ();
1637
1638    fn can_split(&self) -> bool {
1639        true
1640    }
1641
1642    fn clone_on_split(
1643        &self,
1644        _workspace_id: Option<workspace::WorkspaceId>,
1645        window: &mut Window,
1646        cx: &mut Context<Self>,
1647    ) -> Task<Option<Entity<Self>>>
1648    where
1649        Self: Sized,
1650    {
1651        Task::ready(Some(cx.new(|cx| {
1652            Self::new(self.project.clone(), self.notebook_item.clone(), window, cx)
1653        })))
1654    }
1655
1656    fn buffer_kind(&self, _: &App) -> workspace::item::ItemBufferKind {
1657        workspace::item::ItemBufferKind::Singleton
1658    }
1659
1660    fn for_each_project_item(
1661        &self,
1662        cx: &App,
1663        f: &mut dyn FnMut(gpui::EntityId, &dyn project::ProjectItem),
1664    ) {
1665        f(self.notebook_item.entity_id(), self.notebook_item.read(cx))
1666    }
1667
1668    fn tab_content_text(&self, _detail: usize, cx: &App) -> SharedString {
1669        self.notebook_item
1670            .read(cx)
1671            .project_path
1672            .path
1673            .file_name()
1674            .map(|s| s.to_string())
1675            .unwrap_or_default()
1676            .into()
1677    }
1678
1679    fn tab_content(&self, params: TabContentParams, window: &Window, cx: &App) -> AnyElement {
1680        Label::new(self.tab_content_text(params.detail.unwrap_or(0), cx))
1681            .single_line()
1682            .color(params.text_color())
1683            .when(params.preview, |this| this.italic())
1684            .into_any_element()
1685    }
1686
1687    fn tab_icon(&self, _window: &Window, _cx: &App) -> Option<Icon> {
1688        Some(IconName::Book.into())
1689    }
1690
1691    fn show_toolbar(&self) -> bool {
1692        false
1693    }
1694
1695    // TODO
1696    fn pixel_position_of_cursor(&self, _: &App) -> Option<Point<Pixels>> {
1697        None
1698    }
1699
1700    // TODO
1701    fn as_searchable(&self, _: &Entity<Self>, _: &App) -> Option<Box<dyn SearchableItemHandle>> {
1702        None
1703    }
1704
1705    fn set_nav_history(
1706        &mut self,
1707        _: workspace::ItemNavHistory,
1708        _window: &mut Window,
1709        _: &mut Context<Self>,
1710    ) {
1711        // TODO
1712    }
1713
1714    fn can_save(&self, _cx: &App) -> bool {
1715        true
1716    }
1717
1718    fn save(
1719        &mut self,
1720        _options: SaveOptions,
1721        project: Entity<Project>,
1722        _window: &mut Window,
1723        cx: &mut Context<Self>,
1724    ) -> Task<Result<()>> {
1725        let notebook = self.to_notebook(cx);
1726        let path = self.notebook_item.read(cx).path.clone();
1727        let fs = project.read(cx).fs().clone();
1728
1729        self.mark_as_saved(cx);
1730
1731        cx.spawn(async move |_this, _cx| {
1732            let json =
1733                serde_json::to_string_pretty(&notebook).context("Failed to serialize notebook")?;
1734            fs.atomic_write(path, json).await?;
1735            Ok(())
1736        })
1737    }
1738
1739    fn save_as(
1740        &mut self,
1741        project: Entity<Project>,
1742        path: ProjectPath,
1743        _window: &mut Window,
1744        cx: &mut Context<Self>,
1745    ) -> Task<Result<()>> {
1746        let notebook = self.to_notebook(cx);
1747        let fs = project.read(cx).fs().clone();
1748
1749        let abs_path = project.read(cx).absolute_path(&path, cx);
1750
1751        self.mark_as_saved(cx);
1752
1753        cx.spawn(async move |_this, _cx| {
1754            let abs_path = abs_path.context("Failed to get absolute path")?;
1755            let json =
1756                serde_json::to_string_pretty(&notebook).context("Failed to serialize notebook")?;
1757            fs.atomic_write(abs_path, json).await?;
1758            Ok(())
1759        })
1760    }
1761
1762    fn reload(
1763        &mut self,
1764        project: Entity<Project>,
1765        window: &mut Window,
1766        cx: &mut Context<Self>,
1767    ) -> Task<Result<()>> {
1768        let project_path = self.notebook_item.read(cx).project_path.clone();
1769        let languages = self.languages.clone();
1770        let notebook_language = self.notebook_language.clone();
1771
1772        cx.spawn_in(window, async move |this, cx| {
1773            let buffer = this
1774                .update(cx, |this, cx| {
1775                    this.project
1776                        .update(cx, |project, cx| project.open_buffer(project_path, cx))
1777                })?
1778                .await?;
1779
1780            let file_content = buffer.read_with(cx, |buffer, _| buffer.text());
1781
1782            let mut json: serde_json::Value = serde_json::from_str(&file_content)?;
1783            if let Some(cells) = json.get_mut("cells").and_then(|c| c.as_array_mut()) {
1784                for cell in cells {
1785                    if cell.get("id").is_none() {
1786                        cell["id"] = serde_json::Value::String(Uuid::new_v4().to_string());
1787                    }
1788                }
1789            }
1790            let file_content = serde_json::to_string(&json)?;
1791
1792            let notebook = nbformat::parse_notebook(&file_content);
1793            let notebook = match notebook {
1794                Ok(nbformat::Notebook::V4(notebook)) => notebook,
1795                Ok(nbformat::Notebook::Legacy(legacy_notebook)) => {
1796                    nbformat::upgrade_legacy_notebook(legacy_notebook)?
1797                }
1798                Ok(nbformat::Notebook::V3(v3_notebook)) => {
1799                    nbformat::upgrade_v3_notebook(v3_notebook)?
1800                }
1801                Err(e) => {
1802                    anyhow::bail!("Failed to parse notebook: {:?}", e);
1803                }
1804            };
1805
1806            this.update_in(cx, |this, window, cx| {
1807                let mut cell_order = vec![];
1808                let mut cell_map = HashMap::default();
1809
1810                for cell in notebook.cells.iter() {
1811                    let cell_id = cell.id();
1812                    cell_order.push(cell_id.clone());
1813                    let cell_entity =
1814                        Cell::load(cell, &languages, notebook_language.clone(), window, cx);
1815                    cell_map.insert(cell_id.clone(), cell_entity);
1816                }
1817
1818                this.cell_order = cell_order.clone();
1819                this.original_cell_order = cell_order;
1820                this.cell_map = cell_map;
1821                this.cell_list =
1822                    ListState::new(this.cell_order.len(), gpui::ListAlignment::Top, px(1000.));
1823                cx.notify();
1824            })?;
1825
1826            Ok(())
1827        })
1828    }
1829
1830    fn is_dirty(&self, cx: &App) -> bool {
1831        self.has_structural_changes() || self.has_content_changes(cx)
1832    }
1833}
1834
1835impl ProjectItem for NotebookEditor {
1836    type Item = NotebookItem;
1837
1838    fn for_project_item(
1839        project: Entity<Project>,
1840        _pane: Option<&Pane>,
1841        item: Entity<Self::Item>,
1842        window: &mut Window,
1843        cx: &mut Context<Self>,
1844    ) -> Self {
1845        Self::new(project, item, window, cx)
1846    }
1847}
1848
1849impl KernelSession for NotebookEditor {
1850    fn route(&mut self, message: &JupyterMessage, window: &mut Window, cx: &mut Context<Self>) {
1851        // Handle kernel status updates (these are broadcast to all)
1852        if let JupyterMessageContent::Status(status) = &message.content {
1853            self.kernel.set_execution_state(&status.execution_state);
1854            cx.notify();
1855        }
1856
1857        if let JupyterMessageContent::KernelInfoReply(reply) = &message.content {
1858            self.kernel.set_kernel_info(reply);
1859
1860            if let Ok(language_info) = serde_json::from_value::<nbformat::v4::LanguageInfo>(
1861                serde_json::to_value(&reply.language_info).unwrap(),
1862            ) {
1863                self.notebook_item.update(cx, |item, cx| {
1864                    item.notebook.metadata.language_info = Some(language_info);
1865                    cx.emit(());
1866                });
1867            }
1868            cx.notify();
1869        }
1870
1871        // Handle cell-specific messages
1872        if let Some(parent_header) = &message.parent_header {
1873            if let Some(cell_id) = self.execution_requests.get(&parent_header.msg_id) {
1874                if let Some(Cell::Code(cell)) = self.cell_map.get(cell_id) {
1875                    cell.update(cx, |cell, cx| {
1876                        cell.handle_message(message, window, cx);
1877                    });
1878                }
1879            }
1880        }
1881    }
1882
1883    fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>) {
1884        self.kernel = Kernel::ErroredLaunch(error_message);
1885        cx.notify();
1886    }
1887}