zed/crates/repl/src/session.rs

use crate::components::KernelListItem;
use crate::setup_editor_session_actions;
use crate::{
    KernelStatus,
    kernels::{
        Kernel, KernelSession, KernelSpecification, NativeRunningKernel, RemoteRunningKernel,
        SshRunningKernel, WslRunningKernel,
    },
    outputs::{
        ExecutionStatus, ExecutionView, ExecutionViewFinishedEmpty, ExecutionViewFinishedSmall,
        InputReplyEvent,
    },
    repl_settings::ReplSettings,
};
use anyhow::Context as _;
use collections::{HashMap, HashSet};
use editor::SelectionEffects;
use editor::{
    Anchor, AnchorRangeExt as _, Editor, Inlay, MultiBuffer, ToOffset, ToPoint,
    display_map::{
        BlockContext, BlockId, BlockPlacement, BlockProperties, BlockStyle, CustomBlockId,
        RenderBlock,
    },
    scroll::Autoscroll,
};
use project::InlayId;

/// Marker types
enum ReplExecutedRange {}

use futures::FutureExt as _;
use gpui::{
    Context, Entity, EventEmitter, Render, Subscription, Task, WeakEntity, Window, div, prelude::*,
};
use language::Point;
use project::Fs;
use runtimelib::{
    ExecuteRequest, ExecutionState, InputReply, InterruptRequest, JupyterMessage,
    JupyterMessageContent, KernelInfoRequest, ReplyStatus, ShutdownRequest,
};
use settings::Settings as _;
use std::{env::temp_dir, ops::Range, sync::Arc, time::Duration};
use theme::ActiveTheme;
use ui::{IconButtonShape, Tooltip, prelude::*};
use util::ResultExt as _;

pub struct Session {
    fs: Arc<dyn Fs>,
    editor: WeakEntity<Editor>,
    pub kernel: Kernel,
    pub kernel_specification: KernelSpecification,

    blocks: HashMap<String, EditorBlock>,
    result_inlays: HashMap<String, (InlayId, Range<Anchor>, usize)>,
    next_inlay_id: usize,

    _subscriptions: Vec<Subscription>,
}

struct EditorBlock {
    code_range: Range<Anchor>,
    invalidation_anchor: Anchor,
    block_id: CustomBlockId,
    execution_view: Entity<ExecutionView>,
}

type CloseBlockFn =
    Arc<dyn for<'a> Fn(CustomBlockId, &'a mut Window, &mut App) + Send + Sync + 'static>;

impl EditorBlock {
    fn new(
        editor: WeakEntity<Editor>,
        code_range: Range<Anchor>,
        status: ExecutionStatus,
        on_close: CloseBlockFn,
        cx: &mut Context<Session>,
    ) -> anyhow::Result<Self> {
        let editor = editor.upgrade().context("editor is not open")?;
        let workspace = editor.read(cx).workspace().context("workspace dropped")?;

        let execution_view = cx.new(|cx| ExecutionView::new(status, workspace.downgrade(), cx));

        let (block_id, invalidation_anchor) = editor.update(cx, |editor, cx| {
            let buffer = editor.buffer().clone();
            let buffer_snapshot = buffer.read(cx).snapshot(cx);
            let end_point = code_range.end.to_point(&buffer_snapshot);
            let next_row_start = end_point + Point::new(1, 0);
            if next_row_start > buffer_snapshot.max_point() {
                buffer.update(cx, |buffer, cx| {
                    buffer.edit(
                        [(
                            buffer_snapshot.max_point()..buffer_snapshot.max_point(),
                            "\n",
                        )],
                        None,
                        cx,
                    )
                });
            }

            // Re-read snapshot after potential buffer edit and create a fresh anchor for
            // block placement. Using anchor_before (Bias::Left) ensures the anchor stays
            // at the end of the code line regardless of whether we inserted a newline.
            let buffer_snapshot = buffer.read(cx).snapshot(cx);
            let block_placement_anchor = buffer_snapshot.anchor_before(end_point);
            let invalidation_anchor = buffer_snapshot.anchor_before(next_row_start);
            let block = BlockProperties {
                placement: BlockPlacement::Below(block_placement_anchor),
                // Take up at least one height for status, allow the editor to determine the real height based on the content from render
                height: Some(1),
                style: BlockStyle::Sticky,
                render: Self::create_output_area_renderer(execution_view.clone(), on_close.clone()),
                priority: 0,
            };

            let block_id = editor.insert_blocks([block], None, cx)[0];
            (block_id, invalidation_anchor)
        });

        anyhow::Ok(Self {
            code_range,
            invalidation_anchor,
            block_id,
            execution_view,
        })
    }

