tenant.openagents/omega

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