tenant.openagents/omega

1use crate::KERNEL_DOCS_URL;
2use crate::kernels::KernelSpecification;
3use crate::repl_store::ReplStore;
4
5use gpui::{AnyView, DismissEvent, FontWeight, SharedString, Task};
6use picker::{Picker, PickerDelegate};
7use project::WorktreeId;
8use std::sync::Arc;
9use ui::{ListItem, ListItemSpacing, PopoverMenu, PopoverMenuHandle, PopoverTrigger, prelude::*};
10
11type OnSelect = Box<dyn Fn(KernelSpecification, &mut Window, &mut App)>;
12
13#[derive(Clone)]
14pub enum KernelPickerEntry {
15    SectionHeader(SharedString),
16    Kernel {
17        spec: KernelSpecification,
18        is_recommended: bool,
19    },
20}
21
22fn build_grouped_entries(store: &ReplStore, worktree_id: WorktreeId) -> Vec<KernelPickerEntry> {
23    let mut entries = Vec::new();
24    let mut recommended_entry: Option<KernelPickerEntry> = None;
25    let mut found_selected = false;
26    let selected_kernel = store.selected_kernel(worktree_id);
27
28    let mut python_envs = Vec::new();
29    let mut jupyter_kernels = Vec::new();
30    let mut wsl_kernels = Vec::new();
31    let mut remote_kernels = Vec::new();
32
33    for spec in store.kernel_specifications_for_worktree(worktree_id) {
34        let is_recommended = store.is_recommended_kernel(worktree_id, spec);
35        let is_selected = selected_kernel.map_or(false, |s| s == spec);
36
37        if is_selected {
38            recommended_entry = Some(KernelPickerEntry::Kernel {
39                spec: spec.clone(),
40                is_recommended: true,
41            });
42            found_selected = true;
43        } else if is_recommended && !found_selected {
44            recommended_entry = Some(KernelPickerEntry::Kernel {
45                spec: spec.clone(),
46                is_recommended: true,
47            });
48        }
49
50        match spec {
51            KernelSpecification::PythonEnv(_) => {
52                python_envs.push(KernelPickerEntry::Kernel {
53                    spec: spec.clone(),
54                    is_recommended,
55                });
56            }
57            KernelSpecification::Jupyter(_) => {
58                jupyter_kernels.push(KernelPickerEntry::Kernel {
59                    spec: spec.clone(),
60                    is_recommended,
61                });
62            }
63            KernelSpecification::JupyterServer(_) | KernelSpecification::SshRemote(_) => {
64                remote_kernels.push(KernelPickerEntry::Kernel {
65                    spec: spec.clone(),
66                    is_recommended,
67                });
68            }
69            KernelSpecification::WslRemote(_) => {
70                wsl_kernels.push(KernelPickerEntry::Kernel {
71                    spec: spec.clone(),
72                    is_recommended,
73                });
74            }
75        }
76    }
77
78    // Sort Python envs: has_ipykernel first, then by name
79    python_envs.sort_by(|a, b| {
80        let (spec_a, spec_b) = match (a, b) {
81            (
82                KernelPickerEntry::Kernel { spec: sa, .. },
83                KernelPickerEntry::Kernel { spec: sb, .. },
84            ) => (sa, sb),
85            _ => return std::cmp::Ordering::Equal,
86        };
87        spec_b
88            .has_ipykernel()
89            .cmp(&spec_a.has_ipykernel())
90            .then_with(|| spec_a.name().cmp(&spec_b.name()))
91    });
92
93    // Recommended section
94    if let Some(rec) = recommended_entry {
95        entries.push(KernelPickerEntry::SectionHeader("Recommended".into()));
96        entries.push(rec);
97    }
98
99    // Python Environments section
100    if !python_envs.is_empty() {
101        entries.push(KernelPickerEntry::SectionHeader(
102            "Python Environments".into(),
103        ));
104        entries.extend(python_envs);
105    }
106
107    // Jupyter Kernels section
108    if !jupyter_kernels.is_empty() {
109        entries.push(KernelPickerEntry::SectionHeader("Jupyter Kernels".into()));
110        entries.extend(jupyter_kernels);
111    }
112
113    // WSL Kernels section
114    if !wsl_kernels.is_empty() {
115        entries.push(KernelPickerEntry::SectionHeader("WSL Kernels".into()));
116        entries.extend(wsl_kernels);
117    }
118
119    // Remote section
120    if !remote_kernels.is_empty() {
121        entries.push(KernelPickerEntry::SectionHeader("Remote Servers".into()));
122        entries.extend(remote_kernels);
123    }
124
125    entries
126}
127
128#[derive(IntoElement)]
129pub struct KernelSelector<T, TT>
130where
131    T: PopoverTrigger + ButtonCommon,
132    TT: Fn(&mut Window, &mut App) -> AnyView + 'static,
133{
134    handle: Option<PopoverMenuHandle<Picker<KernelPickerDelegate>>>,
135    on_select: OnSelect,
136    trigger: T,
137    tooltip: TT,
138    info_text: Option<SharedString>,
139    worktree_id: WorktreeId,
140}
141
142pub struct KernelPickerDelegate {
143    all_entries: Vec<KernelPickerEntry>,
144    filtered_entries: Vec<KernelPickerEntry>,
145    selected_kernelspec: Option<KernelSpecification>,
146    selected_index: usize,
147    on_select: OnSelect,
148}
149
150impl<T, TT> KernelSelector<T, TT>
151where
152    T: PopoverTrigger + ButtonCommon,
153    TT: Fn(&mut Window, &mut App) -> AnyView + 'static,
154{
155    pub fn new(on_select: OnSelect, worktree_id: WorktreeId, trigger: T, tooltip: TT) -> Self {
156        KernelSelector {
157            on_select,
158            handle: None,
159            trigger,
160            tooltip,
161            info_text: None,
162            worktree_id,
163        }
164    }
165
166    pub fn with_handle(mut self, handle: PopoverMenuHandle<Picker<KernelPickerDelegate>>) -> Self {
167        self.handle = Some(handle);
168        self
169    }
170
171    pub fn with_info_text(mut self, text: impl Into<SharedString>) -> Self {
172        self.info_text = Some(text.into());
173        self
174    }
175}
176
177impl KernelPickerDelegate {
178    fn first_selectable_index(entries: &[KernelPickerEntry]) -> usize {
179        entries
180            .iter()
181            .position(|e| matches!(e, KernelPickerEntry::Kernel { .. }))
182            .unwrap_or(0)
183    }
184
185    fn next_selectable_index(&self, from: usize, direction: i32) -> usize {
186        let len = self.filtered_entries.len();
187        if len == 0 {
188            return 0;
189        }
190
191        let mut index = from as i32 + direction;
192        while index >= 0 && (index as usize) < len {
193            if matches!(
194                self.filtered_entries.get(index as usize),
195                Some(KernelPickerEntry::Kernel { .. })
196            ) {
197                return index as usize;
198            }
199            index += direction;
200        }
201
202        from
203    }
204}
205
206impl PickerDelegate for KernelPickerDelegate {
207    type ListItem = ListItem;
208
209    fn name() -> &'static str {
210        "kernel picker"
211    }
212
213    fn match_count(&self) -> usize {
214        self.filtered_entries.len()
215    }
216
217    fn selected_index(&self) -> usize {
218        self.selected_index
219    }
220
221    fn set_selected_index(&mut self, ix: usize, _: &mut Window, cx: &mut Context<Picker<Self>>) {
222        if matches!(
223            self.filtered_entries.get(ix),
224            Some(KernelPickerEntry::SectionHeader(_))
225        ) {
226            let forward = self.next_selectable_index(ix, 1);
227            if forward != ix {
228                self.selected_index = forward;
229            } else {
230                self.selected_index = self.next_selectable_index(ix, -1);
231            }
232        } else {
233            self.selected_index = ix;
234        }
235
236        if let Some(KernelPickerEntry::Kernel { spec, .. }) =
237            self.filtered_entries.get(self.selected_index)
238        {
239            self.selected_kernelspec = Some(spec.clone());
240        }
241        cx.notify();
242    }
243
244    fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
245        "Select a kernel...".into()
246    }
247
248    fn update_matches(
249        &mut self,
250        query: String,
251        _window: &mut Window,
252        _cx: &mut Context<Picker<Self>>,
253    ) -> Task<()> {
254        if query.is_empty() {
255            self.filtered_entries = self.all_entries.clone();
256        } else {
257            let query_lower = query.to_lowercase();
258            let mut filtered = Vec::new();
259            let mut pending_header: Option<KernelPickerEntry> = None;
260
261            for entry in &self.all_entries {
262                match entry {
263                    KernelPickerEntry::SectionHeader(_) => {
264                        pending_header = Some(entry.clone());
265                    }
266                    KernelPickerEntry::Kernel { spec, .. } => {
267                        if spec.name().to_lowercase().contains(&query_lower) {
268                            if let Some(header) = pending_header.take() {
269                                filtered.push(header);
270                            }
271                            filtered.push(entry.clone());
272                        }
273                    }
274                }
275            }
276
277            self.filtered_entries = filtered;
278        }
279
280        self.selected_index = Self::first_selectable_index(&self.filtered_entries);
281        if let Some(KernelPickerEntry::Kernel { spec, .. }) =
282            self.filtered_entries.get(self.selected_index)
283        {
284            self.selected_kernelspec = Some(spec.clone());
285        }
286
287        Task::ready(())
288    }
289
290    fn separators_after_indices(&self) -> Vec<usize> {
291        let mut separators = Vec::new();
292        for (index, entry) in self.filtered_entries.iter().enumerate() {
293            if matches!(entry, KernelPickerEntry::SectionHeader(_)) && index > 0 {
294                separators.push(index - 1);
295            }
296        }
297        separators
298    }
299
300    fn confirm(&mut self, _secondary: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {
301        if let Some(KernelPickerEntry::Kernel { spec, .. }) =
302            self.filtered_entries.get(self.selected_index)
303        {
304            (self.on_select)(spec.clone(), window, cx);
305            cx.emit(DismissEvent);
306        }
307    }
308
309    fn dismissed(&mut self, _window: &mut Window, _cx: &mut Context<Picker<Self>>) {}
310
311    fn render_match(
312        &self,
313        ix: usize,
314        selected: bool,
315        _: &mut Window,
316        cx: &mut Context<Picker<Self>>,
317    ) -> Option<Self::ListItem> {
318        let entry = self.filtered_entries.get(ix)?;
319
320        match entry {
321            KernelPickerEntry::SectionHeader(title) => Some(
322                ListItem::new(ix)
323                    .inset(true)
324                    .spacing(ListItemSpacing::Dense)
325                    .selectable(false)
326                    .child(
327                        Label::new(title.clone())
328                            .size(LabelSize::Small)
329                            .weight(FontWeight::SEMIBOLD)
330                            .color(Color::Muted),
331                    ),
332            ),
333            KernelPickerEntry::Kernel {
334                spec,
335                is_recommended,
336            } => {
337                let is_currently_selected = self.selected_kernelspec.as_ref() == Some(spec);
338                let icon = spec.icon(cx);
339                let has_ipykernel = spec.has_ipykernel();
340
341                let subtitle = match spec {
342                    KernelSpecification::Jupyter(_) => None,
343                    KernelSpecification::WslRemote(_) => Some(spec.path().to_string()),
344                    KernelSpecification::PythonEnv(_)
345                    | KernelSpecification::JupyterServer(_)
346                    | KernelSpecification::SshRemote(_) => {
347                        let env_kind = spec.environment_kind_label();
348                        let path = spec.path();
349                        match env_kind {
350                            Some(kind) => Some(format!("{} \u{2013} {}", kind, path)),
351                            None => Some(path.to_string()),
352                        }
353                    }
354                };
355
356                Some(
357                    ListItem::new(ix)
358                        .inset(true)
359                        .spacing(ListItemSpacing::Sparse)
360                        .toggle_state(selected)
361                        .child(
362                            h_flex()
363                                .w_full()
364                                .gap_3()
365                                .when(!has_ipykernel, |flex| flex.opacity(0.5))
366                                .child(icon.color(Color::Default).size(IconSize::Medium))
367                                .child(
368                                    v_flex()
369                                        .flex_grow_1()
370                                        .overflow_x_hidden()
371                                        .gap_0p5()
372                                        .child(
373                                            h_flex()
374                                                .gap_1()
375                                                .child(
376                                                    div()
377                                                        .overflow_x_hidden()
378                                                        .flex_shrink_1()
379                                                        .text_ellipsis()
380                                                        .child(
381                                                            Label::new(spec.name())
382                                                                .weight(FontWeight::MEDIUM)
383                                                                .size(LabelSize::Default),
384                                                        ),
385                                                )
386                                                .when(*is_recommended, |flex| {
387                                                    flex.child(
388                                                        Label::new("Recommended")
389                                                            .size(LabelSize::XSmall)
390                                                            .color(Color::Accent),
391                                                    )
392                                                })
393                                                .when(!has_ipykernel, |flex| {
394                                                    flex.child(
395                                                        Label::new("ipykernel not installed")
396                                                            .size(LabelSize::XSmall)
397                                                            .color(Color::Warning),
398                                                    )
399                                                }),
400                                        )
401                                        .when_some(subtitle, |flex, subtitle| {
402                                            flex.child(
403                                                div().overflow_x_hidden().text_ellipsis().child(
404                                                    Label::new(subtitle)
405                                                        .size(LabelSize::Small)
406                                                        .color(Color::Muted),
407                                                ),
408                                            )
409                                        }),
410                                ),
411                        )
412                        .when(is_currently_selected, |item| {
413                            item.end_slot(
414                                Icon::new(IconName::Check)
415                                    .color(Color::Accent)
416                                    .size(IconSize::Small),
417                            )
418                        }),
419                )
420            }
421        }
422    }
423
424    fn render_footer(
425        &self,
426        _: &mut Window,
427        cx: &mut Context<Picker<Self>>,
428    ) -> Option<gpui::AnyElement> {
429        Some(
430            h_flex()
431                .w_full()
432                .border_t_1()
433                .border_color(cx.theme().colors().border_variant)
434                .p_1()
435                .gap_4()
436                .child(
437                    Button::new("kernel-docs", "Kernel Docs")
438                        .end_icon(
439                            Icon::new(IconName::ArrowUpRight)
440                                .size(IconSize::Small)
441                                .color(Color::Muted),
442                        )
443                        .on_click(move |_, _, cx| cx.open_url(KERNEL_DOCS_URL)),
444                )
445                .into_any(),
446        )
447    }
448}
449
450impl<T, TT> RenderOnce for KernelSelector<T, TT>
451where
452    T: PopoverTrigger + ButtonCommon,
453    TT: Fn(&mut Window, &mut App) -> AnyView + 'static,
454{
455    fn render(self, window: &mut Window, cx: &mut App) -> impl IntoElement {
456        let store = ReplStore::global(cx);
457        store.update(cx, |store, cx| store.ensure_kernelspecs(cx));
458        let store = store.read(cx);
459
460        let all_entries = build_grouped_entries(store, self.worktree_id);
461        let selected_kernelspec = store.active_kernelspec(self.worktree_id, None, cx);
462        let selected_index = all_entries
463            .iter()
464            .position(|entry| {
465                if let KernelPickerEntry::Kernel { spec, .. } = entry {
466                    selected_kernelspec.as_ref() == Some(spec)
467                } else {
468                    false
469                }
470            })
471            .unwrap_or_else(|| KernelPickerDelegate::first_selectable_index(&all_entries));
472
473        let delegate = KernelPickerDelegate {
474            on_select: self.on_select,
475            all_entries: all_entries.clone(),
476            filtered_entries: all_entries,
477            selected_kernelspec,
478            selected_index,
479        };
480
481        let picker_view = cx.new(|cx| {
482            Picker::list(delegate, window, cx)
483                .list_measure_all()
484                .popover()
485        });
486
487        PopoverMenu::new("kernel-switcher")
488            .menu(move |_window, _cx| Some(picker_view.clone()))
489            .trigger_with_tooltip(self.trigger, self.tooltip)
490            .attach(gpui::Anchor::BottomLeft)
491            .when_some(self.handle, |menu, handle| menu.with_handle(handle))
492    }
493}
494