    fn handle_message(
        &mut self,
        message: &JupyterMessage,
        window: &mut Window,
        cx: &mut Context<Session>,
    ) {
        self.execution_view.update(cx, |execution_view, cx| {
            if matches!(&message.content, JupyterMessageContent::InputRequest(_)) {
                execution_view.handle_input_request(message, window, cx);
            } else {
                execution_view.push_message(&message.content, window, cx);
            }
        });
    }

    fn create_output_area_renderer(
        execution_view: Entity<ExecutionView>,
        on_close: CloseBlockFn,
    ) -> RenderBlock {
        Arc::new(move |cx: &mut BlockContext| {
            let execution_view = execution_view.clone();
            let text_style = crate::outputs::plain::text_style(cx.window, cx.app);

            let editor_margins = cx.margins;
            let gutter = editor_margins.gutter;

            let block_id = cx.block_id;
            let on_close = on_close.clone();

            let rem_size = cx.window.rem_size();

            let text_line_height = text_style.line_height_in_pixels(rem_size);
            let output_settings = ReplSettings::get_global(cx.app);
            let output_max_height = if output_settings.output_max_height_lines > 0 {
                Some(text_line_height * output_settings.output_max_height_lines as f32)
            } else {
                None
            };

            let close_button = h_flex()
                .flex_none()
                .items_center()
                .justify_center()
                .absolute()
                .top(text_line_height / 2.)
                .right(
                    // 2px is a magic number to nudge the button just a bit closer to
                    // the line number start
                    gutter.full_width() / 2.0 - text_line_height / 2.0 - px(2.),
                )
                .w(text_line_height)
                .h(text_line_height)
                .child(
                    IconButton::new("close_output_area", IconName::Close)
                        .icon_size(IconSize::Small)
                        .icon_color(Color::Muted)
                        .size(ButtonSize::Compact)
                        .shape(IconButtonShape::Square)
                        .tooltip(Tooltip::text("Close output area"))
                        .on_click(move |_, window, cx| {
                            if let BlockId::Custom(block_id) = block_id {
                                (on_close)(block_id, window, cx)
                            }
                        }),
                );

            div()
                .id(cx.block_id)
                .block_mouse_except_scroll()
                .flex()
                .items_start()
                .min_h(text_line_height)
                .w_full()
                .border_y_1()
                .border_color(cx.theme().colors().border)
                .bg(cx.theme().colors().background)
                .child(
                    div()
                        .relative()
                        .w(gutter.full_width())
                        .h(text_line_height * 2)
                        .child(close_button),
                )
                .child(
                    div()
                        .id((ElementId::from(cx.block_id), "output-scroll"))
                        .flex_1()
                        .overflow_x_hidden()
                        .py(text_line_height / 2.)
                        .mr(editor_margins.right)
                        .pr_2()
                        .when_some(output_max_height, |div, max_h| {
                            div.max_h(max_h).overflow_y_scroll()
                        })
                        .child(execution_view),
                )
                .into_any_element()
        })
    }
}

