1use crate::components::KernelListItem;
  2use crate::kernels::RemoteRunningKernel;
  3use crate::setup_editor_session_actions;
  4use crate::{
  5    KernelStatus,
  6    kernels::{Kernel, KernelSession, KernelSpecification, NativeRunningKernel},
  7    outputs::{
  8        ExecutionStatus, ExecutionView, ExecutionViewFinishedEmpty, ExecutionViewFinishedSmall,
  9    },
 10    repl_settings::ReplSettings,
 11};
 12use anyhow::Context as _;
 13use collections::{HashMap, HashSet};
 14use editor::SelectionEffects;
 15use editor::{
 16    Anchor, AnchorRangeExt as _, Editor, Inlay, MultiBuffer, ToOffset, ToPoint,
 17    display_map::{
 18        BlockContext, BlockId, BlockPlacement, BlockProperties, BlockStyle, CustomBlockId,
 19        RenderBlock,
 20    },
 21    scroll::Autoscroll,
 22};
 23use project::InlayId;
 24
 25/// Marker types
 26enum ReplExecutedRange {}
 27
 28use futures::FutureExt as _;
 29use gpui::{
 30    Context, Entity, EventEmitter, Render, Subscription, Task, WeakEntity, Window, div, prelude::*,
 31};
 32use language::Point;
 33use project::Fs;
 34use runtimelib::{
 35    ExecuteRequest, ExecutionState, InterruptRequest, JupyterMessage, JupyterMessageContent,
 36    ShutdownRequest,
 37};
 38use settings::Settings as _;
 39use std::{env::temp_dir, ops::Range, sync::Arc, time::Duration};
 40use theme::ActiveTheme;
 41use ui::{IconButtonShape, Tooltip, prelude::*};
 42use util::ResultExt as _;
 43
 44pub struct Session {
 45    fs: Arc<dyn Fs>,
 46    editor: WeakEntity<Editor>,
 47    pub kernel: Kernel,
 48    pub kernel_specification: KernelSpecification,
 49
 50    blocks: HashMap<String, EditorBlock>,
 51    result_inlays: HashMap<String, (InlayId, Range<Anchor>, usize)>,
 52    next_inlay_id: usize,
 53
 54    _subscriptions: Vec<Subscription>,
 55}
 56
 57struct EditorBlock {
 58    code_range: Range<Anchor>,
 59    invalidation_anchor: Anchor,
 60    block_id: CustomBlockId,
 61    execution_view: Entity<ExecutionView>,
 62}
 63
 64type CloseBlockFn =
 65    Arc<dyn for<'a> Fn(CustomBlockId, &'a mut Window, &mut App) + Send + Sync + 'static>;
 66
 67impl EditorBlock {
 68    fn new(
 69        editor: WeakEntity<Editor>,
 70        code_range: Range<Anchor>,
 71        status: ExecutionStatus,
 72        on_close: CloseBlockFn,
 73        cx: &mut Context<Session>,
 74    ) -> anyhow::Result<Self> {
 75        let editor = editor.upgrade().context("editor is not open")?;
 76        let workspace = editor.read(cx).workspace().context("workspace dropped")?;
 77
 78        let execution_view = cx.new(|cx| ExecutionView::new(status, workspace.downgrade(), cx));
 79
 80        let (block_id, invalidation_anchor) = editor.update(cx, |editor, cx| {
 81            let buffer = editor.buffer().clone();
 82            let buffer_snapshot = buffer.read(cx).snapshot(cx);
 83            let end_point = code_range.end.to_point(&buffer_snapshot);
 84            let next_row_start = end_point + Point::new(1, 0);
 85            if next_row_start > buffer_snapshot.max_point() {
 86                buffer.update(cx, |buffer, cx| {
 87                    buffer.edit(
 88                        [(
 89                            buffer_snapshot.max_point()..buffer_snapshot.max_point(),
 90                            "\n",
 91                        )],
 92                        None,
 93                        cx,
 94                    )
 95                });
 96            }
 97
 98            // Re-read snapshot after potential buffer edit and create a fresh anchor for
 99            // block placement. Using anchor_before (Bias::Left) ensures the anchor stays
100            // at the end of the code line regardless of whether we inserted a newline.
101            let buffer_snapshot = buffer.read(cx).snapshot(cx);
102            let block_placement_anchor = buffer_snapshot.anchor_before(end_point);
103            let invalidation_anchor = buffer_snapshot.anchor_before(next_row_start);
104            let block = BlockProperties {
105                placement: BlockPlacement::Below(block_placement_anchor),
106                // Take up at least one height for status, allow the editor to determine the real height based on the content from render
107                height: Some(1),
108                style: BlockStyle::Sticky,
109                render: Self::create_output_area_renderer(execution_view.clone(), on_close.clone()),
110                priority: 0,
111            };
112
113            let block_id = editor.insert_blocks([block], None, cx)[0];
114            (block_id, invalidation_anchor)
115        });
116
117        anyhow::Ok(Self {
118            code_range,
119            invalidation_anchor,
120            block_id,
121            execution_view,
122        })
123    }
124
125    fn handle_message(
126        &mut self,
127        message: &JupyterMessage,
128        window: &mut Window,
129        cx: &mut Context<Session>,
130    ) {
131        self.execution_view.update(cx, |execution_view, cx| {
132            execution_view.push_message(&message.content, window, cx);
133        });
134    }
135
136    fn create_output_area_renderer(
137        execution_view: Entity<ExecutionView>,
138        on_close: CloseBlockFn,
139    ) -> RenderBlock {
140        Arc::new(move |cx: &mut BlockContext| {
141            let execution_view = execution_view.clone();
142            let text_style = crate::outputs::plain::text_style(cx.window, cx.app);
143
144            let editor_margins = cx.margins;
145            let gutter = editor_margins.gutter;
146
147            let block_id = cx.block_id;
148            let on_close = on_close.clone();
149
150            let rem_size = cx.window.rem_size();
151
152            let text_line_height = text_style.line_height_in_pixels(rem_size);
153            let output_settings = ReplSettings::get_global(cx.app);
154            let output_max_height = if output_settings.output_max_height_lines > 0 {
155                Some(text_line_height * output_settings.output_max_height_lines as f32)
156            } else {
157                None
158            };
159
160            let close_button = h_flex()
161                .flex_none()
162                .items_center()
163                .justify_center()
164                .absolute()
165                .top(text_line_height / 2.)
166                .right(
167                    // 2px is a magic number to nudge the button just a bit closer to
168                    // the line number start
169                    gutter.full_width() / 2.0 - text_line_height / 2.0 - px(2.),
170                )
171                .w(text_line_height)
172                .h(text_line_height)
173                .child(
174                    IconButton::new("close_output_area", IconName::Close)
175                        .icon_size(IconSize::Small)
176                        .icon_color(Color::Muted)
177                        .size(ButtonSize::Compact)
178                        .shape(IconButtonShape::Square)
179                        .tooltip(Tooltip::text("Close output area"))
180                        .on_click(move |_, window, cx| {
181                            if let BlockId::Custom(block_id) = block_id {
182                                (on_close)(block_id, window, cx)
183                            }
184                        }),
185                );
186
187            div()
188                .id(cx.block_id)
189                .block_mouse_except_scroll()
190                .flex()
191                .items_start()
192                .min_h(text_line_height)
193                .w_full()
194                .border_y_1()
195                .border_color(cx.theme().colors().border)
196                .bg(cx.theme().colors().background)
197                .child(
198                    div()
199                        .relative()
200                        .w(gutter.full_width())
201                        .h(text_line_height * 2)
202                        .child(close_button),
203                )
204                .child(
205                    div()
206                        .id((ElementId::from(cx.block_id), "output-scroll"))
207                        .flex_1()
208                        .overflow_x_hidden()
209                        .py(text_line_height / 2.)
210                        .mr(editor_margins.right)
211                        .pr_2()
212                        .when_some(output_max_height, |div, max_h| {
213                            div.max_h(max_h).overflow_y_scroll()
214                        })
215                        .child(execution_view),
216                )
217                .into_any_element()
218        })
219    }
220}
221
222impl Session {
223    pub fn new(
224        editor: WeakEntity<Editor>,
225        fs: Arc<dyn Fs>,
226        kernel_specification: KernelSpecification,
227        window: &mut Window,
228        cx: &mut Context<Self>,
229    ) -> Self {
230        let subscription = match editor.upgrade() {
231            Some(editor) => {
232                let buffer = editor.read(cx).buffer().clone();
233                cx.subscribe(&buffer, Self::on_buffer_event)
234            }
235            None => Subscription::new(|| {}),
236        };
237
238        let editor_handle = editor.clone();
239
240        editor
241            .update(cx, |editor, _cx| {
242                setup_editor_session_actions(editor, editor_handle);
243            })
244            .ok();
245
246        let mut session = Self {
247            fs,
248            editor,
249            kernel: Kernel::StartingKernel(Task::ready(()).shared()),
250            blocks: HashMap::default(),
251            result_inlays: HashMap::default(),
252            next_inlay_id: 0,
253            kernel_specification,
254            _subscriptions: vec![subscription],
255        };
256
257        session.start_kernel(window, cx);
258        session
259    }
260
261    fn start_kernel(&mut self, window: &mut Window, cx: &mut Context<Self>) {
262        let kernel_language = self.kernel_specification.language();
263        let entity_id = self.editor.entity_id();
264        let working_directory = self
265            .editor
266            .upgrade()
267            .and_then(|editor| editor.read(cx).working_directory(cx))
268            .unwrap_or_else(temp_dir);
269
270        telemetry::event!(
271            "Kernel Status Changed",
272            kernel_language,
273            kernel_status = KernelStatus::Starting.to_string(),
274            repl_session_id = cx.entity_id().to_string(),
275        );
276
277        let session_view = cx.entity();
278
279        let kernel = match self.kernel_specification.clone() {
280            KernelSpecification::Jupyter(kernel_specification) => NativeRunningKernel::new(
281                kernel_specification,
282                entity_id,
283                working_directory,
284                self.fs.clone(),
285                session_view,
286                window,
287                cx,
288            ),
289            KernelSpecification::PythonEnv(env_specification) => NativeRunningKernel::new(
290                env_specification.as_local_spec(),
291                entity_id,
292                working_directory,
293                self.fs.clone(),
294                session_view,
295                window,
296                cx,
297            ),
298            KernelSpecification::Remote(remote_kernel_specification) => RemoteRunningKernel::new(
299                remote_kernel_specification,
300                working_directory,
301                session_view,
302                window,
303                cx,
304            ),
305        };
306
307        let pending_kernel = cx
308            .spawn(async move |this, cx| {
309                let kernel = kernel.await;
310
311                match kernel {
312                    Ok(kernel) => {
313                        this.update(cx, |session, cx| {
314                            session.kernel(Kernel::RunningKernel(kernel), cx);
315                        })
316                        .ok();
317                    }
318                    Err(err) => {
319                        this.update(cx, |session, cx| {
320                            session.kernel_errored(err.to_string(), cx);
321                        })
322                        .ok();
323                    }
324                }
325            })
326            .shared();
327
328        self.kernel(Kernel::StartingKernel(pending_kernel), cx);
329        cx.notify();
330    }
331
332    pub fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>) {
333        self.kernel(Kernel::ErroredLaunch(error_message.clone()), cx);
334
335        self.blocks.values().for_each(|block| {
336            block.execution_view.update(cx, |execution_view, cx| {
337                match execution_view.status {
338                    ExecutionStatus::Finished => {
339                        // Do nothing when the output was good
340                    }
341                    _ => {
342                        // All other cases, set the status to errored
343                        execution_view.status =
344                            ExecutionStatus::KernelErrored(error_message.clone())
345                    }
346                }
347                cx.notify();
348            });
349        });
350    }
351
352    fn on_buffer_event(
353        &mut self,
354        buffer: Entity<MultiBuffer>,
355        event: &multi_buffer::Event,
356        cx: &mut Context<Self>,
357    ) {
358        if let multi_buffer::Event::Edited { .. } = event {
359            let snapshot = buffer.read(cx).snapshot(cx);
360
361            let mut blocks_to_remove: HashSet<CustomBlockId> = HashSet::default();
362            let mut gutter_ranges_to_remove: Vec<Range<Anchor>> = Vec::new();
363            let mut keys_to_remove: Vec<String> = Vec::new();
364
365            self.blocks.retain(|id, block| {
366                if block.invalidation_anchor.is_valid(&snapshot) {
367                    true
368                } else {
369                    blocks_to_remove.insert(block.block_id);
370                    gutter_ranges_to_remove.push(block.code_range.clone());
371                    keys_to_remove.push(id.clone());
372                    false
373                }
374            });
375
376            let mut inlays_to_remove: Vec<InlayId> = Vec::new();
377
378            self.result_inlays
379                .retain(|id, (inlay_id, code_range, original_len)| {
380                    let start_offset = code_range.start.to_offset(&snapshot);
381                    let end_offset = code_range.end.to_offset(&snapshot);
382                    let current_len = end_offset.saturating_sub(start_offset);
383
384                    if current_len != *original_len {
385                        inlays_to_remove.push(*inlay_id);
386                        gutter_ranges_to_remove.push(code_range.clone());
387                        keys_to_remove.push(id.clone());
388                        false
389                    } else {
390                        true
391                    }
392                });
393
394            if !blocks_to_remove.is_empty()
395                || !inlays_to_remove.is_empty()
396                || !gutter_ranges_to_remove.is_empty()
397            {
398                self.editor
399                    .update(cx, |editor, cx| {
400                        if !blocks_to_remove.is_empty() {
401                            editor.remove_blocks(blocks_to_remove, None, cx);
402                        }
403                        if !inlays_to_remove.is_empty() {
404                            editor.splice_inlays(&inlays_to_remove, vec![], cx);
405                        }
406                        if !gutter_ranges_to_remove.is_empty() {
407                            editor.remove_gutter_highlights::<ReplExecutedRange>(
408                                gutter_ranges_to_remove,
409                                cx,
410                            );
411                        }
412                    })
413                    .ok();
414                cx.notify();
415            }
416        }
417    }
418
419    fn send(&mut self, message: JupyterMessage, _cx: &mut Context<Self>) -> anyhow::Result<()> {
420        if let Kernel::RunningKernel(kernel) = &mut self.kernel {
421            kernel.request_tx().try_send(message).ok();
422        }
423
424        anyhow::Ok(())
425    }
426
427    fn replace_block_with_inlay(&mut self, message_id: &str, text: &str, cx: &mut Context<Self>) {
428        let Some(block) = self.blocks.remove(message_id) else {
429            return;
430        };
431
432        let Some(editor) = self.editor.upgrade() else {
433            return;
434        };
435
436        let code_range = block.code_range.clone();
437
438        editor.update(cx, |editor, cx| {
439            let mut block_ids = HashSet::default();
440            block_ids.insert(block.block_id);
441            editor.remove_blocks(block_ids, None, cx);
442
443            let buffer = editor.buffer().read(cx).snapshot(cx);
444            let start_offset = code_range.start.to_offset(&buffer);
445            let end_offset = code_range.end.to_offset(&buffer);
446            let original_len = end_offset.saturating_sub(start_offset);
447
448            let end_point = code_range.end.to_point(&buffer);
449            let inlay_position = buffer.anchor_after(end_point);
450
451            let inlay_id = self.next_inlay_id;
452            self.next_inlay_id += 1;
453
454            let inlay = Inlay::repl_result(inlay_id, inlay_position, format!("    {}", text));
455
456            editor.splice_inlays(&[], vec![inlay], cx);
457            self.result_inlays.insert(
458                message_id.to_string(),
459                (
460                    InlayId::ReplResult(inlay_id),
461                    code_range.clone(),
462                    original_len,
463                ),
464            );
465
466            editor.insert_gutter_highlight::<ReplExecutedRange>(
467                code_range,
468                |cx| cx.theme().status().success,
469                cx,
470            );
471        });
472
473        cx.notify();
474    }
475
476    pub fn clear_outputs(&mut self, cx: &mut Context<Self>) {
477        let blocks_to_remove: HashSet<CustomBlockId> =
478            self.blocks.values().map(|block| block.block_id).collect();
479
480        let inlays_to_remove: Vec<InlayId> =
481            self.result_inlays.values().map(|(id, _, _)| *id).collect();
482
483        self.editor
484            .update(cx, |editor, cx| {
485                editor.remove_blocks(blocks_to_remove, None, cx);
486                editor.splice_inlays(&inlays_to_remove, vec![], cx);
487                editor.clear_gutter_highlights::<ReplExecutedRange>(cx);
488            })
489            .ok();
490
491        self.blocks.clear();
492        self.result_inlays.clear();
493    }
494
495    pub fn execute(
496        &mut self,
497        code: String,
498        anchor_range: Range<Anchor>,
499        next_cell: Option<Anchor>,
500        move_down: bool,
501        window: &mut Window,
502        cx: &mut Context<Self>,
503    ) {
504        let Some(editor) = self.editor.upgrade() else {
505            return;
506        };
507
508        if code.is_empty() {
509            return;
510        }
511
512        let execute_request = ExecuteRequest {
513            code,
514            ..ExecuteRequest::default()
515        };
516
517        let message: JupyterMessage = execute_request.into();
518
519        let mut blocks_to_remove: HashSet<CustomBlockId> = HashSet::default();
520        let mut inlays_to_remove: Vec<InlayId> = Vec::new();
521        let mut gutter_ranges_to_remove: Vec<Range<Anchor>> = Vec::new();
522
523        let buffer = editor.read(cx).buffer().read(cx).snapshot(cx);
524
525        self.blocks.retain(|_key, block| {
526            if anchor_range.overlaps(&block.code_range, &buffer) {
527                blocks_to_remove.insert(block.block_id);
528                false
529            } else {
530                true
531            }
532        });
533
534        self.result_inlays
535            .retain(|_key, (inlay_id, inlay_range, _)| {
536                if anchor_range.overlaps(inlay_range, &buffer) {
537                    inlays_to_remove.push(*inlay_id);
538                    gutter_ranges_to_remove.push(inlay_range.clone());
539                    false
540                } else {
541                    true
542                }
543            });
544
545        self.editor
546            .update(cx, |editor, cx| {
547                editor.remove_blocks(blocks_to_remove, None, cx);
548                if !inlays_to_remove.is_empty() {
549                    editor.splice_inlays(&inlays_to_remove, vec![], cx);
550                }
551                if !gutter_ranges_to_remove.is_empty() {
552                    editor
553                        .remove_gutter_highlights::<ReplExecutedRange>(gutter_ranges_to_remove, cx);
554                }
555            })
556            .ok();
557
558        let status = match &self.kernel {
559            Kernel::Restarting => ExecutionStatus::Restarting,
560            Kernel::RunningKernel(_) => ExecutionStatus::Queued,
561            Kernel::StartingKernel(_) => ExecutionStatus::ConnectingToKernel,
562            Kernel::ErroredLaunch(error) => ExecutionStatus::KernelErrored(error.clone()),
563            Kernel::ShuttingDown => ExecutionStatus::ShuttingDown,
564            Kernel::Shutdown => ExecutionStatus::Shutdown,
565        };
566
567        let parent_message_id = message.header.msg_id.clone();
568        let session_view = cx.entity().downgrade();
569        let weak_editor = self.editor.clone();
570        let code_range_for_close = anchor_range.clone();
571
572        let on_close: CloseBlockFn = Arc::new(
573            move |block_id: CustomBlockId, _: &mut Window, cx: &mut App| {
574                if let Some(session) = session_view.upgrade() {
575                    session.update(cx, |session, cx| {
576                        session.blocks.remove(&parent_message_id);
577                        cx.notify();
578                    });
579                }
580
581                if let Some(editor) = weak_editor.upgrade() {
582                    editor.update(cx, |editor, cx| {
583                        let mut block_ids = HashSet::default();
584                        block_ids.insert(block_id);
585                        editor.remove_blocks(block_ids, None, cx);
586                        editor.remove_gutter_highlights::<ReplExecutedRange>(
587                            vec![code_range_for_close.clone()],
588                            cx,
589                        );
590                    });
591                }
592            },
593        );
594
595        let Ok(editor_block) = EditorBlock::new(
596            self.editor.clone(),
597            anchor_range.clone(),
598            status,
599            on_close,
600            cx,
601        ) else {
602            return;
603        };
604
605        self.editor
606            .update(cx, |editor, cx| {
607                editor.insert_gutter_highlight::<ReplExecutedRange>(
608                    anchor_range.clone(),
609                    |cx| cx.theme().status().success,
610                    cx,
611                );
612            })
613            .ok();
614
615        let new_cursor_pos = if let Some(next_cursor) = next_cell {
616            next_cursor
617        } else {
618            editor_block.invalidation_anchor
619        };
620
621        let msg_id = message.header.msg_id.clone();
622        let subscription = cx.subscribe(
623            &editor_block.execution_view,
624            move |session, _execution_view, _event: &ExecutionViewFinishedEmpty, cx| {
625                session.replace_block_with_inlay(&msg_id, "✓", cx);
626            },
627        );
628        self._subscriptions.push(subscription);
629
630        let msg_id = message.header.msg_id.clone();
631        let subscription = cx.subscribe(
632            &editor_block.execution_view,
633            move |session, _execution_view, event: &ExecutionViewFinishedSmall, cx| {
634                session.replace_block_with_inlay(&msg_id, &event.0, cx);
635            },
636        );
637        self._subscriptions.push(subscription);
638
639        self.blocks
640            .insert(message.header.msg_id.clone(), editor_block);
641
642        match &self.kernel {
643            Kernel::RunningKernel(_) => {
644                self.send(message, cx).ok();
645            }
646            Kernel::StartingKernel(task) => {
647                // Queue up the execution as a task to run after the kernel starts
648                let task = task.clone();
649
650                cx.spawn(async move |this, cx| {
651                    task.await;
652                    this.update(cx, |session, cx| {
653                        session.send(message, cx).ok();
654                    })
655                    .ok();
656                })
657                .detach();
658            }
659            _ => {}
660        }
661
662        if move_down {
663            editor.update(cx, move |editor, cx| {
664                editor.change_selections(
665                    SelectionEffects::scroll(Autoscroll::top_relative(8)),
666                    window,
667                    cx,
668                    |selections| {
669                        selections.select_ranges([new_cursor_pos..new_cursor_pos]);
670                    },
671                );
672            });
673        }
674    }
675
676    pub fn interrupt(&mut self, cx: &mut Context<Self>) {
677        match &mut self.kernel {
678            Kernel::RunningKernel(_kernel) => {
679                self.send(InterruptRequest {}.into(), cx).ok();
680            }
681            Kernel::StartingKernel(_task) => {
682                // NOTE: If we switch to a literal queue instead of chaining on to the task, clear all queued executions
683            }
684            _ => {}
685        }
686    }
687
688    pub fn kernel(&mut self, kernel: Kernel, cx: &mut Context<Self>) {
689        if let Kernel::Shutdown = kernel {
690            cx.emit(SessionEvent::Shutdown(self.editor.clone()));
691        }
692
693        let kernel_status = KernelStatus::from(&kernel).to_string();
694        let kernel_language = self.kernel_specification.language();
695
696        telemetry::event!(
697            "Kernel Status Changed",
698            kernel_language,
699            kernel_status,
700            repl_session_id = cx.entity_id().to_string(),
701        );
702
703        self.kernel = kernel;
704    }
705
706    pub fn shutdown(&mut self, window: &mut Window, cx: &mut Context<Self>) {
707        let kernel = std::mem::replace(&mut self.kernel, Kernel::ShuttingDown);
708
709        match kernel {
710            Kernel::RunningKernel(mut kernel) => {
711                let mut request_tx = kernel.request_tx();
712
713                let forced = kernel.force_shutdown(window, cx);
714
715                cx.spawn(async move |this, cx| {
716                    let message: JupyterMessage = ShutdownRequest { restart: false }.into();
717                    request_tx.try_send(message).ok();
718
719                    forced.await.log_err();
720
721                    // Give the kernel a bit of time to clean up
722                    cx.background_executor().timer(Duration::from_secs(3)).await;
723
724                    this.update(cx, |session, cx| {
725                        session.clear_outputs(cx);
726                        session.kernel(Kernel::Shutdown, cx);
727                        cx.notify();
728                    })
729                    .ok();
730                })
731                .detach();
732            }
733            _ => {
734                self.kernel(Kernel::Shutdown, cx);
735            }
736        }
737        cx.notify();
738    }
739
740    pub fn restart(&mut self, window: &mut Window, cx: &mut Context<Self>) {
741        let kernel = std::mem::replace(&mut self.kernel, Kernel::Restarting);
742
743        match kernel {
744            Kernel::Restarting => {
745                // Do nothing if already restarting
746            }
747            Kernel::RunningKernel(mut kernel) => {
748                let mut request_tx = kernel.request_tx();
749
750                let forced = kernel.force_shutdown(window, cx);
751
752                cx.spawn_in(window, async move |this, cx| {
753                    // Send shutdown request with restart flag
754                    log::debug!("restarting kernel");
755                    let message: JupyterMessage = ShutdownRequest { restart: true }.into();
756                    request_tx.try_send(message).ok();
757
758                    // Wait for kernel to shutdown
759                    cx.background_executor().timer(Duration::from_secs(1)).await;
760
761                    // Force kill the kernel if it hasn't shut down
762                    forced.await.log_err();
763
764                    // Start a new kernel
765                    this.update_in(cx, |session, window, cx| {
766                        // TODO: Differentiate between restart and restart+clear-outputs
767                        session.clear_outputs(cx);
768                        session.start_kernel(window, cx);
769                    })
770                    .ok();
771                })
772                .detach();
773            }
774            _ => {
775                self.clear_outputs(cx);
776                self.start_kernel(window, cx);
777            }
778        }
779        cx.notify();
780    }
781}
782
783pub enum SessionEvent {
784    Shutdown(WeakEntity<Editor>),
785}
786
787impl EventEmitter<SessionEvent> for Session {}
788
789impl Render for Session {
790    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
791        let (status_text, interrupt_button) = match &self.kernel {
792            Kernel::RunningKernel(kernel) => (
793                kernel
794                    .kernel_info()
795                    .as_ref()
796                    .map(|info| info.language_info.name.clone()),
797                Some(
798                    Button::new("interrupt", "Interrupt")
799                        .style(ButtonStyle::Subtle)
800                        .on_click(cx.listener(move |session, _, _, cx| {
801                            session.interrupt(cx);
802                        })),
803                ),
804            ),
805            Kernel::StartingKernel(_) => (Some("Starting".into()), None),
806            Kernel::ErroredLaunch(err) => (Some(format!("Error: {err}")), None),
807            Kernel::ShuttingDown => (Some("Shutting Down".into()), None),
808            Kernel::Shutdown => (Some("Shutdown".into()), None),
809            Kernel::Restarting => (Some("Restarting".into()), None),
810        };
811
812        KernelListItem::new(self.kernel_specification.clone())
813            .status_color(match &self.kernel {
814                Kernel::RunningKernel(kernel) => match kernel.execution_state() {
815                    ExecutionState::Idle => Color::Success,
816                    ExecutionState::Busy => Color::Modified,
817                    ExecutionState::Unknown => Color::Modified,
818                    ExecutionState::Starting => Color::Modified,
819                    ExecutionState::Restarting => Color::Modified,
820                    ExecutionState::Terminating => Color::Disabled,
821                    ExecutionState::AutoRestarting => Color::Modified,
822                    ExecutionState::Dead => Color::Disabled,
823                    ExecutionState::Other(_) => Color::Modified,
824                },
825                Kernel::StartingKernel(_) => Color::Modified,
826                Kernel::ErroredLaunch(_) => Color::Error,
827                Kernel::ShuttingDown => Color::Modified,
828                Kernel::Shutdown => Color::Disabled,
829                Kernel::Restarting => Color::Modified,
830            })
831            .child(Label::new(self.kernel_specification.name()))
832            .children(status_text.map(|status_text| Label::new(format!("({status_text})"))))
833            .button(
834                Button::new("shutdown", "Shutdown")
835                    .style(ButtonStyle::Subtle)
836                    .disabled(self.kernel.is_shutting_down())
837                    .on_click(cx.listener(move |session, _, window, cx| {
838                        session.shutdown(window, cx);
839                    })),
840            )
841            .buttons(interrupt_button)
842    }
843}
844
845impl KernelSession for Session {
846    fn route(&mut self, message: &JupyterMessage, window: &mut Window, cx: &mut Context<Self>) {
847        let parent_message_id = match message.parent_header.as_ref() {
848            Some(header) => &header.msg_id,
849            None => return,
850        };
851
852        match &message.content {
853            JupyterMessageContent::Status(status) => {
854                self.kernel.set_execution_state(&status.execution_state);
855
856                telemetry::event!(
857                    "Kernel Status Changed",
858                    kernel_language = self.kernel_specification.language(),
859                    kernel_status = KernelStatus::from(&self.kernel).to_string(),
860                    repl_session_id = cx.entity_id().to_string(),
861                );
862
863                cx.notify();
864            }
865            JupyterMessageContent::KernelInfoReply(reply) => {
866                self.kernel.set_kernel_info(reply);
867                cx.notify();
868            }
869            JupyterMessageContent::UpdateDisplayData(update) => {
870                let display_id = if let Some(display_id) = update.transient.display_id.clone() {
871                    display_id
872                } else {
873                    return;
874                };
875
876                self.blocks.iter_mut().for_each(|(_, block)| {
877                    block.execution_view.update(cx, |execution_view, cx| {
878                        execution_view.update_display_data(&update.data, &display_id, window, cx);
879                    });
880                });
881                return;
882            }
883            _ => {}
884        }
885
886        if let Some(block) = self.blocks.get_mut(parent_message_id) {
887            block.handle_message(message, window, cx);
888        }
889    }
890
891    fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>) {
892        self.kernel_errored(error_message, cx);
893    }
894}