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