impl Session {
    pub fn new(
        editor: WeakEntity<Editor>,
        fs: Arc<dyn Fs>,
        kernel_specification: KernelSpecification,
        window: &mut Window,
        cx: &mut Context<Self>,
    ) -> Self {
        let subscription = match editor.upgrade() {
            Some(editor) => {
                let buffer = editor.read(cx).buffer().clone();
                cx.subscribe(&buffer, Self::on_buffer_event)
            }
            None => Subscription::new(|| {}),
        };

        let editor_handle = editor.clone();

        editor
            .update(cx, |editor, _cx| {
                setup_editor_session_actions(editor, editor_handle);
            })
            .ok();

        let mut session = Self {
            fs,
            editor,
            kernel: Kernel::StartingKernel(Task::ready(()).shared()),
            blocks: HashMap::default(),
            result_inlays: HashMap::default(),
            next_inlay_id: 0,
            kernel_specification,
            _subscriptions: vec![subscription],
        };

        session.start_kernel(window, cx);
        session
    }

    fn start_kernel(&mut self, window: &mut Window, cx: &mut Context<Self>) {
        let kernel_language = self.kernel_specification.language();
        let entity_id = self.editor.entity_id();

        // For WSL Remote kernels, use project root instead of potentially temporary working directory
        // which causes .venv/bin/python checks to fail
        let is_remote_execution = matches!(
            self.kernel_specification,
            crate::KernelSpecification::WslRemote(_) | crate::KernelSpecification::SshRemote(_)
        );

        let working_directory = if is_remote_execution {
            // For WSL Remote kernels, use project root instead of potentially temporary working directory
            // which causes .venv/bin/python checks to fail
            self.editor
                .upgrade()
                .and_then(|editor| editor.read(cx).project().cloned())
                .and_then(|project| {
                    project
                        .read(cx)
                        .worktrees(cx)
                        .next()
                        .map(|worktree| worktree.read(cx).abs_path().to_path_buf())
                })
                .unwrap_or_else(temp_dir)
        } else {
            self.editor
                .upgrade()
                .and_then(|editor| editor.read(cx).working_directory(cx))
                .unwrap_or_else(temp_dir)
        };

        telemetry::event!(
            "Kernel Status Changed",
            kernel_language,
            kernel_status = KernelStatus::Starting.to_string(),
            repl_session_id = cx.entity_id().to_string(),
        );

        let session_view = cx.entity();

        let kernel = match self.kernel_specification.clone() {
            KernelSpecification::Jupyter(kernel_specification) => NativeRunningKernel::new(
                kernel_specification,
                entity_id,
                working_directory,
                self.fs.clone(),
                session_view,
                window,
                cx,
            ),
            KernelSpecification::PythonEnv(env_specification) => NativeRunningKernel::new(
                env_specification.as_local_spec(),
                entity_id,
                working_directory,
                self.fs.clone(),
                session_view,
                window,
                cx,
            ),
            KernelSpecification::JupyterServer(remote_kernel_specification) => {
                RemoteRunningKernel::new(
                    remote_kernel_specification,
                    working_directory,
                    session_view,
                    window,
                    cx,
                )
            }
            KernelSpecification::SshRemote(spec) => {
                let project = self
                    .editor
                    .upgrade()
                    .and_then(|editor| editor.read(cx).project().cloned());
                if let Some(project) = project {
                    SshRunningKernel::new(
                        spec,
                        working_directory,
                        project,
                        session_view,
                        window,
                        cx,
                    )
                } else {
                    Task::ready(Err(anyhow::anyhow!("No project associated with editor")))
                }
            }
            KernelSpecification::WslRemote(spec) => WslRunningKernel::new(
                spec,
                entity_id,
                working_directory,
                self.fs.clone(),
                session_view,
                window,
                cx,
            ),
        };

        let pending_kernel = cx
            .spawn(async move |this, cx| {
                let kernel: anyhow::Result<Box<dyn crate::kernels::RunningKernel>> = kernel.await;

                match kernel {
                    Ok(kernel) => {
                        this.update(cx, |session, cx| {
                            session.kernel(Kernel::RunningKernel(kernel), cx);
                            let request =
                                JupyterMessageContent::KernelInfoRequest(KernelInfoRequest {});
                            session.send(request.into(), cx).log_err();
                        })
                        .ok();
                    }
                    Err(err) => {
                        this.update(cx, |session, cx| {
                            session.kernel_errored(err.to_string(), cx);
                        })
                        .ok();
                    }
                }
            })
            .shared();

        self.kernel(Kernel::StartingKernel(pending_kernel), cx);
        cx.notify();
    }

