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 clear_output_at_position(&mut self, position: Anchor, cx: &mut Context<Self>) {
518        let Some(editor) = self.editor.upgrade() else {
519            return;
520        };
521
522        let (block_id, code_range, msg_id) = {
523            let snapshot = editor.read(cx).buffer().read(cx).read(cx);
524            let pos_range = position..position;
525
526            let block_to_remove = self
527                .blocks
528                .iter()
529                .find(|(_, block)| block.code_range.includes(&pos_range, &snapshot));
530
531            let Some((msg_id, block)) = block_to_remove else {
532                return;
533            };
534
535            (block.block_id, block.code_range.clone(), msg_id.clone())
536        };
537
538        let inlay_to_remove = self.result_inlays.get(&msg_id).map(|(id, _, _)| *id);
539
540        self.blocks.remove(&msg_id);
541        if inlay_to_remove.is_some() {
542            self.result_inlays.remove(&msg_id);
543        }
544
545        self.editor
546            .update(cx, |editor, cx| {
547                let mut block_ids = HashSet::default();
548                block_ids.insert(block_id);
549                editor.remove_blocks(block_ids, None, cx);
550
551                if let Some(inlay_id) = inlay_to_remove {
552                    editor.splice_inlays(&[inlay_id], vec![], cx);
553                }
554
555                editor.remove_gutter_highlights::<ReplExecutedRange>(vec![code_range], cx);
556            })
557            .ok();
558
559        cx.notify();
560    }
561
562    pub fn execute(
563        &mut self,
564        code: String,
565        anchor_range: Range<Anchor>,
566        next_cell: Option<Anchor>,
567        move_down: bool,
568        window: &mut Window,
569        cx: &mut Context<Self>,
570    ) {
571        let Some(editor) = self.editor.upgrade() else {
572            return;
573        };
574
575        if code.is_empty() {
576            return;
577        }
578
579        let execute_request = ExecuteRequest {
580            code,
581            allow_stdin: true,
582            ..ExecuteRequest::default()
583        };
584
585        let message: JupyterMessage = execute_request.into();
586
587        let mut blocks_to_remove: HashSet<CustomBlockId> = HashSet::default();
588        let mut inlays_to_remove: Vec<InlayId> = Vec::new();
589        let mut gutter_ranges_to_remove: Vec<Range<Anchor>> = Vec::new();
590
591        let buffer = editor.read(cx).buffer().read(cx).snapshot(cx);
592
593        self.blocks.retain(|_key, block| {
594            if anchor_range.overlaps(&block.code_range, &buffer) {
595                blocks_to_remove.insert(block.block_id);
596                false
597            } else {
598                true
599            }
600        });
601
602        self.result_inlays
603            .retain(|_key, (inlay_id, inlay_range, _)| {
604                if anchor_range.overlaps(inlay_range, &buffer) {
605                    inlays_to_remove.push(*inlay_id);
606                    gutter_ranges_to_remove.push(inlay_range.clone());
607                    false
608                } else {
609                    true
610                }
611            });
612
613        self.editor
614            .update(cx, |editor, cx| {
615                editor.remove_blocks(blocks_to_remove, None, cx);
616                if !inlays_to_remove.is_empty() {
617                    editor.splice_inlays(&inlays_to_remove, vec![], cx);
618                }
619                if !gutter_ranges_to_remove.is_empty() {
620                    editor
621                        .remove_gutter_highlights::<ReplExecutedRange>(gutter_ranges_to_remove, cx);
622                }
623            })
624            .ok();
625
626        let status = match &self.kernel {
627            Kernel::Restarting => ExecutionStatus::Restarting,
628            Kernel::RunningKernel(_) => ExecutionStatus::Queued,
629            Kernel::StartingKernel(_) => ExecutionStatus::ConnectingToKernel,
630            Kernel::ErroredLaunch(error) => ExecutionStatus::KernelErrored(error.clone()),
631            Kernel::ShuttingDown => ExecutionStatus::ShuttingDown,
632            Kernel::Shutdown => ExecutionStatus::Shutdown,
633        };
634
635        let parent_message_id = message.header.msg_id.clone();
636        let session_view = cx.entity().downgrade();
637        let weak_editor = self.editor.clone();
638        let code_range_for_close = anchor_range.clone();
639
640        let on_close: CloseBlockFn = Arc::new(
641            move |block_id: CustomBlockId, _: &mut Window, cx: &mut App| {
642                if let Some(session) = session_view.upgrade() {
643                    session.update(cx, |session, cx| {
644                        session.blocks.remove(&parent_message_id);
645                        cx.notify();
646                    });
647                }
648
649                if let Some(editor) = weak_editor.upgrade() {
650                    editor.update(cx, |editor, cx| {
651                        let mut block_ids = HashSet::default();
652                        block_ids.insert(block_id);
653                        editor.remove_blocks(block_ids, None, cx);
654                        editor.remove_gutter_highlights::<ReplExecutedRange>(
655                            vec![code_range_for_close.clone()],
656                            cx,
657                        );
658                    });
659                }
660            },
661        );
662
663        let Ok(editor_block) = EditorBlock::new(
664            self.editor.clone(),
665            anchor_range.clone(),
666            status,
667            on_close,
668            cx,
669        ) else {
670            return;
671        };
672
673        self.editor
674            .update(cx, |editor, cx| {
675                editor.insert_gutter_highlight::<ReplExecutedRange>(
676                    anchor_range.clone(),
677                    |cx| cx.theme().status().success,
678                    cx,
679                );
680            })
681            .ok();
682
683        let new_cursor_pos = if let Some(next_cursor) = next_cell {
684            next_cursor
685        } else {
686            editor_block.invalidation_anchor
687        };
688
689        let msg_id = message.header.msg_id.clone();
690        let subscription = cx.subscribe(
691            &editor_block.execution_view,
692            move |session, _execution_view, _event: &ExecutionViewFinishedEmpty, cx| {
693                session.replace_block_with_inlay(&msg_id, "✓", cx);
694            },
695        );
696        self._subscriptions.push(subscription);
697
698        let msg_id = message.header.msg_id.clone();
699        let subscription = cx.subscribe(
700            &editor_block.execution_view,
701            move |session, _execution_view, event: &ExecutionViewFinishedSmall, cx| {
702                session.replace_block_with_inlay(&msg_id, &event.0, cx);
703            },
704        );
705        self._subscriptions.push(subscription);
706
707        let subscription = cx.subscribe(
708            &editor_block.execution_view,
709            |session, _execution_view, event: &InputReplyEvent, cx| {
710                session.send_stdin_reply(event.value.clone(), &event.parent_message, cx);
711            },
712        );
713        self._subscriptions.push(subscription);
714
715        self.blocks
716            .insert(message.header.msg_id.clone(), editor_block);
717
718        match &self.kernel {
719            Kernel::RunningKernel(_) => {
720                self.send(message, cx).ok();
721            }
722            Kernel::StartingKernel(task) => {
723                // Queue up the execution as a task to run after the kernel starts
724                let task = task.clone();
725
726                cx.spawn(async move |this, cx| {
727                    task.await;
728                    this.update(cx, |session, cx| {
729                        session.send(message, cx).ok();
730                    })
731                    .ok();
732                })
733                .detach();
734            }
735            _ => {}
736        }
737
738        if move_down {
739            editor.update(cx, move |editor, cx| {
740                editor.change_selections(
741                    SelectionEffects::scroll(Autoscroll::top_relative(8)),
742                    window,
743                    cx,
744                    |selections| {
745                        selections.select_ranges([new_cursor_pos..new_cursor_pos]);
746                    },
747                );
748            });
749        }
750    }
751
752    pub fn interrupt(&mut self, cx: &mut Context<Self>) {
753        match &mut self.kernel {
754            Kernel::RunningKernel(_kernel) => {
755                self.send(InterruptRequest {}.into(), cx).ok();
756            }
757            Kernel::StartingKernel(_task) => {
758                // NOTE: If we switch to a literal queue instead of chaining on to the task, clear all queued executions
759            }
760            _ => {}
761        }
762    }
763
764    pub fn kernel(&mut self, kernel: Kernel, cx: &mut Context<Self>) {
765        if let Kernel::Shutdown = kernel {
766            cx.emit(SessionEvent::Shutdown(self.editor.clone()));
767        }
768
769        let kernel_status = KernelStatus::from(&kernel).to_string();
770        let kernel_language = self.kernel_specification.language();
771
772        telemetry::event!(
773            "Kernel Status Changed",
774            kernel_language,
775            kernel_status,
776            repl_session_id = cx.entity_id().to_string(),
777        );
778
779        self.kernel = kernel;
780    }
781
782    pub fn shutdown(&mut self, window: &mut Window, cx: &mut Context<Self>) {
783        let kernel = std::mem::replace(&mut self.kernel, Kernel::ShuttingDown);
784
785        match kernel {
786            Kernel::RunningKernel(mut kernel) => {
787                let mut request_tx = kernel.request_tx();
788
789                let forced = kernel.force_shutdown(window, cx);
790
791                cx.spawn(async move |this, cx| {
792                    let message: JupyterMessage = ShutdownRequest { restart: false }.into();
793                    request_tx.try_send(message).ok();
794
795                    forced.await.log_err();
796
797                    // Give the kernel a bit of time to clean up
798                    cx.background_executor().timer(Duration::from_secs(3)).await;
799
800                    this.update(cx, |session, cx| {
801                        session.clear_outputs(cx);
802                        session.kernel(Kernel::Shutdown, cx);
803                        cx.notify();
804                    })
805                    .ok();
806                })
807                .detach();
808            }
809            _ => {
810                self.kernel(Kernel::Shutdown, cx);
811            }
812        }
813        cx.notify();
814    }
815
816    pub fn restart(&mut self, window: &mut Window, cx: &mut Context<Self>) {
817        let kernel = std::mem::replace(&mut self.kernel, Kernel::Restarting);
818
819        match kernel {
820            Kernel::Restarting => {
821                // Do nothing if already restarting
822            }
823            Kernel::RunningKernel(mut kernel) => {
824                let mut request_tx = kernel.request_tx();
825
826                let forced = kernel.force_shutdown(window, cx);
827
828                cx.spawn_in(window, async move |this, cx| {
829                    // Send shutdown request with restart flag
830                    log::debug!("restarting kernel");
831                    let message: JupyterMessage = ShutdownRequest { restart: true }.into();
832                    request_tx.try_send(message).ok();
833
834                    // Wait for kernel to shutdown
835                    cx.background_executor().timer(Duration::from_secs(1)).await;
836
837                    // Force kill the kernel if it hasn't shut down
838                    forced.await.log_err();
839
840                    // Start a new kernel
841                    this.update_in(cx, |session, window, cx| {
842                        // TODO: Differentiate between restart and restart+clear-outputs
843                        session.clear_outputs(cx);
844                        session.start_kernel(window, cx);
845                    })
846                    .ok();
847                })
848                .detach();
849            }
850            _ => {
851                self.clear_outputs(cx);
852                self.start_kernel(window, cx);
853            }
854        }
855        cx.notify();
856    }
857}
858
859pub enum SessionEvent {
860    Shutdown(WeakEntity<Editor>),
861}
862
863impl EventEmitter<SessionEvent> for Session {}
864
865impl Render for Session {
866    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
867        let (status_text, interrupt_button) = match &self.kernel {
868            Kernel::RunningKernel(kernel) => (
869                kernel
870                    .kernel_info()
871                    .as_ref()
872                    .map(|info| info.language_info.name.clone()),
873                Some(
874                    Button::new("interrupt", "Interrupt")
875                        .style(ButtonStyle::Subtle)
876                        .on_click(cx.listener(move |session, _, _, cx| {
877                            session.interrupt(cx);
878                        })),
879                ),
880            ),
881            Kernel::StartingKernel(_) => (Some("Starting".into()), None),
882            Kernel::ErroredLaunch(err) => (Some(format!("Error: {err}")), None),
883            Kernel::ShuttingDown => (Some("Shutting Down".into()), None),
884            Kernel::Shutdown => (Some("Shutdown".into()), None),
885            Kernel::Restarting => (Some("Restarting".into()), None),
886        };
887
888        KernelListItem::new(self.kernel_specification.clone())
889            .status_color(match &self.kernel {
890                Kernel::RunningKernel(kernel) => match kernel.execution_state() {
891                    ExecutionState::Idle => Color::Success,
892                    ExecutionState::Busy => Color::Modified,
893                    ExecutionState::Unknown => Color::Modified,
894                    ExecutionState::Starting => Color::Modified,
895                    ExecutionState::Restarting => Color::Modified,
896                    ExecutionState::Terminating => Color::Disabled,
897                    ExecutionState::AutoRestarting => Color::Modified,
898                    ExecutionState::Dead => Color::Disabled,
899                    ExecutionState::Other(_) => Color::Modified,
900                },
901                Kernel::StartingKernel(_) => Color::Modified,
902                Kernel::ErroredLaunch(_) => Color::Error,
903                Kernel::ShuttingDown => Color::Modified,
904                Kernel::Shutdown => Color::Disabled,
905                Kernel::Restarting => Color::Modified,
906            })
907            .child(Label::new(self.kernel_specification.name()))
908            .children(status_text.map(|status_text| Label::new(format!("({status_text})"))))
909            .button(
910                Button::new("shutdown", "Shutdown")
911                    .style(ButtonStyle::Subtle)
912                    .disabled(self.kernel.is_shutting_down())
913                    .on_click(cx.listener(move |session, _, window, cx| {
914                        session.shutdown(window, cx);
915                    })),
916            )
917            .buttons(interrupt_button)
918    }
919}
920
921impl KernelSession for Session {
922    fn route(&mut self, message: &JupyterMessage, window: &mut Window, cx: &mut Context<Self>) {
923        let parent_message_id = match message.parent_header.as_ref() {
924            Some(header) => &header.msg_id,
925            None => return,
926        };
927
928        match &message.content {
929            JupyterMessageContent::Status(status) => {
930                self.kernel.set_execution_state(&status.execution_state);
931
932                telemetry::event!(
933                    "Kernel Status Changed",
934                    kernel_language = self.kernel_specification.language(),
935                    kernel_status = KernelStatus::from(&self.kernel).to_string(),
936                    repl_session_id = cx.entity_id().to_string(),
937                );
938
939                cx.notify();
940            }
941            JupyterMessageContent::KernelInfoReply(reply) => {
942                self.kernel.set_kernel_info(reply);
943                cx.notify();
944            }
945            JupyterMessageContent::UpdateDisplayData(update) => {
946                let display_id = if let Some(display_id) = update.transient.display_id.clone() {
947                    display_id
948                } else {
949                    return;
950                };
951
952                self.blocks.iter_mut().for_each(|(_, block)| {
953                    block.execution_view.update(cx, |execution_view, cx| {
954                        execution_view.update_display_data(&update.data, &display_id, window, cx);
955                    });
956                });
957                return;
958            }
959            _ => {}
960        }
961
962        if let Some(block) = self.blocks.get_mut(parent_message_id) {
963            block.handle_message(message, window, cx);
964        }
965    }
966
967    fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>) {
968        self.kernel_errored(error_message, cx);
969    }
970}