tenant.openagents/omega

1use std::future::Future;
2use std::sync::Arc;
3
4use anyhow::{Context as _, Result};
5use collections::{HashMap, HashSet};
6use command_palette_hooks::CommandPaletteFilter;
7use gpui::{
8    App, Context, Entity, EntityId, Global, SharedString, Subscription, Task, TaskExt, prelude::*,
9};
10use jupyter_websocket_client::RemoteServer;
11use language::{Language, LanguageName};
12use project::{Fs, Project, ProjectPath, WorktreeId};
13use remote::RemoteConnectionOptions;
14use settings::{Settings, SettingsStore};
15use util::rel_path::RelPath;
16
17use crate::kernels::{
18    Kernel, PythonEnvKernelSpecification, list_remote_kernelspecs, local_kernel_specifications,
19    python_env_kernel_specifications, wsl_kernel_specifications,
20};
21use crate::{JupyterSettings, KernelSpecification, Session};
22
23struct GlobalReplStore(Entity<ReplStore>);
24
25impl Global for GlobalReplStore {}
26
27pub struct ReplStore {
28    fs: Arc<dyn Fs>,
29    enabled: bool,
30    sessions: HashMap<EntityId, Entity<Session>>,
31    kernel_specifications: Vec<KernelSpecification>,
32    kernelspecs_initialized: bool,
33    selected_kernel_for_worktree: HashMap<WorktreeId, KernelSpecification>,
34    kernel_specifications_for_worktree: HashMap<WorktreeId, Vec<KernelSpecification>>,
35    active_python_toolchain_for_worktree: HashMap<WorktreeId, SharedString>,
36    remote_worktrees: HashSet<WorktreeId>,
37    fetching_python_kernelspecs: HashSet<WorktreeId>,
38    _subscriptions: Vec<Subscription>,
39}
40
41impl ReplStore {
42    const NAMESPACE: &'static str = "repl";
43
44    pub(crate) fn init(fs: Arc<dyn Fs>, cx: &mut App) {
45        let store = cx.new(move |cx| Self::new(fs, cx));
46        cx.set_global(GlobalReplStore(store))
47    }
48
49    pub fn global(cx: &App) -> Entity<Self> {
50        cx.global::<GlobalReplStore>().0.clone()
51    }
52
53    pub fn new(fs: Arc<dyn Fs>, cx: &mut Context<Self>) -> Self {
54        let subscriptions = vec![
55            cx.observe_global::<SettingsStore>(move |this, cx| {
56                this.set_enabled(JupyterSettings::enabled(cx), cx);
57            }),
58            cx.on_app_quit(Self::shutdown_all_sessions),
59        ];
60
61        let this = Self {
62            fs,
63            enabled: JupyterSettings::enabled(cx),
64            sessions: HashMap::default(),
65            kernel_specifications: Vec::new(),
66            kernelspecs_initialized: false,
67            _subscriptions: subscriptions,
68            kernel_specifications_for_worktree: HashMap::default(),
69            selected_kernel_for_worktree: HashMap::default(),
70            active_python_toolchain_for_worktree: HashMap::default(),
71            remote_worktrees: HashSet::default(),
72            fetching_python_kernelspecs: HashSet::default(),
73        };
74        this.on_enabled_changed(cx);
75        this
76    }
77
78    pub fn fs(&self) -> &Arc<dyn Fs> {
79        &self.fs
80    }
81
82    pub fn is_enabled(&self) -> bool {
83        self.enabled
84    }
85
86    pub fn has_python_kernelspecs(&self, worktree_id: WorktreeId) -> bool {
87        self.kernel_specifications_for_worktree
88            .contains_key(&worktree_id)
89    }
90
91    pub fn kernel_specifications_for_worktree(
92        &self,
93        worktree_id: WorktreeId,
94    ) -> impl Iterator<Item = &KernelSpecification> {
95        let global_specs = if self.remote_worktrees.contains(&worktree_id) {
96            None
97        } else {
98            Some(self.kernel_specifications.iter())
99        };
100
101        self.kernel_specifications_for_worktree
102            .get(&worktree_id)
103            .into_iter()
104            .flat_map(|specs| specs.iter())
105            .chain(global_specs.into_iter().flatten())
106    }
107
108    pub fn pure_jupyter_kernel_specifications(&self) -> impl Iterator<Item = &KernelSpecification> {
109        self.kernel_specifications.iter()
110    }
111
112    pub fn sessions(&self) -> impl Iterator<Item = &Entity<Session>> {
113        self.sessions.values()
114    }
115
116    fn set_enabled(&mut self, enabled: bool, cx: &mut Context<Self>) {
117        if self.enabled == enabled {
118            return;
119        }
120
121        self.enabled = enabled;
122        self.on_enabled_changed(cx);
123    }
124
125    fn on_enabled_changed(&self, cx: &mut Context<Self>) {
126        if !self.enabled {
127            CommandPaletteFilter::update_global(cx, |filter, _cx| {
128                filter.hide_namespace(Self::NAMESPACE);
129            });
130
131            return;
132        }
133
134        CommandPaletteFilter::update_global(cx, |filter, _cx| {
135            filter.show_namespace(Self::NAMESPACE);
136        });
137
138        cx.notify();
139    }
140
141    pub fn mark_ipykernel_installed(
142        &mut self,
143        cx: &mut Context<Self>,
144        spec: &PythonEnvKernelSpecification,
145    ) {
146        for specs in self.kernel_specifications_for_worktree.values_mut() {
147            for kernel_spec in specs.iter_mut() {
148                if let KernelSpecification::PythonEnv(env_spec) = kernel_spec {
149                    if env_spec == spec {
150                        env_spec.has_ipykernel = true;
151                    }
152                }
153            }
154        }
155        cx.notify();
156    }
157
158    pub fn refresh_python_kernelspecs(
159        &mut self,
160        worktree_id: WorktreeId,
161        project: &Entity<Project>,
162        cx: &mut Context<Self>,
163    ) -> Task<Result<()>> {
164        if !self.fetching_python_kernelspecs.insert(worktree_id) {
165            return Task::ready(Ok(()));
166        }
167
168        let is_remote = project.read(cx).is_remote();
169        // WSL does require access to global kernel specs, so we only exclude remote worktrees that aren't WSL.
170        // TODO: a better way to handle WSL vs SSH/remote projects,
171        let is_wsl_remote = project
172            .read(cx)
173            .remote_connection_options(cx)
174            .map_or(false, |opts| {
175                matches!(opts, RemoteConnectionOptions::Wsl(_))
176            });
177        let kernel_specifications_task = python_env_kernel_specifications(project, worktree_id, cx);
178        let active_toolchain = project.read(cx).active_toolchain(
179            ProjectPath {
180                worktree_id,
181                path: RelPath::empty_arc(),
182            },
183            LanguageName::new_static("Python"),
184            cx,
185        );
186
187        cx.spawn(async move |this, cx| {
188            let kernel_specifications_res = kernel_specifications_task.await;
189
190            this.update(cx, |this, _cx| {
191                this.fetching_python_kernelspecs.remove(&worktree_id);
192            })
193            .ok();
194
195            let kernel_specifications =
196                kernel_specifications_res.context("getting python kernelspecs")?;
197
198            let active_toolchain_path = active_toolchain.await.map(|toolchain| toolchain.path);
199
200            this.update(cx, |this, cx| {
201                this.kernel_specifications_for_worktree
202                    .insert(worktree_id, kernel_specifications);
203                if let Some(path) = active_toolchain_path {
204                    this.active_python_toolchain_for_worktree
205                        .insert(worktree_id, path);
206                }
207                if is_remote && !is_wsl_remote {
208                    this.remote_worktrees.insert(worktree_id);
209                } else {
210                    this.remote_worktrees.remove(&worktree_id);
211                }
212                cx.notify();
213            })
214        })
215    }
216
217    fn get_remote_kernel_specifications(
218        &self,
219        cx: &mut Context<Self>,
220    ) -> Option<Task<Result<Vec<KernelSpecification>>>> {
221        match (
222            std::env::var("JUPYTER_SERVER"),
223            std::env::var("JUPYTER_TOKEN"),
224        ) {
225            (Ok(server), Ok(token)) => {
226                let remote_server = RemoteServer {
227                    base_url: server,
228                    token,
229                };
230                let http_client = cx.http_client();
231                Some(cx.spawn(async move |_, _| {
232                    list_remote_kernelspecs(remote_server, http_client)
233                        .await
234                        .map(|specs| {
235                            specs
236                                .into_iter()
237                                .map(KernelSpecification::JupyterServer)
238                                .collect()
239                        })
240                }))
241            }
242            _ => None,
243        }
244    }
245
246    pub fn ensure_kernelspecs(&mut self, cx: &mut Context<Self>) {
247        if self.kernelspecs_initialized {
248            return;
249        }
250        self.kernelspecs_initialized = true;
251        self.refresh_kernelspecs(cx).detach_and_log_err(cx);
252    }
253
254    pub fn refresh_kernelspecs(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
255        let local_kernel_specifications = local_kernel_specifications(self.fs.clone());
256        let wsl_kernel_specifications = wsl_kernel_specifications(cx.background_executor().clone());
257        let remote_kernel_specifications = self.get_remote_kernel_specifications(cx);
258
259        let all_specs = cx.background_spawn(async move {
260            let mut all_specs = local_kernel_specifications
261                .await?
262                .into_iter()
263                .map(KernelSpecification::Jupyter)
264                .collect::<Vec<_>>();
265
266            if let Ok(wsl_specs) = wsl_kernel_specifications.await {
267                all_specs.extend(wsl_specs);
268            }
269
270            if let Some(remote_task) = remote_kernel_specifications
271                && let Ok(remote_specs) = remote_task.await
272            {
273                all_specs.extend(remote_specs);
274            }
275
276            anyhow::Ok(all_specs)
277        });
278
279        cx.spawn(async move |this, cx| {
280            let all_specs = all_specs.await;
281
282            if let Ok(specs) = all_specs {
283                this.update(cx, |this, cx| {
284                    this.kernel_specifications = specs;
285                    cx.notify();
286                })
287                .ok();
288            }
289
290            anyhow::Ok(())
291        })
292    }
293
294    pub fn set_active_kernelspec(
295        &mut self,
296        worktree_id: WorktreeId,
297        kernelspec: KernelSpecification,
298        _cx: &mut Context<Self>,
299    ) {
300        self.selected_kernel_for_worktree
301            .insert(worktree_id, kernelspec);
302    }
303
304    pub fn active_python_toolchain_path(&self, worktree_id: WorktreeId) -> Option<&SharedString> {
305        self.active_python_toolchain_for_worktree.get(&worktree_id)
306    }
307
308    pub fn selected_kernel(&self, worktree_id: WorktreeId) -> Option<&KernelSpecification> {
309        self.selected_kernel_for_worktree.get(&worktree_id)
310    }
311
312    pub fn is_recommended_kernel(
313        &self,
314        worktree_id: WorktreeId,
315        spec: &KernelSpecification,
316    ) -> bool {
317        if let Some(active_path) = self.active_python_toolchain_path(worktree_id) {
318            spec.path().as_ref() == active_path.as_ref()
319        } else {
320            false
321        }
322    }
323
324    pub fn active_kernelspec(
325        &self,
326        worktree_id: WorktreeId,
327        language_at_cursor: Option<Arc<Language>>,
328        cx: &App,
329    ) -> Option<KernelSpecification> {
330        if let Some(selected) = self.selected_kernel_for_worktree.get(&worktree_id).cloned() {
331            return Some(selected);
332        }
333
334        let language_at_cursor = language_at_cursor?;
335
336        // Prefer the recommended (active toolchain) kernel if it has ipykernel
337        if let Some(active_path) = self.active_python_toolchain_path(worktree_id) {
338            let recommended = self
339                .kernel_specifications_for_worktree(worktree_id)
340                .find(|spec| {
341                    spec.has_ipykernel()
342                        && language_at_cursor.matches_kernel_language(spec.language().as_ref())
343                        && spec.path().as_ref() == active_path.as_ref()
344                })
345                .cloned();
346            if recommended.is_some() {
347                return recommended;
348            }
349        }
350
351        // Then try the first PythonEnv with ipykernel matching the language
352        let python_env = self
353            .kernel_specifications_for_worktree(worktree_id)
354            .find(|spec| {
355                matches!(spec, KernelSpecification::PythonEnv(_))
356                    && spec.has_ipykernel()
357                    && language_at_cursor.matches_kernel_language(spec.language().as_ref())
358            })
359            .cloned();
360        if python_env.is_some() {
361            return python_env;
362        }
363
364        // Fall back to legacy name-based and language-based matching
365        self.kernelspec_legacy_by_lang_only(worktree_id, language_at_cursor, cx)
366    }
367
368    fn kernelspec_legacy_by_lang_only(
369        &self,
370        worktree_id: WorktreeId,
371        language_at_cursor: Arc<Language>,
372        cx: &App,
373    ) -> Option<KernelSpecification> {
374        let settings = JupyterSettings::get_global(cx);
375        let selected_kernel = settings
376            .kernel_selections
377            .get(language_at_cursor.code_fence_block_name().as_ref());
378
379        let found_by_name = self
380            .kernel_specifications_for_worktree(worktree_id)
381            .find(|runtime_specification| {
382                if let (Some(selected), KernelSpecification::Jupyter(runtime_specification)) =
383                    (selected_kernel, runtime_specification)
384                {
385                    return runtime_specification.name.to_lowercase() == selected.to_lowercase();
386                }
387                false
388            })
389            .cloned();
390
391        if let Some(found_by_name) = found_by_name {
392            return Some(found_by_name);
393        }
394
395        self.kernel_specifications_for_worktree(worktree_id)
396            .find(|spec| {
397                spec.has_ipykernel()
398                    && language_at_cursor.matches_kernel_language(spec.language().as_ref())
399            })
400            .cloned()
401    }
402
403    pub fn get_session(&self, entity_id: EntityId) -> Option<&Entity<Session>> {
404        self.sessions.get(&entity_id)
405    }
406
407    pub fn insert_session(&mut self, entity_id: EntityId, session: Entity<Session>) {
408        self.sessions.insert(entity_id, session);
409    }
410
411    pub fn remove_session(&mut self, entity_id: EntityId) {
412        self.sessions.remove(&entity_id);
413    }
414
415    fn shutdown_all_sessions(
416        &mut self,
417        cx: &mut Context<Self>,
418    ) -> impl Future<Output = ()> + use<> {
419        for session in self.sessions.values() {
420            session.update(cx, |session, _cx| {
421                if let Kernel::RunningKernel(mut kernel) =
422                    std::mem::replace(&mut session.kernel, Kernel::Shutdown)
423                {
424                    kernel.kill();
425                }
426            });
427        }
428        self.sessions.clear();
429        futures::future::ready(())
430    }
431
432    #[cfg(test)]
433    pub fn set_kernel_specs_for_testing(
434        &mut self,
435        specs: Vec<KernelSpecification>,
436        cx: &mut Context<Self>,
437    ) {
438        self.kernel_specifications = specs;
439        cx.notify();
440    }
441}
442