    pub fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>) {
        self.kernel(Kernel::ErroredLaunch(error_message.clone()), cx);

        self.blocks.values().for_each(|block| {
            block.execution_view.update(cx, |execution_view, cx| {
                match execution_view.status {
                    ExecutionStatus::Finished => {
                        // Do nothing when the output was good
                    }
                    _ => {
                        // All other cases, set the status to errored
                        execution_view.status =
                            ExecutionStatus::KernelErrored(error_message.clone())
                    }
                }
                cx.notify();
            });
        });
    }

    fn on_buffer_event(
        &mut self,
        buffer: Entity<MultiBuffer>,
        event: &multi_buffer::Event,
        cx: &mut Context<Self>,
    ) {
        if let multi_buffer::Event::Edited { .. } = event {
            let snapshot = buffer.read(cx).snapshot(cx);

            let mut blocks_to_remove: HashSet<CustomBlockId> = HashSet::default();
            let mut gutter_ranges_to_remove: Vec<Range<Anchor>> = Vec::new();
            let mut keys_to_remove: Vec<String> = Vec::new();

            self.blocks.retain(|id, block| {
                if block.invalidation_anchor.is_valid(&snapshot) {
                    true
                } else {
                    blocks_to_remove.insert(block.block_id);
                    gutter_ranges_to_remove.push(block.code_range.clone());
                    keys_to_remove.push(id.clone());
                    false
                }
            });

            let mut inlays_to_remove: Vec<InlayId> = Vec::new();

            self.result_inlays
                .retain(|id, (inlay_id, code_range, original_len)| {
                    let start_offset = code_range.start.to_offset(&snapshot);
                    let end_offset = code_range.end.to_offset(&snapshot);
                    let current_len = end_offset.saturating_sub(start_offset);

                    if current_len != *original_len {
                        inlays_to_remove.push(*inlay_id);
                        gutter_ranges_to_remove.push(code_range.clone());
                        keys_to_remove.push(id.clone());
                        false
                    } else {
                        true
                    }
                });

            if !blocks_to_remove.is_empty()
                || !inlays_to_remove.is_empty()
                || !gutter_ranges_to_remove.is_empty()
            {
                self.editor
                    .update(cx, |editor, cx| {
                        if !blocks_to_remove.is_empty() {
                            editor.remove_blocks(blocks_to_remove, None, cx);
                        }
                        if !inlays_to_remove.is_empty() {
                            editor.splice_inlays(&inlays_to_remove, vec![], cx);
                        }
                        if !gutter_ranges_to_remove.is_empty() {
                            editor.remove_gutter_highlights::<ReplExecutedRange>(
                                gutter_ranges_to_remove,
                                cx,
                            );
                        }
                    })
                    .ok();
                cx.notify();
            }
        }
    }

    fn send(&mut self, message: JupyterMessage, _cx: &mut Context<Self>) -> anyhow::Result<()> {
        if let Kernel::RunningKernel(kernel) = &mut self.kernel {
            kernel.request_tx().try_send(message).ok();
        }

        anyhow::Ok(())
    }

    fn send_stdin_reply(
        &mut self,
        value: String,
        parent_message: &JupyterMessage,
        _cx: &mut Context<Self>,
    ) {
        if let Kernel::RunningKernel(kernel) = &mut self.kernel {
            let reply = InputReply {
                value,
                status: ReplyStatus::Ok,
                error: None,
            };
            let message = reply.as_child_of(parent_message);
            kernel.stdin_tx().try_send(message).log_err();
        }
    }

    fn replace_block_with_inlay(&mut self, message_id: &str, text: &str, cx: &mut Context<Self>) {
        let Some(block) = self.blocks.remove(message_id) else {
            return;
        };

        let Some(editor) = self.editor.upgrade() else {
            return;
        };

        let code_range = block.code_range.clone();

        editor.update(cx, |editor, cx| {
            let mut block_ids = HashSet::default();
            block_ids.insert(block.block_id);
            editor.remove_blocks(block_ids, None, cx);

            let buffer = editor.buffer().read(cx).snapshot(cx);
            let start_offset = code_range.start.to_offset(&buffer);
            let end_offset = code_range.end.to_offset(&buffer);
            let original_len = end_offset.saturating_sub(start_offset);

            let end_point = code_range.end.to_point(&buffer);
            let inlay_position = buffer.anchor_after(end_point);

            let inlay_id = self.next_inlay_id;
            self.next_inlay_id += 1;

            let inlay = Inlay::repl_result(inlay_id, inlay_position, format!("    {}", text));

            editor.splice_inlays(&[], vec![inlay], cx);
            self.result_inlays.insert(
                message_id.to_string(),
                (
                    InlayId::ReplResult(inlay_id),
                    code_range.clone(),
                    original_len,
                ),
            );

            editor.insert_gutter_highlight::<ReplExecutedRange>(
                code_range,
                |cx| cx.theme().status().success,
                cx,
            );
        });

        cx.notify();
    }

    pub fn clear_outputs(&mut self, cx: &mut Context<Self>) {
        let blocks_to_remove: HashSet<CustomBlockId> =
            self.blocks.values().map(|block| block.block_id).collect();

        let inlays_to_remove: Vec<InlayId> =
            self.result_inlays.values().map(|(id, _, _)| *id).collect();

        self.editor
            .update(cx, |editor, cx| {
                editor.remove_blocks(blocks_to_remove, None, cx);
                editor.splice_inlays(&inlays_to_remove, vec![], cx);
                editor.clear_gutter_highlights::<ReplExecutedRange>(cx);
            })
            .ok();

        self.blocks.clear();
        self.result_inlays.clear();
    }

    pub fn clear_output_at_position(&mut self, position: Anchor, cx: &mut Context<Self>) {
        let Some(editor) = self.editor.upgrade() else {
            return;
        };

        let (block_id, code_range, msg_id) = {
            let snapshot = editor.read(cx).buffer().read(cx).read(cx);
            let pos_range = position..position;

            let block_to_remove = self
                .blocks
                .iter()
                .find(|(_, block)| block.code_range.includes(&pos_range, &snapshot));

            let Some((msg_id, block)) = block_to_remove else {
                return;
            };

            (block.block_id, block.code_range.clone(), msg_id.clone())
        };

        let inlay_to_remove = self.result_inlays.get(&msg_id).map(|(id, _, _)| *id);

        self.blocks.remove(&msg_id);
        if inlay_to_remove.is_some() {
            self.result_inlays.remove(&msg_id);
        }

        self.editor
            .update(cx, |editor, cx| {
                let mut block_ids = HashSet::default();
                block_ids.insert(block_id);
                editor.remove_blocks(block_ids, None, cx);

                if let Some(inlay_id) = inlay_to_remove {
                    editor.splice_inlays(&[inlay_id], vec![], cx);
                }

                editor.remove_gutter_highlights::<ReplExecutedRange>(vec![code_range], cx);
            })
            .ok();

        cx.notify();
    }

    pub fn execute(
        &mut self,
        code: String,
        anchor_range: Range<Anchor>,
        next_cell: Option<Anchor>,
        move_down: bool,
        window: &mut Window,
        cx: &mut Context<Self>,
    ) {
        let Some(editor) = self.editor.upgrade() else {
            return;
        };

        if code.is_empty() {
            return;
        }

        let execute_request = ExecuteRequest {
            code,
            allow_stdin: true,
            ..ExecuteRequest::default()
        };

        let message: JupyterMessage = execute_request.into();

        let mut blocks_to_remove: HashSet<CustomBlockId> = HashSet::default();
        let mut inlays_to_remove: Vec<InlayId> = Vec::new();
        let mut gutter_ranges_to_remove: Vec<Range<Anchor>> = Vec::new();

        let buffer = editor.read(cx).buffer().read(cx).snapshot(cx);

        self.blocks.retain(|_key, block| {
            if anchor_range.overlaps(&block.code_range, &buffer) {
                blocks_to_remove.insert(block.block_id);
                false
            } else {
                true
            }
        });

        self.result_inlays
            .retain(|_key, (inlay_id, inlay_range, _)| {
                if anchor_range.overlaps(inlay_range, &buffer) {
                    inlays_to_remove.push(*inlay_id);
                    gutter_ranges_to_remove.push(inlay_range.clone());
                    false
                } else {
                    true
                }
            });

        self.editor
            .update(cx, |editor, cx| {
                editor.remove_blocks(blocks_to_remove, None, cx);
                if !inlays_to_remove.is_empty() {
                    editor.splice_inlays(&inlays_to_remove, vec![], cx);
                }
                if !gutter_ranges_to_remove.is_empty() {
                    editor
                        .remove_gutter_highlights::<ReplExecutedRange>(gutter_ranges_to_remove, cx);
                }
            })
            .ok();

        let status = match &self.kernel {
            Kernel::Restarting => ExecutionStatus::Restarting,
            Kernel::RunningKernel(_) => ExecutionStatus::Queued,
            Kernel::StartingKernel(_) => ExecutionStatus::ConnectingToKernel,
            Kernel::ErroredLaunch(error) => ExecutionStatus::KernelErrored(error.clone()),
            Kernel::ShuttingDown => ExecutionStatus::ShuttingDown,
            Kernel::Shutdown => ExecutionStatus::Shutdown,
        };

        let parent_message_id = message.header.msg_id.clone();
        let session_view = cx.entity().downgrade();
        let weak_editor = self.editor.clone();
        let code_range_for_close = anchor_range.clone();

        let on_close: CloseBlockFn = Arc::new(
            move |block_id: CustomBlockId, _: &mut Window, cx: &mut App| {
                if let Some(session) = session_view.upgrade() {
                    session.update(cx, |session, cx| {
                        session.blocks.remove(&parent_message_id);
                        cx.notify();
                    });
                }

                if let Some(editor) = weak_editor.upgrade() {
                    editor.update(cx, |editor, cx| {
                        let mut block_ids = HashSet::default();
                        block_ids.insert(block_id);
                        editor.remove_blocks(block_ids, None, cx);
                        editor.remove_gutter_highlights::<ReplExecutedRange>(
                            vec![code_range_for_close.clone()],
                            cx,
                        );
                    });
                }
            },
        );

        let Ok(editor_block) = EditorBlock::new(
            self.editor.clone(),
            anchor_range.clone(),
            status,
            on_close,
            cx,
        ) else {
            return;
        };

        self.editor
            .update(cx, |editor, cx| {
                editor.insert_gutter_highlight::<ReplExecutedRange>(
                    anchor_range.clone(),
                    |cx| cx.theme().status().success,
                    cx,
                );
            })
            .ok();

        let new_cursor_pos = if let Some(next_cursor) = next_cell {
            next_cursor
        } else {
            editor_block.invalidation_anchor
        };

        let msg_id = message.header.msg_id.clone();
        let subscription = cx.subscribe(
            &editor_block.execution_view,
            move |session, _execution_view, _event: &ExecutionViewFinishedEmpty, cx| {
                session.replace_block_with_inlay(&msg_id, "✓", cx);
            },
        );
        self._subscriptions.push(subscription);

        let msg_id = message.header.msg_id.clone();
        let subscription = cx.subscribe(
            &editor_block.execution_view,
            move |session, _execution_view, event: &ExecutionViewFinishedSmall, cx| {
                session.replace_block_with_inlay(&msg_id, &event.0, cx);
            },
        );
        self._subscriptions.push(subscription);

        let subscription = cx.subscribe(
            &editor_block.execution_view,
            |session, _execution_view, event: &InputReplyEvent, cx| {
                session.send_stdin_reply(event.value.clone(), &event.parent_message, cx);
            },
        );
        self._subscriptions.push(subscription);

        self.blocks
            .insert(message.header.msg_id.clone(), editor_block);

        match &self.kernel {
            Kernel::RunningKernel(_) => {
                self.send(message, cx).ok();
            }
            Kernel::StartingKernel(task) => {
                // Queue up the execution as a task to run after the kernel starts
                let task = task.clone();

                cx.spawn(async move |this, cx| {
                    task.await;
                    this.update(cx, |session, cx| {
                        session.send(message, cx).ok();
                    })
                    .ok();
                })
                .detach();
            }
            _ => {}
        }

        if move_down {
            editor.update(cx, move |editor, cx| {
                editor.change_selections(
                    SelectionEffects::scroll(Autoscroll::top_relative(8.0)),
                    window,
                    cx,
                    |selections| {
                        selections.select_ranges([new_cursor_pos..new_cursor_pos]);
                    },
                );
            });
        }
    }

    pub fn interrupt(&mut self, cx: &mut Context<Self>) {
        match &mut self.kernel {
            Kernel::RunningKernel(_kernel) => {
                self.send(InterruptRequest {}.into(), cx).ok();
            }
            Kernel::StartingKernel(_task) => {
                // NOTE: If we switch to a literal queue instead of chaining on to the task, clear all queued executions
            }
            _ => {}
        }
    }

    pub fn kernel(&mut self, kernel: Kernel, cx: &mut Context<Self>) {
        if let Kernel::Shutdown = kernel {
            cx.emit(SessionEvent::Shutdown(self.editor.clone()));
        }

        let kernel_status = KernelStatus::from(&kernel).to_string();
        let kernel_language = self.kernel_specification.language();

        telemetry::event!(
            "Kernel Status Changed",
            kernel_language,
            kernel_status,
            repl_session_id = cx.entity_id().to_string(),
        );

        self.kernel = kernel;
    }

    pub fn shutdown(&mut self, window: &mut Window, cx: &mut Context<Self>) {
        let kernel = std::mem::replace(&mut self.kernel, Kernel::ShuttingDown);

        match kernel {
            Kernel::RunningKernel(mut kernel) => {
                let mut request_tx = kernel.request_tx();

                let forced = kernel.force_shutdown(window, cx);

                cx.spawn(async move |this, cx| {
                    let message: JupyterMessage = ShutdownRequest { restart: false }.into();
                    request_tx.try_send(message).ok();

                    forced.await.log_err();

                    // Give the kernel a bit of time to clean up
                    cx.background_executor().timer(Duration::from_secs(3)).await;

                    this.update(cx, |session, cx| {
                        session.clear_outputs(cx);
                        session.kernel(Kernel::Shutdown, cx);
                        cx.notify();
                    })
                    .ok();
                })
                .detach();
            }
            _ => {
                self.kernel(Kernel::Shutdown, cx);
            }
        }
        cx.notify();
    }

    pub fn restart(&mut self, window: &mut Window, cx: &mut Context<Self>) {
        let kernel = std::mem::replace(&mut self.kernel, Kernel::Restarting);

        match kernel {
            Kernel::Restarting => {
                // Do nothing if already restarting
            }
            Kernel::RunningKernel(mut kernel) => {
                let mut request_tx = kernel.request_tx();

                let forced = kernel.force_shutdown(window, cx);

                cx.spawn_in(window, async move |this, cx| {
                    // Send shutdown request with restart flag
                    log::debug!("restarting kernel");
                    let message: JupyterMessage = ShutdownRequest { restart: true }.into();
                    request_tx.try_send(message).ok();

                    // Wait for kernel to shutdown
                    cx.background_executor().timer(Duration::from_secs(1)).await;

                    // Force kill the kernel if it hasn't shut down
                    forced.await.log_err();

                    // Start a new kernel
                    this.update_in(cx, |session, window, cx| {
                        // TODO: Differentiate between restart and restart+clear-outputs
                        session.clear_outputs(cx);
                        session.start_kernel(window, cx);
                    })
                    .ok();
                })
                .detach();
            }
            _ => {
                self.clear_outputs(cx);
                self.start_kernel(window, cx);
            }
        }
        cx.notify();
    }
}

