Bring Jupyter to Zed Editing (#12062)

Run any Jupyter kernel in Zed on any buffer (editor):

<img width="1074" alt="image"
src="https://github.com/zed-industries/zed/assets/836375/eac8ed69-d02b-4d46-b379-6186d8f59470">

## TODO

### Lifecycle

* [x] Launch kernels on demand
* [x] Wait for kernel to be started
* [x] Request Kernel info on start
* [x] Show in progress indicator
* [ ] Allow picking kernel (it defaults to first matching language name)
* [ ] Menu for interrupting and shutting down the kernel
* [ ] Drop running kernels once editor is dropped

### Media Outputs

* [x] Render text and tracebacks with ANSI color handling
* [x] Render markdown as text
* [x] Render PNG and JPEG images using an explicit height based on
line-height
* ~~Render SVG~~ -- not happening for this PR due to lack of text in SVG
support
* [ ] Process `update_display_data` message and related `display_id`
* [x] Process `page` data from payloads as outputs
* [ ] Render markdown as, well, rendered markdown -- Note: unsure if we
can get line heights here

### Document

* [x] Select code and run
* [x] Run current line
* [x] Clear previous overlapping runs
* [ ] Support running markdown code blocks
* [ ] Action to export session as notebook or output files
* [ ] Action to clear all outputs
* [ ] Delete outputs when lines are deleted

## Other missing features

The following is a list of missing functionality or expectations that
are out of scope for this PR.

### Python Environments

Detecting python environments should probably be done in a separate PR
in tandem with how they're used with LSP. Users likely want to pick an
environment for their project, whether a virtualenv, conda env, pyenv,
poetry backed virtualenv, or the system. Related issues:

* https://github.com/zed-industries/zed/issues/7646
* https://github.com/zed-industries/zed/issues/7808
* https://github.com/zed-industries/zed/issues/7296

### LSP Integration

* Submit `complete_request` messages for completions to interleave
interactive variables with LSP
* LSP for IPython semantics (`%%timeit`, `!ls`, `get_ipython`, etc.)

## Future release notes

- Run code in any editor, whether it's a script or a markdown document

Release Notes:

- N/A

crates/repl/src/repl.rs

@@ -0,0 +1,571 @@
+use anyhow::{anyhow, Context as _, Result};
+use async_dispatcher::{set_dispatcher, timeout, Dispatcher, Runnable};
+use collections::{HashMap, HashSet};
+use editor::{
+    display_map::{
+        BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, RenderBlock,
+    },
+    Anchor, AnchorRangeExt, Editor,
+};
+use futures::{
+    channel::mpsc::{self, UnboundedSender},
+    future::Shared,
+    Future, FutureExt, SinkExt as _, StreamExt,
+};
+use gpui::prelude::*;
+use gpui::{
+    actions, AppContext, Context, EntityId, Global, Model, ModelContext, PlatformDispatcher, Task,
+    WeakView,
+};
+use gpui::{Entity, View};
+use language::Point;
+use outputs::{ExecutionStatus, ExecutionView, LineHeight as _};
+use project::Fs;
+use runtime_settings::JupyterSettings;
+use runtimelib::JupyterMessageContent;
+use settings::{Settings as _, SettingsStore};
+use std::{ops::Range, time::Instant};
+use std::{sync::Arc, time::Duration};
+use theme::{ActiveTheme, ThemeSettings};
+use ui::prelude::*;
+use workspace::Workspace;
+
+mod outputs;
+// mod runtime_panel;
+mod runtime_settings;
+mod runtimes;
+mod stdio;
+
+use runtimes::{get_runtime_specifications, Request, RunningKernel, RuntimeSpecification};
+
+actions!(repl, [Run]);
+
+#[derive(Clone)]
+pub struct RuntimeManagerGlobal(Model<RuntimeManager>);
+
+impl Global for RuntimeManagerGlobal {}
+
+pub fn zed_dispatcher(cx: &mut AppContext) -> impl Dispatcher {
+    struct ZedDispatcher {
+        dispatcher: Arc<dyn PlatformDispatcher>,
+    }
+
+    // PlatformDispatcher is _super_ close to the same interface we put in
+    // async-dispatcher, except for the task label in dispatch. Later we should
+    // just make that consistent so we have this dispatcher ready to go for
+    // other crates in Zed.
+    impl Dispatcher for ZedDispatcher {
+        fn dispatch(&self, runnable: Runnable) {
+            self.dispatcher.dispatch(runnable, None)
+        }
+
+        fn dispatch_after(&self, duration: Duration, runnable: Runnable) {
+            self.dispatcher.dispatch_after(duration, runnable);
+        }
+    }
+
+    ZedDispatcher {
+        dispatcher: cx.background_executor().dispatcher.clone(),
+    }
+}
+
+pub fn init(fs: Arc<dyn Fs>, cx: &mut AppContext) {
+    set_dispatcher(zed_dispatcher(cx));
+    JupyterSettings::register(cx);
+
+    observe_jupyter_settings_changes(fs.clone(), cx);
+
+    cx.observe_new_views(
+        |workspace: &mut Workspace, _: &mut ViewContext<Workspace>| {
+            workspace.register_action(run);
+        },
+    )
+    .detach();
+
+    let settings = JupyterSettings::get_global(cx);
+
+    if !settings.enabled {
+        return;
+    }
+
+    initialize_runtime_manager(fs, cx);
+}
+
+fn initialize_runtime_manager(fs: Arc<dyn Fs>, cx: &mut AppContext) {
+    let runtime_manager = cx.new_model(|cx| RuntimeManager::new(fs.clone(), cx));
+    RuntimeManager::set_global(runtime_manager.clone(), cx);
+
+    cx.spawn(|mut cx| async move {
+        let fs = fs.clone();
+
+        let runtime_specifications = get_runtime_specifications(fs).await?;
+
+        runtime_manager.update(&mut cx, |this, _cx| {
+            this.runtime_specifications = runtime_specifications;
+        })?;
+
+        anyhow::Ok(())
+    })
+    .detach_and_log_err(cx);
+}
+
+fn observe_jupyter_settings_changes(fs: Arc<dyn Fs>, cx: &mut AppContext) {
+    cx.observe_global::<SettingsStore>(move |cx| {
+        let settings = JupyterSettings::get_global(cx);
+        if settings.enabled && RuntimeManager::global(cx).is_none() {
+            initialize_runtime_manager(fs.clone(), cx);
+        } else {
+            RuntimeManager::remove_global(cx);
+            // todo!(): Remove action from workspace(s)
+        }
+    })
+    .detach();
+}
+
+#[derive(Debug)]
+pub enum Kernel {
+    RunningKernel(RunningKernel),
+    StartingKernel(Shared<Task<()>>),
+    FailedLaunch,
+}
+
+// Per workspace
+pub struct RuntimeManager {
+    fs: Arc<dyn Fs>,
+    runtime_specifications: Vec<RuntimeSpecification>,
+
+    instances: HashMap<EntityId, Kernel>,
+    editors: HashMap<WeakView<Editor>, EditorRuntimeState>,
+    // todo!(): Next
+    // To reduce the number of open tasks and channels we have, let's feed the response
+    // messages by ID over to the paired ExecutionView
+    _execution_views_by_id: HashMap<String, View<ExecutionView>>,
+}
+
+#[derive(Debug, Clone)]
+struct EditorRuntimeState {
+    blocks: Vec<EditorRuntimeBlock>,
+    // todo!(): Store a subscription to the editor so we can drop them when the editor is dropped
+    // subscription: gpui::Subscription,
+}
+
+#[derive(Debug, Clone)]
+struct EditorRuntimeBlock {
+    code_range: Range<Anchor>,
+    _execution_id: String,
+    block_id: BlockId,
+    _execution_view: View<ExecutionView>,
+}
+
+impl RuntimeManager {
+    pub fn new(fs: Arc<dyn Fs>, _cx: &mut AppContext) -> Self {
+        Self {
+            fs,
+            runtime_specifications: Default::default(),
+            instances: Default::default(),
+            editors: Default::default(),
+            _execution_views_by_id: Default::default(),
+        }
+    }
+
+    fn get_or_launch_kernel(
+        &mut self,
+        entity_id: EntityId,
+        language_name: Arc<str>,
+        cx: &mut ModelContext<Self>,
+    ) -> Task<Result<UnboundedSender<Request>>> {
+        let kernel = self.instances.get(&entity_id);
+        let pending_kernel_start = match kernel {
+            Some(Kernel::RunningKernel(running_kernel)) => {
+                return Task::ready(anyhow::Ok(running_kernel.request_tx.clone()));
+            }
+            Some(Kernel::StartingKernel(task)) => task.clone(),
+            Some(Kernel::FailedLaunch) | None => {
+                self.instances.remove(&entity_id);
+
+                let kernel = self.launch_kernel(entity_id, language_name, cx);
+                let pending_kernel = cx
+                    .spawn(|this, mut cx| async move {
+                        let running_kernel = kernel.await;
+
+                        match running_kernel {
+                            Ok(running_kernel) => {
+                                let _ = this.update(&mut cx, |this, _cx| {
+                                    this.instances
+                                        .insert(entity_id, Kernel::RunningKernel(running_kernel));
+                                });
+                            }
+                            Err(_err) => {
+                                let _ = this.update(&mut cx, |this, _cx| {
+                                    this.instances.insert(entity_id, Kernel::FailedLaunch);
+                                });
+                            }
+                        }
+                    })
+                    .shared();
+
+                self.instances
+                    .insert(entity_id, Kernel::StartingKernel(pending_kernel.clone()));
+
+                pending_kernel
+            }
+        };
+
+        cx.spawn(|this, mut cx| async move {
+            pending_kernel_start.await;
+
+            this.update(&mut cx, |this, _cx| {
+                let kernel = this
+                    .instances
+                    .get(&entity_id)
+                    .ok_or(anyhow!("unable to get a running kernel"))?;
+
+                match kernel {
+                    Kernel::RunningKernel(running_kernel) => Ok(running_kernel.request_tx.clone()),
+                    _ => Err(anyhow!("unable to get a running kernel")),
+                }
+            })?
+        })
+    }
+
+    fn launch_kernel(
+        &mut self,
+        entity_id: EntityId,
+        language_name: Arc<str>,
+        cx: &mut ModelContext<Self>,
+    ) -> Task<Result<RunningKernel>> {
+        // Get first runtime that matches the language name (for now)
+        let runtime_specification =
+            self.runtime_specifications
+                .iter()
+                .find(|runtime_specification| {
+                    runtime_specification.kernelspec.language == language_name.to_string()
+                });
+
+        let runtime_specification = match runtime_specification {
+            Some(runtime_specification) => runtime_specification,
+            None => {
+                return Task::ready(Err(anyhow::anyhow!(
+                    "No runtime found for language {}",
+                    language_name
+                )));
+            }
+        };
+
+        let runtime_specification = runtime_specification.clone();
+
+        let fs = self.fs.clone();
+
+        cx.spawn(|_, cx| async move {
+            let running_kernel =
+                RunningKernel::new(runtime_specification, entity_id, fs.clone(), cx);
+
+            let running_kernel = running_kernel.await?;
+
+            let mut request_tx = running_kernel.request_tx.clone();
+
+            let overall_timeout_duration = Duration::from_secs(10);
+
+            let start_time = Instant::now();
+
+            loop {
+                if start_time.elapsed() > overall_timeout_duration {
+                    // todo!(): Kill the kernel
+                    return Err(anyhow::anyhow!("Kernel did not respond in time"));
+                }
+
+                let (tx, rx) = mpsc::unbounded();
+                match request_tx
+                    .send(Request {
+                        request: runtimelib::KernelInfoRequest {}.into(),
+                        responses_rx: tx,
+                    })
+                    .await
+                {
+                    Ok(_) => {}
+                    Err(_err) => {
+                        break;
+                    }
+                };
+
+                let mut rx = rx.fuse();
+
+                let kernel_info_timeout = Duration::from_secs(1);
+
+                let mut got_kernel_info = false;
+                while let Ok(Some(message)) = timeout(kernel_info_timeout, rx.next()).await {
+                    match message {
+                        JupyterMessageContent::KernelInfoReply(_) => {
+                            got_kernel_info = true;
+                        }
+                        _ => {}
+                    }
+                }
+
+                if got_kernel_info {
+                    break;
+                }
+            }
+
+            anyhow::Ok(running_kernel)
+        })
+    }
+
+    fn execute_code(
+        &mut self,
+        entity_id: EntityId,
+        language_name: Arc<str>,
+        code: String,
+        cx: &mut ModelContext<Self>,
+    ) -> impl Future<Output = Result<mpsc::UnboundedReceiver<JupyterMessageContent>>> {
+        let (tx, rx) = mpsc::unbounded();
+
+        let request_tx = self.get_or_launch_kernel(entity_id, language_name, cx);
+
+        async move {
+            let request_tx = request_tx.await?;
+
+            request_tx
+                .unbounded_send(Request {
+                    request: runtimelib::ExecuteRequest {
+                        code,
+                        allow_stdin: false,
+                        silent: false,
+                        store_history: true,
+                        stop_on_error: true,
+                        ..Default::default()
+                    }
+                    .into(),
+                    responses_rx: tx,
+                })
+                .context("Failed to send execution request")?;
+
+            Ok(rx)
+        }
+    }
+
+    pub fn global(cx: &AppContext) -> Option<Model<Self>> {
+        cx.try_global::<RuntimeManagerGlobal>()
+            .map(|runtime_manager| runtime_manager.0.clone())
+    }
+
+    pub fn set_global(runtime_manager: Model<Self>, cx: &mut AppContext) {
+        cx.set_global(RuntimeManagerGlobal(runtime_manager));
+    }
+
+    pub fn remove_global(cx: &mut AppContext) {
+        if RuntimeManager::global(cx).is_some() {
+            cx.remove_global::<RuntimeManagerGlobal>();
+        }
+    }
+}
+
+pub fn get_active_editor(
+    workspace: &mut Workspace,
+    cx: &mut ViewContext<Workspace>,
+) -> Option<View<Editor>> {
+    workspace
+        .active_item(cx)
+        .and_then(|item| item.act_as::<Editor>(cx))
+}
+
+// Gets the active selection in the editor or the current line
+pub fn selection(editor: View<Editor>, cx: &mut ViewContext<Workspace>) -> Range<Anchor> {
+    let editor = editor.read(cx);
+    let selection = editor.selections.newest::<usize>(cx);
+    let buffer = editor.buffer().read(cx).snapshot(cx);
+
+    let range = if selection.is_empty() {
+        let cursor = selection.head();
+
+        let line_start = buffer.offset_to_point(cursor).row;
+        let mut start_offset = buffer.point_to_offset(Point::new(line_start, 0));
+
+        // Iterate backwards to find the start of the line
+        while start_offset > 0 {
+            let ch = buffer.chars_at(start_offset - 1).next().unwrap_or('\0');
+            if ch == '\n' {
+                break;
+            }
+            start_offset -= 1;
+        }
+
+        let mut end_offset = cursor;
+
+        // Iterate forwards to find the end of the line
+        while end_offset < buffer.len() {
+            let ch = buffer.chars_at(end_offset).next().unwrap_or('\0');
+            if ch == '\n' {
+                break;
+            }
+            end_offset += 1;
+        }
+
+        // Create a range from the start to the end of the line
+        start_offset..end_offset
+    } else {
+        selection.range()
+    };
+
+    let anchor_range = buffer.anchor_before(range.start)..buffer.anchor_after(range.end);
+    anchor_range
+}
+
+pub fn run(workspace: &mut Workspace, _: &Run, cx: &mut ViewContext<Workspace>) {
+    let (editor, runtime_manager) = if let (Some(editor), Some(runtime_manager)) =
+        (get_active_editor(workspace, cx), RuntimeManager::global(cx))
+    {
+        (editor, runtime_manager)
+    } else {
+        log::warn!("No active editor or runtime manager found");
+        return;
+    };
+
+    let anchor_range = selection(editor.clone(), cx);
+
+    let buffer = editor.read(cx).buffer().read(cx).snapshot(cx);
+
+    let selected_text = buffer
+        .text_for_range(anchor_range.clone())
+        .collect::<String>();
+
+    let start_language = buffer.language_at(anchor_range.start);
+    let end_language = buffer.language_at(anchor_range.end);
+
+    let language_name = if start_language == end_language {
+        start_language
+            .map(|language| language.code_fence_block_name())
+            .filter(|lang| **lang != *"markdown")
+    } else {
+        // If the selection spans multiple languages, don't run it
+        return;
+    };
+
+    let language_name = if let Some(language_name) = language_name {
+        language_name
+    } else {
+        return;
+    };
+
+    let entity_id = editor.entity_id();
+
+    let execution_view = cx.new_view(|cx| ExecutionView::new(cx));
+
+    // If any block overlaps with the new block, remove it
+    // TODO: When inserting a new block, put it in order so that search is efficient
+    let blocks_to_remove = runtime_manager.update(cx, |runtime_manager, _cx| {
+        // Get the current `EditorRuntimeState` for this runtime_manager, inserting it if it doesn't exist
+        let editor_runtime_state = runtime_manager
+            .editors
+            .entry(editor.downgrade())
+            .or_insert_with(|| EditorRuntimeState { blocks: Vec::new() });
+
+        let mut blocks_to_remove: HashSet<BlockId> = HashSet::default();
+
+        editor_runtime_state.blocks.retain(|block| {
+            if anchor_range.overlaps(&block.code_range, &buffer) {
+                blocks_to_remove.insert(block.block_id);
+                // Drop this block
+                false
+            } else {
+                true
+            }
+        });
+
+        blocks_to_remove
+    });
+
+    let blocks_to_remove = blocks_to_remove.clone();
+
+    let block_id = editor.update(cx, |editor, cx| {
+        editor.remove_blocks(blocks_to_remove, None, cx);
+        let block = BlockProperties {
+            position: anchor_range.end,
+            height: execution_view.num_lines(cx).saturating_add(1),
+            style: BlockStyle::Sticky,
+            render: create_output_area_render(execution_view.clone()),
+            disposition: BlockDisposition::Below,
+        };
+
+        editor.insert_blocks([block], None, cx)[0]
+    });
+
+    let receiver = runtime_manager.update(cx, |runtime_manager, cx| {
+        let editor_runtime_state = runtime_manager
+            .editors
+            .entry(editor.downgrade())
+            .or_insert_with(|| EditorRuntimeState { blocks: Vec::new() });
+
+        let editor_runtime_block = EditorRuntimeBlock {
+            code_range: anchor_range.clone(),
+            block_id,
+            _execution_view: execution_view.clone(),
+            _execution_id: Default::default(),
+        };
+
+        editor_runtime_state
+            .blocks
+            .push(editor_runtime_block.clone());
+
+        runtime_manager.execute_code(entity_id, language_name, selected_text.clone(), cx)
+    });
+
+    cx.spawn(|_this, mut cx| async move {
+        execution_view.update(&mut cx, |execution_view, cx| {
+            execution_view.set_status(ExecutionStatus::ConnectingToKernel, cx);
+        })?;
+        let mut receiver = receiver.await?;
+
+        let execution_view = execution_view.clone();
+        while let Some(content) = receiver.next().await {
+            execution_view.update(&mut cx, |execution_view, cx| {
+                execution_view.push_message(&content, cx)
+            })?;
+
+            editor.update(&mut cx, |editor, cx| {
+                let mut replacements = HashMap::default();
+                replacements.insert(
+                    block_id,
+                    (
+                        Some(execution_view.num_lines(cx).saturating_add(1)),
+                        create_output_area_render(execution_view.clone()),
+                    ),
+                );
+                editor.replace_blocks(replacements, None, cx);
+            })?;
+        }
+        anyhow::Ok(())
+    })
+    .detach_and_log_err(cx);
+}
+
+fn create_output_area_render(execution_view: View<ExecutionView>) -> RenderBlock {
+    let render = move |cx: &mut BlockContext| {
+        let execution_view = execution_view.clone();
+        let text_font = ThemeSettings::get_global(cx).buffer_font.family.clone();
+        // Note: we'll want to use `cx.anchor_x` when someone runs something with no output -- just show a checkmark and not make the full block below the line
+
+        let gutter_width = cx.gutter_dimensions.width;
+
+        h_flex()
+            .w_full()
+            .bg(cx.theme().colors().background)
+            .border_y_1()
+            .border_color(cx.theme().colors().border)
+            .pl(gutter_width)
+            .child(
+                div()
+                    .font_family(text_font)
+                    // .ml(gutter_width)
+                    .mx_1()
+                    .my_2()
+                    .h_full()
+                    .w_full()
+                    .mr(gutter_width)
+                    .child(execution_view),
+            )
+            .into_any_element()
+    };
+
+    Box::new(render)
+}