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