pub enum SessionEvent {
    Shutdown(WeakEntity<Editor>),
}

impl EventEmitter<SessionEvent> for Session {}

impl Render for Session {
    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
        let (status_text, interrupt_button) = match &self.kernel {
            Kernel::RunningKernel(kernel) => (
                kernel
                    .kernel_info()
                    .as_ref()
                    .map(|info| info.language_info.name.clone()),
                Some(
                    Button::new("interrupt", "Interrupt")
                        .style(ButtonStyle::Subtle)
                        .on_click(cx.listener(move |session, _, _, cx| {
                            session.interrupt(cx);
                        })),
                ),
            ),
            Kernel::StartingKernel(_) => (Some("Starting".into()), None),
            Kernel::ErroredLaunch(err) => (Some(format!("Error: {err}")), None),
            Kernel::ShuttingDown => (Some("Shutting Down".into()), None),
            Kernel::Shutdown => (Some("Shutdown".into()), None),
            Kernel::Restarting => (Some("Restarting".into()), None),
        };

        KernelListItem::new(self.kernel_specification.clone())
            .status_color(match &self.kernel {
                Kernel::RunningKernel(kernel) => match kernel.execution_state() {
                    ExecutionState::Idle => Color::Success,
                    ExecutionState::Busy => Color::Modified,
                    ExecutionState::Unknown => Color::Modified,
                    ExecutionState::Starting => Color::Modified,
                    ExecutionState::Restarting => Color::Modified,
                    ExecutionState::Terminating => Color::Disabled,
                    ExecutionState::AutoRestarting => Color::Modified,
                    ExecutionState::Dead => Color::Disabled,
                    ExecutionState::Other(_) => Color::Modified,
                },
                Kernel::StartingKernel(_) => Color::Modified,
                Kernel::ErroredLaunch(_) => Color::Error,
                Kernel::ShuttingDown => Color::Modified,
                Kernel::Shutdown => Color::Disabled,
                Kernel::Restarting => Color::Modified,
            })
            .child(Label::new(self.kernel_specification.name()))
            .children(status_text.map(|status_text| Label::new(format!("({status_text})"))))
            .button(
                Button::new("shutdown", "Shutdown")
                    .style(ButtonStyle::Subtle)
                    .disabled(self.kernel.is_shutting_down())
                    .on_click(cx.listener(move |session, _, window, cx| {
                        session.shutdown(window, cx);
                    })),
            )
            .buttons(interrupt_button)
    }
}

