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)),
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}