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