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