impl KernelSession for Session {
    fn route(&mut self, message: &JupyterMessage, window: &mut Window, cx: &mut Context<Self>) {
        let parent_message_id = match message.parent_header.as_ref() {
            Some(header) => &header.msg_id,
            None => return,
        };

        match &message.content {
            JupyterMessageContent::Status(status) => {
                self.kernel.set_execution_state(&status.execution_state);

                telemetry::event!(
                    "Kernel Status Changed",
                    kernel_language = self.kernel_specification.language(),
                    kernel_status = KernelStatus::from(&self.kernel).to_string(),
                    repl_session_id = cx.entity_id().to_string(),
                );

                cx.notify();
            }
            JupyterMessageContent::KernelInfoReply(reply) => {
                self.kernel.set_kernel_info(reply);
                cx.notify();
            }
            JupyterMessageContent::UpdateDisplayData(update) => {
                let display_id = if let Some(display_id) = update.transient.display_id.clone() {
                    display_id
                } else {
                    return;
                };

                self.blocks.iter_mut().for_each(|(_, block)| {
                    block.execution_view.update(cx, |execution_view, cx| {
                        execution_view.update_display_data(&update.data, &display_id, window, cx);
                    });
                });
                return;
            }
            _ => {}
        }

        if let Some(block) = self.blocks.get_mut(parent_message_id) {
            block.handle_message(message, window, cx);
        }
    }

    fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>) {
        self.kernel_errored(error_message, cx);
    }
}