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)
281            | KernelSpecification::PythonEnv(kernel_specification) => NativeRunningKernel::new(
282                kernel_specification,
283                entity_id,
284                working_directory,
285                self.fs.clone(),
286                session_view,
287                window,
288                cx,
289            ),
290            KernelSpecification::Remote(remote_kernel_specification) => RemoteRunningKernel::new(
291                remote_kernel_specification,
292                working_directory,
293                session_view,
294                window,
295                cx,
296            ),
297        };
298
299        let pending_kernel = cx
300            .spawn(async move |this, cx| {
301                let kernel = kernel.await;
302
303                match kernel {
304                    Ok(kernel) => {
305                        this.update(cx, |session, cx| {
306                            session.kernel(Kernel::RunningKernel(kernel), cx);
307                        })
308                        .ok();
309                    }
310                    Err(err) => {
311                        this.update(cx, |session, cx| {
312                            session.kernel_errored(err.to_string(), cx);
313                        })
314                        .ok();
315                    }
316                }
317            })
318            .shared();
319
320        self.kernel(Kernel::StartingKernel(pending_kernel), cx);
321        cx.notify();
322    }
323
324    pub fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>) {
325        self.kernel(Kernel::ErroredLaunch(error_message.clone()), cx);
326
327        self.blocks.values().for_each(|block| {
328            block.execution_view.update(cx, |execution_view, cx| {
329                match execution_view.status {
330                    ExecutionStatus::Finished => {
331                        // Do nothing when the output was good
332                    }
333                    _ => {
334                        // All other cases, set the status to errored
335                        execution_view.status =
336                            ExecutionStatus::KernelErrored(error_message.clone())
337                    }
338                }
339                cx.notify();
340            });
341        });
342    }
343
344    fn on_buffer_event(
345        &mut self,
346        buffer: Entity<MultiBuffer>,
347        event: &multi_buffer::Event,
348        cx: &mut Context<Self>,
349    ) {
350        if let multi_buffer::Event::Edited { .. } = event {
351            let snapshot = buffer.read(cx).snapshot(cx);
352
353            let mut blocks_to_remove: HashSet<CustomBlockId> = HashSet::default();
354            let mut gutter_ranges_to_remove: Vec<Range<Anchor>> = Vec::new();
355            let mut keys_to_remove: Vec<String> = Vec::new();
356
357            self.blocks.retain(|id, block| {
358                if block.invalidation_anchor.is_valid(&snapshot) {
359                    true
360                } else {
361                    blocks_to_remove.insert(block.block_id);
362                    gutter_ranges_to_remove.push(block.code_range.clone());
363                    keys_to_remove.push(id.clone());
364                    false
365                }
366            });
367
368            let mut inlays_to_remove: Vec<InlayId> = Vec::new();
369
370            self.result_inlays
371                .retain(|id, (inlay_id, code_range, original_len)| {
372                    let start_offset = code_range.start.to_offset(&snapshot);
373                    let end_offset = code_range.end.to_offset(&snapshot);
374                    let current_len = end_offset.saturating_sub(start_offset);
375
376                    if current_len != *original_len {
377                        inlays_to_remove.push(*inlay_id);
378                        gutter_ranges_to_remove.push(code_range.clone());
379                        keys_to_remove.push(id.clone());
380                        false
381                    } else {
382                        true
383                    }
384                });
385
386            if !blocks_to_remove.is_empty()
387                || !inlays_to_remove.is_empty()
388                || !gutter_ranges_to_remove.is_empty()
389            {
390                self.editor
391                    .update(cx, |editor, cx| {
392                        if !blocks_to_remove.is_empty() {
393                            editor.remove_blocks(blocks_to_remove, None, cx);
394                        }
395                        if !inlays_to_remove.is_empty() {
396                            editor.splice_inlays(&inlays_to_remove, vec![], cx);
397                        }
398                        if !gutter_ranges_to_remove.is_empty() {
399                            editor.remove_gutter_highlights::<ReplExecutedRange>(
400                                gutter_ranges_to_remove,
401                                cx,
402                            );
403                        }
404                    })
405                    .ok();
406                cx.notify();
407            }
408        }
409    }
410
411    fn send(&mut self, message: JupyterMessage, _cx: &mut Context<Self>) -> anyhow::Result<()> {
412        if let Kernel::RunningKernel(kernel) = &mut self.kernel {
413            kernel.request_tx().try_send(message).ok();
414        }
415
416        anyhow::Ok(())
417    }
418
419    fn replace_block_with_inlay(&mut self, message_id: &str, text: &str, cx: &mut Context<Self>) {
420        let Some(block) = self.blocks.remove(message_id) else {
421            return;
422        };
423
424        let Some(editor) = self.editor.upgrade() else {
425            return;
426        };
427
428        let code_range = block.code_range.clone();
429
430        editor.update(cx, |editor, cx| {
431            let mut block_ids = HashSet::default();
432            block_ids.insert(block.block_id);
433            editor.remove_blocks(block_ids, None, cx);
434
435            let buffer = editor.buffer().read(cx).snapshot(cx);
436            let start_offset = code_range.start.to_offset(&buffer);
437            let end_offset = code_range.end.to_offset(&buffer);
438            let original_len = end_offset.saturating_sub(start_offset);
439
440            let end_point = code_range.end.to_point(&buffer);
441            let inlay_position = buffer.anchor_after(end_point);
442
443            let inlay_id = self.next_inlay_id;
444            self.next_inlay_id += 1;
445
446            let inlay = Inlay::repl_result(inlay_id, inlay_position, format!("    {}", text));
447
448            editor.splice_inlays(&[], vec![inlay], cx);
449            self.result_inlays.insert(
450                message_id.to_string(),
451                (
452                    InlayId::ReplResult(inlay_id),
453                    code_range.clone(),
454                    original_len,
455                ),
456            );
457
458            editor.insert_gutter_highlight::<ReplExecutedRange>(
459                code_range,
460                |cx| cx.theme().status().success,
461                cx,
462            );
463        });
464
465        cx.notify();
466    }
467
468    pub fn clear_outputs(&mut self, cx: &mut Context<Self>) {
469        let blocks_to_remove: HashSet<CustomBlockId> =
470            self.blocks.values().map(|block| block.block_id).collect();
471
472        let inlays_to_remove: Vec<InlayId> =
473            self.result_inlays.values().map(|(id, _, _)| *id).collect();
474
475        self.editor
476            .update(cx, |editor, cx| {
477                editor.remove_blocks(blocks_to_remove, None, cx);
478                editor.splice_inlays(&inlays_to_remove, vec![], cx);
479                editor.clear_gutter_highlights::<ReplExecutedRange>(cx);
480            })
481            .ok();
482
483        self.blocks.clear();
484        self.result_inlays.clear();
485    }
486
487    pub fn execute(
488        &mut self,
489        code: String,
490        anchor_range: Range<Anchor>,
491        next_cell: Option<Anchor>,
492        move_down: bool,
493        window: &mut Window,
494        cx: &mut Context<Self>,
495    ) {
496        let Some(editor) = self.editor.upgrade() else {
497            return;
498        };
499
500        if code.is_empty() {
501            return;
502        }
503
504        let execute_request = ExecuteRequest {
505            code,
506            ..ExecuteRequest::default()
507        };
508
509        let message: JupyterMessage = execute_request.into();
510
511        let mut blocks_to_remove: HashSet<CustomBlockId> = HashSet::default();
512        let mut inlays_to_remove: Vec<InlayId> = Vec::new();
513        let mut gutter_ranges_to_remove: Vec<Range<Anchor>> = Vec::new();
514
515        let buffer = editor.read(cx).buffer().read(cx).snapshot(cx);
516
517        self.blocks.retain(|_key, block| {
518            if anchor_range.overlaps(&block.code_range, &buffer) {
519                blocks_to_remove.insert(block.block_id);
520                false
521            } else {
522                true
523            }
524        });
525
526        self.result_inlays
527            .retain(|_key, (inlay_id, inlay_range, _)| {
528                if anchor_range.overlaps(inlay_range, &buffer) {
529                    inlays_to_remove.push(*inlay_id);
530                    gutter_ranges_to_remove.push(inlay_range.clone());
531                    false
532                } else {
533                    true
534                }
535            });
536
537        self.editor
538            .update(cx, |editor, cx| {
539                editor.remove_blocks(blocks_to_remove, None, cx);
540                if !inlays_to_remove.is_empty() {
541                    editor.splice_inlays(&inlays_to_remove, vec![], cx);
542                }
543                if !gutter_ranges_to_remove.is_empty() {
544                    editor
545                        .remove_gutter_highlights::<ReplExecutedRange>(gutter_ranges_to_remove, cx);
546                }
547            })
548            .ok();
549
550        let status = match &self.kernel {
551            Kernel::Restarting => ExecutionStatus::Restarting,
552            Kernel::RunningKernel(_) => ExecutionStatus::Queued,
553            Kernel::StartingKernel(_) => ExecutionStatus::ConnectingToKernel,
554            Kernel::ErroredLaunch(error) => ExecutionStatus::KernelErrored(error.clone()),
555            Kernel::ShuttingDown => ExecutionStatus::ShuttingDown,
556            Kernel::Shutdown => ExecutionStatus::Shutdown,
557        };
558
559        let parent_message_id = message.header.msg_id.clone();
560        let session_view = cx.entity().downgrade();
561        let weak_editor = self.editor.clone();
562        let code_range_for_close = anchor_range.clone();
563
564        let on_close: CloseBlockFn = Arc::new(
565            move |block_id: CustomBlockId, _: &mut Window, cx: &mut App| {
566                if let Some(session) = session_view.upgrade() {
567                    session.update(cx, |session, cx| {
568                        session.blocks.remove(&parent_message_id);
569                        cx.notify();
570                    });
571                }
572
573                if let Some(editor) = weak_editor.upgrade() {
574                    editor.update(cx, |editor, cx| {
575                        let mut block_ids = HashSet::default();
576                        block_ids.insert(block_id);
577                        editor.remove_blocks(block_ids, None, cx);
578                        editor.remove_gutter_highlights::<ReplExecutedRange>(
579                            vec![code_range_for_close.clone()],
580                            cx,
581                        );
582                    });
583                }
584            },
585        );
586
587        let Ok(editor_block) = EditorBlock::new(
588            self.editor.clone(),
589            anchor_range.clone(),
590            status,
591            on_close,
592            cx,
593        ) else {
594            return;
595        };
596
597        self.editor
598            .update(cx, |editor, cx| {
599                editor.insert_gutter_highlight::<ReplExecutedRange>(
600                    anchor_range.clone(),
601                    |cx| cx.theme().status().success,
602                    cx,
603                );
604            })
605            .ok();
606
607        let new_cursor_pos = if let Some(next_cursor) = next_cell {
608            next_cursor
609        } else {
610            editor_block.invalidation_anchor
611        };
612
613        let msg_id = message.header.msg_id.clone();
614        let subscription = cx.subscribe(
615            &editor_block.execution_view,
616            move |session, _execution_view, _event: &ExecutionViewFinishedEmpty, cx| {
617                session.replace_block_with_inlay(&msg_id, "✓", cx);
618            },
619        );
620        self._subscriptions.push(subscription);
621
622        let msg_id = message.header.msg_id.clone();
623        let subscription = cx.subscribe(
624            &editor_block.execution_view,
625            move |session, _execution_view, event: &ExecutionViewFinishedSmall, cx| {
626                session.replace_block_with_inlay(&msg_id, &event.0, cx);
627            },
628        );
629        self._subscriptions.push(subscription);
630
631        self.blocks
632            .insert(message.header.msg_id.clone(), editor_block);
633
634        match &self.kernel {
635            Kernel::RunningKernel(_) => {
636                self.send(message, cx).ok();
637            }
638            Kernel::StartingKernel(task) => {
639                // Queue up the execution as a task to run after the kernel starts
640                let task = task.clone();
641
642                cx.spawn(async move |this, cx| {
643                    task.await;
644                    this.update(cx, |session, cx| {
645                        session.send(message, cx).ok();
646                    })
647                    .ok();
648                })
649                .detach();
650            }
651            _ => {}
652        }
653
654        if move_down {
655            editor.update(cx, move |editor, cx| {
656                editor.change_selections(
657                    SelectionEffects::scroll(Autoscroll::top_relative(8)),
658                    window,
659                    cx,
660                    |selections| {
661                        selections.select_ranges([new_cursor_pos..new_cursor_pos]);
662                    },
663                );
664            });
665        }
666    }
667
668    pub fn interrupt(&mut self, cx: &mut Context<Self>) {
669        match &mut self.kernel {
670            Kernel::RunningKernel(_kernel) => {
671                self.send(InterruptRequest {}.into(), cx).ok();
672            }
673            Kernel::StartingKernel(_task) => {
674                // NOTE: If we switch to a literal queue instead of chaining on to the task, clear all queued executions
675            }
676            _ => {}
677        }
678    }
679
680    pub fn kernel(&mut self, kernel: Kernel, cx: &mut Context<Self>) {
681        if let Kernel::Shutdown = kernel {
682            cx.emit(SessionEvent::Shutdown(self.editor.clone()));
683        }
684
685        let kernel_status = KernelStatus::from(&kernel).to_string();
686        let kernel_language = self.kernel_specification.language();
687
688        telemetry::event!(
689            "Kernel Status Changed",
690            kernel_language,
691            kernel_status,
692            repl_session_id = cx.entity_id().to_string(),
693        );
694
695        self.kernel = kernel;
696    }
697
698    pub fn shutdown(&mut self, window: &mut Window, cx: &mut Context<Self>) {
699        let kernel = std::mem::replace(&mut self.kernel, Kernel::ShuttingDown);
700
701        match kernel {
702            Kernel::RunningKernel(mut kernel) => {
703                let mut request_tx = kernel.request_tx();
704
705                let forced = kernel.force_shutdown(window, cx);
706
707                cx.spawn(async move |this, cx| {
708                    let message: JupyterMessage = ShutdownRequest { restart: false }.into();
709                    request_tx.try_send(message).ok();
710
711                    forced.await.log_err();
712
713                    // Give the kernel a bit of time to clean up
714                    cx.background_executor().timer(Duration::from_secs(3)).await;
715
716                    this.update(cx, |session, cx| {
717                        session.clear_outputs(cx);
718                        session.kernel(Kernel::Shutdown, cx);
719                        cx.notify();
720                    })
721                    .ok();
722                })
723                .detach();
724            }
725            _ => {
726                self.kernel(Kernel::Shutdown, cx);
727            }
728        }
729        cx.notify();
730    }
731
732    pub fn restart(&mut self, window: &mut Window, cx: &mut Context<Self>) {
733        let kernel = std::mem::replace(&mut self.kernel, Kernel::Restarting);
734
735        match kernel {
736            Kernel::Restarting => {
737                // Do nothing if already restarting
738            }
739            Kernel::RunningKernel(mut kernel) => {
740                let mut request_tx = kernel.request_tx();
741
742                let forced = kernel.force_shutdown(window, cx);
743
744                cx.spawn_in(window, async move |this, cx| {
745                    // Send shutdown request with restart flag
746                    log::debug!("restarting kernel");
747                    let message: JupyterMessage = ShutdownRequest { restart: true }.into();
748                    request_tx.try_send(message).ok();
749
750                    // Wait for kernel to shutdown
751                    cx.background_executor().timer(Duration::from_secs(1)).await;
752
753                    // Force kill the kernel if it hasn't shut down
754                    forced.await.log_err();
755
756                    // Start a new kernel
757                    this.update_in(cx, |session, window, cx| {
758                        // TODO: Differentiate between restart and restart+clear-outputs
759                        session.clear_outputs(cx);
760                        session.start_kernel(window, cx);
761                    })
762                    .ok();
763                })
764                .detach();
765            }
766            _ => {
767                self.clear_outputs(cx);
768                self.start_kernel(window, cx);
769            }
770        }
771        cx.notify();
772    }
773}
774
775pub enum SessionEvent {
776    Shutdown(WeakEntity<Editor>),
777}
778
779impl EventEmitter<SessionEvent> for Session {}
780
781impl Render for Session {
782    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
783        let (status_text, interrupt_button) = match &self.kernel {
784            Kernel::RunningKernel(kernel) => (
785                kernel
786                    .kernel_info()
787                    .as_ref()
788                    .map(|info| info.language_info.name.clone()),
789                Some(
790                    Button::new("interrupt", "Interrupt")
791                        .style(ButtonStyle::Subtle)
792                        .on_click(cx.listener(move |session, _, _, cx| {
793                            session.interrupt(cx);
794                        })),
795                ),
796            ),
797            Kernel::StartingKernel(_) => (Some("Starting".into()), None),
798            Kernel::ErroredLaunch(err) => (Some(format!("Error: {err}")), None),
799            Kernel::ShuttingDown => (Some("Shutting Down".into()), None),
800            Kernel::Shutdown => (Some("Shutdown".into()), None),
801            Kernel::Restarting => (Some("Restarting".into()), None),
802        };
803
804        KernelListItem::new(self.kernel_specification.clone())
805            .status_color(match &self.kernel {
806                Kernel::RunningKernel(kernel) => match kernel.execution_state() {
807                    ExecutionState::Idle => Color::Success,
808                    ExecutionState::Busy => Color::Modified,
809                    ExecutionState::Unknown => Color::Modified,
810                    ExecutionState::Starting => Color::Modified,
811                    ExecutionState::Restarting => Color::Modified,
812                    ExecutionState::Terminating => Color::Disabled,
813                    ExecutionState::AutoRestarting => Color::Modified,
814                    ExecutionState::Dead => Color::Disabled,
815                    ExecutionState::Other(_) => Color::Modified,
816                },
817                Kernel::StartingKernel(_) => Color::Modified,
818                Kernel::ErroredLaunch(_) => Color::Error,
819                Kernel::ShuttingDown => Color::Modified,
820                Kernel::Shutdown => Color::Disabled,
821                Kernel::Restarting => Color::Modified,
822            })
823            .child(Label::new(self.kernel_specification.name()))
824            .children(status_text.map(|status_text| Label::new(format!("({status_text})"))))
825            .button(
826                Button::new("shutdown", "Shutdown")
827                    .style(ButtonStyle::Subtle)
828                    .disabled(self.kernel.is_shutting_down())
829                    .on_click(cx.listener(move |session, _, window, cx| {
830                        session.shutdown(window, cx);
831                    })),
832            )
833            .buttons(interrupt_button)
834    }
835}
836
837impl KernelSession for Session {
838    fn route(&mut self, message: &JupyterMessage, window: &mut Window, cx: &mut Context<Self>) {
839        let parent_message_id = match message.parent_header.as_ref() {
840            Some(header) => &header.msg_id,
841            None => return,
842        };
843
844        match &message.content {
845            JupyterMessageContent::Status(status) => {
846                self.kernel.set_execution_state(&status.execution_state);
847
848                telemetry::event!(
849                    "Kernel Status Changed",
850                    kernel_language = self.kernel_specification.language(),
851                    kernel_status = KernelStatus::from(&self.kernel).to_string(),
852                    repl_session_id = cx.entity_id().to_string(),
853                );
854
855                cx.notify();
856            }
857            JupyterMessageContent::KernelInfoReply(reply) => {
858                self.kernel.set_kernel_info(reply);
859                cx.notify();
860            }
861            JupyterMessageContent::UpdateDisplayData(update) => {
862                let display_id = if let Some(display_id) = update.transient.display_id.clone() {
863                    display_id
864                } else {
865                    return;
866                };
867
868                self.blocks.iter_mut().for_each(|(_, block)| {
869                    block.execution_view.update(cx, |execution_view, cx| {
870                        execution_view.update_display_data(&update.data, &display_id, window, cx);
871                    });
872                });
873                return;
874            }
875            _ => {}
876        }
877
878        if let Some(block) = self.blocks.get_mut(parent_message_id) {
879            block.handle_message(message, window, cx);
880        }
881    }
882
883    fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>) {
884        self.kernel_errored(error_message, cx);
885    }
886}