tenant.openagents/omega

1use editor::Editor;
2use gpui::{
3    AnyElement, App, Entity, EventEmitter, FocusHandle, Focusable, Subscription, TaskExt, actions,
4    prelude::*,
5};
6use language::{Buffer, BufferEvent};
7use project::{Project, ProjectItem as _};
8use ui::{ButtonLike, ElevationIndex, KeyBinding, prelude::*};
9use util::ResultExt as _;
10use workspace::item::ItemEvent;
11use workspace::{Workspace, item::Item};
12
13use crate::jupyter_settings::JupyterSettings;
14use crate::repl_store::ReplStore;
15
16fn refresh_python_kernelspecs_for_buffer(
17    buffer: &Entity<Buffer>,
18    project: &Entity<Project>,
19    cx: &mut App,
20) {
21    let buffer = buffer.read(cx);
22
23    if buffer
24        .language()
25        .is_none_or(|language| language.name() != "Python")
26    {
27        return;
28    };
29
30    let Some(project_path) = buffer.project_path(cx) else {
31        return;
32    };
33    let store = ReplStore::global(cx);
34    store.update(cx, |store, cx| {
35        store
36            .refresh_python_kernelspecs(project_path.worktree_id, project, cx)
37            .detach_and_log_err(cx);
38    });
39}
40
41actions!(
42    repl,
43    [
44        /// Runs the current cell and advances to the next one.
45        Run,
46        /// Runs the current cell without advancing.
47        RunInPlace,
48        /// Clears all outputs in the REPL.
49        ClearOutputs,
50        /// Clears the output of the cell at the current cursor position.
51        ClearCurrentOutput,
52        /// Opens the REPL sessions panel.
53        Sessions,
54        /// Interrupts the currently running kernel.
55        Interrupt,
56        /// Shuts down the current kernel.
57        Shutdown,
58        /// Restarts the current kernel.
59        Restart,
60        /// Refreshes the list of available kernelspecs.
61        RefreshKernelspecs
62    ]
63);
64
65pub fn init(cx: &mut App) {
66    cx.observe_new(
67        |workspace: &mut Workspace, _window, _cx: &mut Context<Workspace>| {
68            workspace.register_action(|workspace, _: &Sessions, window, cx| {
69                let existing = workspace
70                    .active_pane()
71                    .read(cx)
72                    .items()
73                    .find_map(|item| item.downcast::<ReplSessionsPage>());
74
75                if let Some(existing) = existing {
76                    workspace.activate_item(&existing, true, true, window, cx);
77                } else {
78                    let repl_sessions_page = ReplSessionsPage::new(window, cx);
79                    workspace.add_item_to_active_pane(
80                        Box::new(repl_sessions_page),
81                        None,
82                        true,
83                        window,
84                        cx,
85                    )
86                }
87            });
88
89            workspace.register_action(|_workspace, _: &RefreshKernelspecs, _, cx| {
90                let store = ReplStore::global(cx);
91                store.update(cx, |store, cx| {
92                    store.refresh_kernelspecs(cx).detach();
93                });
94            });
95        },
96    )
97    .detach();
98
99    cx.observe_new(
100        move |editor: &mut Editor, window, cx: &mut Context<Editor>| {
101            let Some(window) = window else {
102                return;
103            };
104
105            if !editor.use_modal_editing() || !editor.buffer().read(cx).is_singleton() {
106                return;
107            }
108
109            cx.defer_in(window, |editor, _window, cx| {
110                let project = editor.project().cloned();
111
112                let is_valid_project = project
113                    .as_ref()
114                    .map(|project| {
115                        let p = project.read(cx);
116                        !p.is_via_collab()
117                    })
118                    .unwrap_or(false);
119
120                if !is_valid_project {
121                    return;
122                }
123
124                let buffer = editor.buffer().read(cx).as_singleton();
125
126                let editor_handle = cx.entity().downgrade();
127
128                // Subscribe to the buffer's `LanguageChanged` events so remote projects,
129                // where language detection can complete after the editor is observed,
130                // still trigger a kernelspec refresh. Without this the REPL UI stays
131                // hidden until something else populates the global kernel list.
132                if let Some((buffer, project)) = buffer.zip(project) {
133                    refresh_python_kernelspecs_for_buffer(&buffer, &project, cx);
134
135                    cx.subscribe(&buffer, move |_editor, buffer, event, cx| {
136                        if let BufferEvent::LanguageChanged(_) = event {
137                            refresh_python_kernelspecs_for_buffer(&buffer, &project, cx);
138                        }
139                    })
140                    .detach();
141                }
142
143                editor
144                    .register_action({
145                        let editor_handle = editor_handle.clone();
146                        move |_: &Run, window, cx| {
147                            if !JupyterSettings::enabled(cx) {
148                                return;
149                            }
150
151                            crate::run(editor_handle.clone(), true, window, cx).log_err();
152                        }
153                    })
154                    .detach();
155
156                editor
157                    .register_action({
158                        move |_: &RunInPlace, window, cx| {
159                            if !JupyterSettings::enabled(cx) {
160                                return;
161                            }
162
163                            crate::run(editor_handle.clone(), false, window, cx).log_err();
164                        }
165                    })
166                    .detach();
167            });
168        },
169    )
170    .detach();
171}
172
173pub struct ReplSessionsPage {
174    focus_handle: FocusHandle,
175    _subscriptions: Vec<Subscription>,
176}
177
178impl ReplSessionsPage {
179    pub fn new(window: &mut Window, cx: &mut Context<Workspace>) -> Entity<Self> {
180        cx.new(|cx| {
181            let focus_handle = cx.focus_handle();
182
183            let subscriptions = vec![
184                cx.on_focus_in(&focus_handle, window, |_this, _window, cx| cx.notify()),
185                cx.on_focus_out(&focus_handle, window, |_this, _event, _window, cx| {
186                    cx.notify()
187                }),
188            ];
189
190            Self {
191                focus_handle,
192                _subscriptions: subscriptions,
193            }
194        })
195    }
196}
197
198impl EventEmitter<ItemEvent> for ReplSessionsPage {}
199
200impl Focusable for ReplSessionsPage {
201    fn focus_handle(&self, _cx: &App) -> FocusHandle {
202        self.focus_handle.clone()
203    }
204}
205
206impl Item for ReplSessionsPage {
207    type Event = ItemEvent;
208
209    fn tab_content_text(&self, _detail: usize, _cx: &App) -> SharedString {
210        "REPL Sessions".into()
211    }
212
213    fn telemetry_event_text(&self) -> Option<&'static str> {
214        Some("REPL Session Started")
215    }
216
217    fn show_toolbar(&self) -> bool {
218        false
219    }
220
221    fn to_item_events(event: &Self::Event, f: &mut dyn FnMut(workspace::item::ItemEvent)) {
222        f(*event)
223    }
224}
225
226impl Render for ReplSessionsPage {
227    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
228        let store = ReplStore::global(cx);
229
230        let (kernel_specifications, sessions) = store.update(cx, |store, cx| {
231            store.ensure_kernelspecs(cx);
232            (
233                store
234                    .pure_jupyter_kernel_specifications()
235                    .cloned()
236                    .collect::<Vec<_>>(),
237                store.sessions().cloned().collect::<Vec<_>>(),
238            )
239        });
240
241        // When there are no kernel specifications, show a link to the Zed docs explaining how to
242        // install kernels. It can be assumed they don't have a running kernel if we have no
243        // specifications.
244        if kernel_specifications.is_empty() {
245            let instructions = "To start interactively running code in your editor, you need to install and configure Jupyter kernels.";
246
247            return ReplSessionsContainer::new("No Jupyter Kernels Available")
248                .child(Label::new(instructions))
249                .child(
250                    h_flex().w_full().p_4().justify_center().gap_2().child(
251                        ButtonLike::new("install-kernels")
252                            .style(ButtonStyle::Filled)
253                            .size(ButtonSize::Large)
254                            .layer(ElevationIndex::ModalSurface)
255                            .child(Label::new("Install Kernels"))
256                            .on_click(move |_, _, cx| {
257                                cx.open_url(
258                                    "https://zed.dev/docs/repl#language-specific-instructions",
259                                )
260                            }),
261                    ),
262                );
263        }
264
265        // When there are no sessions, show the command to run code in an editor
266        if sessions.is_empty() {
267            let instructions = "To run code in a Jupyter kernel, select some code and use the 'repl::Run' command.";
268
269            return ReplSessionsContainer::new("No Jupyter Kernel Sessions").child(
270                v_flex()
271                    .child(Label::new(instructions))
272                    .child(KeyBinding::for_action(&Run, cx)),
273            );
274        }
275
276        ReplSessionsContainer::new("Jupyter Kernel Sessions").children(sessions)
277    }
278}
279
280#[derive(IntoElement)]
281struct ReplSessionsContainer {
282    title: SharedString,
283    children: Vec<AnyElement>,
284}
285
286impl ReplSessionsContainer {
287    pub fn new(title: impl Into<SharedString>) -> Self {
288        Self {
289            title: title.into(),
290            children: Vec::new(),
291        }
292    }
293}
294
295impl ParentElement for ReplSessionsContainer {
296    fn extend(&mut self, elements: impl IntoIterator<Item = AnyElement>) {
297        self.children.extend(elements)
298    }
299}
300
301impl RenderOnce for ReplSessionsContainer {
302    fn render(self, _window: &mut Window, _cx: &mut App) -> impl IntoElement {
303        v_flex()
304            .p_4()
305            .gap_2()
306            .size_full()
307            .child(Label::new(self.title).size(LabelSize::Large))
308            .children(self.children)
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315
316    use std::{path::PathBuf, sync::Arc};
317
318    use async_trait::async_trait;
319    use collections::HashMap;
320    use editor::EditorMode;
321    use gpui::TestAppContext;
322    use language::{
323        Language, LanguageConfig, LanguageMatcher, LanguageName, ManifestName, Toolchain,
324        ToolchainList, ToolchainLister, ToolchainMetadata,
325    };
326    use multi_buffer::MultiBuffer;
327    use task::ShellKind;
328    use util::path;
329
330    struct TestPythonToolchainLister;
331
332    #[async_trait]
333    impl ToolchainLister for TestPythonToolchainLister {
334        async fn list(
335            &self,
336            _worktree_root: PathBuf,
337            _subroot_relative_path: Arc<util::rel_path::RelPath>,
338            _project_env: Option<HashMap<String, String>>,
339        ) -> ToolchainList {
340            ToolchainList {
341                toolchains: vec![Toolchain {
342                    name: SharedString::new_static("Test Python"),
343                    path: SharedString::new_static("/test/python"),
344                    language_name: LanguageName::new_static("Python"),
345                    as_json: serde_json::Value::Null,
346                }],
347                ..Default::default()
348            }
349        }
350
351        async fn resolve(
352            &self,
353            _path: PathBuf,
354            _project_env: Option<HashMap<String, String>>,
355        ) -> anyhow::Result<Toolchain> {
356            anyhow::bail!("not implemented")
357        }
358
359        fn activation_script(
360            &self,
361            _toolchain: &Toolchain,
362            _shell: ShellKind,
363            _cx: &App,
364        ) -> futures::future::BoxFuture<'static, Vec<String>> {
365            Box::pin(async { Vec::new() })
366        }
367
368        fn meta(&self) -> ToolchainMetadata {
369            ToolchainMetadata {
370                term: SharedString::new_static("Python"),
371                new_toolchain_placeholder: SharedString::default(),
372                manifest_name: ManifestName::from(SharedString::new_static("pyproject.toml")),
373            }
374        }
375    }
376
377    #[gpui::test]
378    async fn test_refreshes_python_kernelspecs_when_buffer_language_changes(
379        cx: &mut TestAppContext,
380    ) {
381        cx.update(|cx| {
382            settings::init(cx);
383            theme_settings::init(theme::LoadThemes::JustBase, cx);
384            editor::init(cx);
385        });
386
387        let fs = project::FakeFs::new(cx.background_executor.clone());
388        fs.insert_tree(
389            path!("/project"),
390            serde_json::json!({
391                "main.txt": "print('hi')",
392            }),
393        )
394        .await;
395
396        let project = project::Project::test(fs.clone(), [path!("/project").as_ref()], cx).await;
397        let python = Arc::new(
398            Language::new(
399                LanguageConfig {
400                    name: "Python".into(),
401                    matcher: (LanguageMatcher {
402                        path_suffixes: vec!["py".to_string()],
403                        ..Default::default()
404                    })
405                    .into(),
406                    ..Default::default()
407                },
408                None,
409            )
410            .with_manifest(Some(ManifestName::from(SharedString::new_static(
411                "pyproject.toml",
412            ))))
413            .with_toolchain_lister(Some(Arc::new(TestPythonToolchainLister))),
414        );
415        project.read_with(cx, |project, _cx| {
416            project.languages().add(python.clone());
417        });
418
419        cx.update(|cx| crate::init(fs, cx));
420
421        let buffer = project
422            .update(cx, |project, cx| {
423                project.open_local_buffer(path!("/project/main.txt"), cx)
424            })
425            .await
426            .expect("failed to open buffer");
427
428        let worktree_id = buffer
429            .read_with(cx, |buffer, cx| {
430                buffer.project_path(cx).map(|path| path.worktree_id)
431            })
432            .expect("buffer should have a project path");
433
434        cx.add_window(|window, cx| {
435            let multi_buffer = MultiBuffer::build_from_buffer(buffer.clone(), cx);
436            Editor::new(
437                EditorMode::full(),
438                multi_buffer,
439                Some(project.clone()),
440                window,
441                cx,
442            )
443        });
444        cx.run_until_parked();
445
446        let store = cx.update(|cx| ReplStore::global(cx));
447        assert!(!cx.update(|cx| store.read(cx).has_python_kernelspecs(worktree_id)));
448
449        buffer.update(cx, |buffer, cx| {
450            buffer.set_language(Some(python), cx);
451        });
452        cx.run_until_parked();
453
454        assert!(cx.update(|cx| store.read(cx).has_python_kernelspecs(worktree_id)));
455    }
456}
457