1use std::future::Future;
2use std::sync::Arc;
3
4use anyhow::{Context as _, Result};
5use collections::HashMap;
6use command_palette_hooks::CommandPaletteFilter;
7use gpui::{App, Context, Entity, EntityId, Global, Subscription, Task, prelude::*};
8use jupyter_websocket_client::RemoteServer;
9use language::Language;
10use project::{Fs, Project, WorktreeId};
11use settings::{Settings, SettingsStore};
12
13use crate::kernels::{
14 Kernel, list_remote_kernelspecs, local_kernel_specifications, python_env_kernel_specifications,
15};
16use crate::{JupyterSettings, KernelSpecification, Session};
17
18struct GlobalReplStore(Entity<ReplStore>);
19
20impl Global for GlobalReplStore {}
21
22pub struct ReplStore {
23 fs: Arc<dyn Fs>,
24 enabled: bool,
25 sessions: HashMap<EntityId, Entity<Session>>,
26 kernel_specifications: Vec<KernelSpecification>,
27 selected_kernel_for_worktree: HashMap<WorktreeId, KernelSpecification>,
28 kernel_specifications_for_worktree: HashMap<WorktreeId, Vec<KernelSpecification>>,
29 _subscriptions: Vec<Subscription>,
30}
31
32impl ReplStore {
33 const NAMESPACE: &'static str = "repl";
34
35 pub(crate) fn init(fs: Arc<dyn Fs>, cx: &mut App) {
36 let store = cx.new(move |cx| Self::new(fs, cx));
37
38 #[cfg(not(feature = "test-support"))]
39 store
40 .update(cx, |store, cx| store.refresh_kernelspecs(cx))
41 .detach_and_log_err(cx);
42
43 cx.set_global(GlobalReplStore(store))
44 }
45
46 pub fn global(cx: &App) -> Entity<Self> {
47 cx.global::<GlobalReplStore>().0.clone()
48 }
49
50 pub fn new(fs: Arc<dyn Fs>, cx: &mut Context<Self>) -> Self {
51 let subscriptions = vec![
52 cx.observe_global::<SettingsStore>(move |this, cx| {
53 this.set_enabled(JupyterSettings::enabled(cx), cx);
54 }),
55 cx.on_app_quit(Self::shutdown_all_sessions),
56 ];
57
58 let this = Self {
59 fs,
60 enabled: JupyterSettings::enabled(cx),
61 sessions: HashMap::default(),
62 kernel_specifications: Vec::new(),
63 _subscriptions: subscriptions,
64 kernel_specifications_for_worktree: HashMap::default(),
65 selected_kernel_for_worktree: HashMap::default(),
66 };
67 this.on_enabled_changed(cx);
68 this
69 }
70
71 pub fn fs(&self) -> &Arc<dyn Fs> {
72 &self.fs
73 }
74
75 pub fn is_enabled(&self) -> bool {
76 self.enabled
77 }
78
79 pub fn kernel_specifications_for_worktree(
80 &self,
81 worktree_id: WorktreeId,
82 ) -> impl Iterator<Item = &KernelSpecification> {
83 self.kernel_specifications_for_worktree
84 .get(&worktree_id)
85 .into_iter()
86 .flat_map(|specs| specs.iter())
87 .chain(self.kernel_specifications.iter())
88 }
89
90 pub fn pure_jupyter_kernel_specifications(&self) -> impl Iterator<Item = &KernelSpecification> {
91 self.kernel_specifications.iter()
92 }
93
94 pub fn sessions(&self) -> impl Iterator<Item = &Entity<Session>> {
95 self.sessions.values()
96 }
97
98 fn set_enabled(&mut self, enabled: bool, cx: &mut Context<Self>) {
99 if self.enabled == enabled {
100 return;
101 }
102
103 self.enabled = enabled;
104 self.on_enabled_changed(cx);
105 }
106
107 fn on_enabled_changed(&self, cx: &mut Context<Self>) {
108 if !self.enabled {
109 CommandPaletteFilter::update_global(cx, |filter, _cx| {
110 filter.hide_namespace(Self::NAMESPACE);
111 });
112
113 return;
114 }
115
116 CommandPaletteFilter::update_global(cx, |filter, _cx| {
117 filter.show_namespace(Self::NAMESPACE);
118 });
119
120 cx.notify();
121 }
122
123 pub fn refresh_python_kernelspecs(
124 &mut self,
125 worktree_id: WorktreeId,
126 project: &Entity<Project>,
127 cx: &mut Context<Self>,
128 ) -> Task<Result<()>> {
129 let kernel_specifications = python_env_kernel_specifications(project, worktree_id, cx);
130 cx.spawn(async move |this, cx| {
131 let kernel_specifications = kernel_specifications
132 .await
133 .context("getting python kernelspecs")?;
134
135 this.update(cx, |this, cx| {
136 this.kernel_specifications_for_worktree
137 .insert(worktree_id, kernel_specifications);
138 cx.notify();
139 })
140 })
141 }
142
143 fn get_remote_kernel_specifications(
144 &self,
145 cx: &mut Context<Self>,
146 ) -> Option<Task<Result<Vec<KernelSpecification>>>> {
147 match (
148 std::env::var("JUPYTER_SERVER"),
149 std::env::var("JUPYTER_TOKEN"),
150 ) {
151 (Ok(server), Ok(token)) => {
152 let remote_server = RemoteServer {
153 base_url: server,
154 token,
155 };
156 let http_client = cx.http_client();
157 Some(cx.spawn(async move |_, _| {
158 list_remote_kernelspecs(remote_server, http_client)
159 .await
160 .map(|specs| specs.into_iter().map(KernelSpecification::Remote).collect())
161 }))
162 }
163 _ => None,
164 }
165 }
166
167 pub fn refresh_kernelspecs(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
168 let local_kernel_specifications = local_kernel_specifications(self.fs.clone());
169
170 let remote_kernel_specifications = self.get_remote_kernel_specifications(cx);
171
172 let all_specs = cx.background_spawn(async move {
173 let mut all_specs = local_kernel_specifications
174 .await?
175 .into_iter()
176 .map(KernelSpecification::Jupyter)
177 .collect::<Vec<_>>();
178
179 if let Some(remote_task) = remote_kernel_specifications
180 && let Ok(remote_specs) = remote_task.await
181 {
182 all_specs.extend(remote_specs);
183 }
184
185 anyhow::Ok(all_specs)
186 });
187
188 cx.spawn(async move |this, cx| {
189 let all_specs = all_specs.await;
190
191 if let Ok(specs) = all_specs {
192 this.update(cx, |this, cx| {
193 this.kernel_specifications = specs;
194 cx.notify();
195 })
196 .ok();
197 }
198
199 anyhow::Ok(())
200 })
201 }
202
203 pub fn set_active_kernelspec(
204 &mut self,
205 worktree_id: WorktreeId,
206 kernelspec: KernelSpecification,
207 _cx: &mut Context<Self>,
208 ) {
209 self.selected_kernel_for_worktree
210 .insert(worktree_id, kernelspec);
211 }
212
213 pub fn active_kernelspec(
214 &self,
215 worktree_id: WorktreeId,
216 language_at_cursor: Option<Arc<Language>>,
217 cx: &App,
218 ) -> Option<KernelSpecification> {
219 let selected_kernelspec = self.selected_kernel_for_worktree.get(&worktree_id).cloned();
220
221 if let Some(language_at_cursor) = language_at_cursor {
222 selected_kernelspec.or_else(|| {
223 self.kernelspec_legacy_by_lang_only(worktree_id, language_at_cursor, cx)
224 })
225 } else {
226 selected_kernelspec
227 }
228 }
229
230 fn kernelspec_legacy_by_lang_only(
231 &self,
232 worktree_id: WorktreeId,
233 language_at_cursor: Arc<Language>,
234 cx: &App,
235 ) -> Option<KernelSpecification> {
236 let settings = JupyterSettings::get_global(cx);
237 let selected_kernel = settings
238 .kernel_selections
239 .get(language_at_cursor.code_fence_block_name().as_ref());
240
241 let found_by_name = self
242 .kernel_specifications_for_worktree(worktree_id)
243 .find(|runtime_specification| {
244 if let (Some(selected), KernelSpecification::Jupyter(runtime_specification)) =
245 (selected_kernel, runtime_specification)
246 {
247 // Top priority is the selected kernel
248 return runtime_specification.name.to_lowercase() == selected.to_lowercase();
249 }
250 false
251 })
252 .cloned();
253
254 if let Some(found_by_name) = found_by_name {
255 return Some(found_by_name);
256 }
257
258 self.kernel_specifications_for_worktree(worktree_id)
259 .find(|kernel_option| match kernel_option {
260 KernelSpecification::Jupyter(runtime_specification) => {
261 runtime_specification.kernelspec.language.to_lowercase()
262 == language_at_cursor.code_fence_block_name().to_lowercase()
263 }
264 KernelSpecification::PythonEnv(runtime_specification) => {
265 runtime_specification.kernelspec.language.to_lowercase()
266 == language_at_cursor.code_fence_block_name().to_lowercase()
267 }
268 KernelSpecification::Remote(remote_spec) => {
269 remote_spec.kernelspec.language.to_lowercase()
270 == language_at_cursor.code_fence_block_name().to_lowercase()
271 }
272 })
273 .cloned()
274 }
275
276 pub fn get_session(&self, entity_id: EntityId) -> Option<&Entity<Session>> {
277 self.sessions.get(&entity_id)
278 }
279
280 pub fn insert_session(&mut self, entity_id: EntityId, session: Entity<Session>) {
281 self.sessions.insert(entity_id, session);
282 }
283
284 pub fn remove_session(&mut self, entity_id: EntityId) {
285 self.sessions.remove(&entity_id);
286 }
287
288 fn shutdown_all_sessions(
289 &mut self,
290 cx: &mut Context<Self>,
291 ) -> impl Future<Output = ()> + use<> {
292 for session in self.sessions.values() {
293 session.update(cx, |session, _cx| {
294 if let Kernel::RunningKernel(mut kernel) =
295 std::mem::replace(&mut session.kernel, Kernel::Shutdown)
296 {
297 kernel.kill();
298 }
299 });
300 }
301 self.sessions.clear();
302 futures::future::ready(())
303 }
304
305 #[cfg(test)]
306 pub fn set_kernel_specs_for_testing(
307 &mut self,
308 specs: Vec<KernelSpecification>,
309 cx: &mut Context<Self>,
310 ) {
311 self.kernel_specifications = specs;
312 cx.notify();
313 }
314}