use std::future::Future;use std::sync::Arc; use anyhow::{Context as _, Result};use collections::{HashMap, HashSet};use command_palette_hooks::CommandPaletteFilter;use gpui::{App, Context, Entity, EntityId, Global, SharedString, Subscription, Task, prelude::*};use jupyter_websocket_client::RemoteServer;use language::{Language, LanguageName};use project::{Fs, Project, ProjectPath, WorktreeId};use remote::RemoteConnectionOptions;use settings::{Settings, SettingsStore};use util::rel_path::RelPath; use crate::kernels::{ Kernel, list_remote_kernelspecs, local_kernel_specifications, python_env_kernel_specifications, wsl_kernel_specifications,};use crate::{JupyterSettings, KernelSpecification, Session}; struct GlobalReplStore(Entity<ReplStore>); impl Global for GlobalReplStore {} pub struct ReplStore { fs: Arc<dyn Fs>, enabled: bool, sessions: HashMap<EntityId, Entity<Session>>, kernel_specifications: Vec<KernelSpecification>, kernelspecs_initialized: bool, selected_kernel_for_worktree: HashMap<WorktreeId, KernelSpecification>, kernel_specifications_for_worktree: HashMap<WorktreeId, Vec<KernelSpecification>>, active_python_toolchain_for_worktree: HashMap<WorktreeId, SharedString>, remote_worktrees: HashSet<WorktreeId>, fetching_python_kernelspecs: HashSet<WorktreeId>, _subscriptions: Vec<Subscription>,} impl ReplStore { const NAMESPACE: &'static str = "repl"; pub(crate) fn init(fs: Arc<dyn Fs>, cx: &mut App) { let store = cx.new(move |cx| Self::new(fs, cx)); cx.set_global(GlobalReplStore(store)) } pub fn global(cx: &App) -> Entity<Self> { cx.global::<GlobalReplStore>().0.clone() } pub fn new(fs: Arc<dyn Fs>, cx: &mut Context<Self>) -> Self { let subscriptions = vec![ cx.observe_global::<SettingsStore>(move |this, cx| { this.set_enabled(JupyterSettings::enabled(cx), cx); }), cx.on_app_quit(Self::shutdown_all_sessions), ]; let this = Self { fs, enabled: JupyterSettings::enabled(cx), sessions: HashMap::default(), kernel_specifications: Vec::new(), kernelspecs_initialized: false, _subscriptions: subscriptions, kernel_specifications_for_worktree: HashMap::default(), selected_kernel_for_worktree: HashMap::default(), active_python_toolchain_for_worktree: HashMap::default(), remote_worktrees: HashSet::default(), fetching_python_kernelspecs: HashSet::default(), }; this.on_enabled_changed(cx); this } pub fn fs(&self) -> &Arc<dyn Fs> { &self.fs } pub fn is_enabled(&self) -> bool { self.enabled } pub fn has_python_kernelspecs(&self, worktree_id: WorktreeId) -> bool { self.kernel_specifications_for_worktree .contains_key(&worktree_id) } pub fn kernel_specifications_for_worktree( &self, worktree_id: WorktreeId, ) -> impl Iterator<Item = &KernelSpecification> { let global_specs = if self.remote_worktrees.contains(&worktree_id) { None } else { Some(self.kernel_specifications.iter()) }; self.kernel_specifications_for_worktree .get(&worktree_id) .into_iter() .flat_map(|specs| specs.iter()) .chain(global_specs.into_iter().flatten()) } pub fn pure_jupyter_kernel_specifications(&self) -> impl Iterator<Item = &KernelSpecification> { self.kernel_specifications.iter() } pub fn sessions(&self) -> impl Iterator<Item = &Entity<Session>> { self.sessions.values() } fn set_enabled(&mut self, enabled: bool, cx: &mut Context<Self>) { if self.enabled == enabled { return; } self.enabled = enabled; self.on_enabled_changed(cx); } fn on_enabled_changed(&self, cx: &mut Context<Self>) { if !self.enabled { CommandPaletteFilter::update_global(cx, |filter, _cx| { filter.hide_namespace(Self::NAMESPACE); }); return; } CommandPaletteFilter::update_global(cx, |filter, _cx| { filter.show_namespace(Self::NAMESPACE); }); cx.notify(); } pub fn refresh_python_kernelspecs( &mut self, worktree_id: WorktreeId, project: &Entity<Project>, cx: &mut Context<Self>, ) -> Task<Result<()>> { if !self.fetching_python_kernelspecs.insert(worktree_id) { return Task::ready(Ok(())); } let is_remote = project.read(cx).is_remote(); // WSL does require access to global kernel specs, so we only exclude remote worktrees that aren't WSL. // TODO: a better way to handle WSL vs SSH/remote projects, let is_wsl_remote = project .read(cx) .remote_connection_options(cx) .map_or(false, |opts| { matches!(opts, RemoteConnectionOptions::Wsl(_)) }); let kernel_specifications_task = python_env_kernel_specifications(project, worktree_id, cx); let active_toolchain = project.read(cx).active_toolchain( ProjectPath { worktree_id, path: RelPath::empty().into(), }, LanguageName::new_static("Python"), cx, ); cx.spawn(async move |this, cx| { let kernel_specifications_res = kernel_specifications_task.await; this.update(cx, |this, _cx| { this.fetching_python_kernelspecs.remove(&worktree_id); }) .ok(); let kernel_specifications = kernel_specifications_res.context("getting python kernelspecs")?; let active_toolchain_path = active_toolchain.await.map(|toolchain| toolchain.path); this.update(cx, |this, cx| { this.kernel_specifications_for_worktree .insert(worktree_id, kernel_specifications); if let Some(path) = active_toolchain_path { this.active_python_toolchain_for_worktree .insert(worktree_id, path); } if is_remote && !is_wsl_remote { this.remote_worktrees.insert(worktree_id); } else { this.remote_worktrees.remove(&worktree_id); } cx.notify(); }) }) } fn get_remote_kernel_specifications( &self, cx: &mut Context<Self>, ) -> Option<Task<Result<Vec<KernelSpecification>>>> { match ( std::env::var("JUPYTER_SERVER"), std::env::var("JUPYTER_TOKEN"), ) { (Ok(server), Ok(token)) => { let remote_server = RemoteServer { base_url: server, token, }; let http_client = cx.http_client(); Some(cx.spawn(async move |_, _| { list_remote_kernelspecs(remote_server, http_client) .await .map(|specs| { specs .into_iter() .map(KernelSpecification::JupyterServer) .collect() }) })) } _ => None, } } pub fn ensure_kernelspecs(&mut self, cx: &mut Context<Self>) { if self.kernelspecs_initialized { return; } self.kernelspecs_initialized = true; self.refresh_kernelspecs(cx).detach_and_log_err(cx); } pub fn refresh_kernelspecs(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> { let local_kernel_specifications = local_kernel_specifications(self.fs.clone()); let wsl_kernel_specifications = wsl_kernel_specifications(cx.background_executor().clone()); let remote_kernel_specifications = self.get_remote_kernel_specifications(cx); let all_specs = cx.background_spawn(async move { let mut all_specs = local_kernel_specifications .await? .into_iter() .map(KernelSpecification::Jupyter) .collect::<Vec<_>>(); if let Ok(wsl_specs) = wsl_kernel_specifications.await { all_specs.extend(wsl_specs); } if let Some(remote_task) = remote_kernel_specifications && let Ok(remote_specs) = remote_task.await { all_specs.extend(remote_specs); } anyhow::Ok(all_specs) }); cx.spawn(async move |this, cx| { let all_specs = all_specs.await; if let Ok(specs) = all_specs { this.update(cx, |this, cx| { this.kernel_specifications = specs; cx.notify(); }) .ok(); } anyhow::Ok(()) }) } pub fn set_active_kernelspec( &mut self, worktree_id: WorktreeId, kernelspec: KernelSpecification, _cx: &mut Context<Self>, ) { self.selected_kernel_for_worktree .insert(worktree_id, kernelspec); } pub fn active_python_toolchain_path(&self, worktree_id: WorktreeId) -> Option<&SharedString> { self.active_python_toolchain_for_worktree.get(&worktree_id) } pub fn selected_kernel(&self, worktree_id: WorktreeId) -> Option<&KernelSpecification> { self.selected_kernel_for_worktree.get(&worktree_id) } pub fn is_recommended_kernel( &self, worktree_id: WorktreeId, spec: &KernelSpecification, ) -> bool { if let Some(active_path) = self.active_python_toolchain_path(worktree_id) { spec.path().as_ref() == active_path.as_ref() } else { false } } pub fn active_kernelspec( &self, worktree_id: WorktreeId, language_at_cursor: Option<Arc<Language>>, cx: &App, ) -> Option<KernelSpecification> { if let Some(selected) = self.selected_kernel_for_worktree.get(&worktree_id).cloned() { return Some(selected); } let language_at_cursor = language_at_cursor?; // Prefer the recommended (active toolchain) kernel if it has ipykernel if let Some(active_path) = self.active_python_toolchain_path(worktree_id) { let recommended = self .kernel_specifications_for_worktree(worktree_id) .find(|spec| { spec.has_ipykernel() && language_at_cursor.matches_kernel_language(spec.language().as_ref()) && spec.path().as_ref() == active_path.as_ref() }) .cloned(); if recommended.is_some() { return recommended; } } // Then try the first PythonEnv with ipykernel matching the language let python_env = self .kernel_specifications_for_worktree(worktree_id) .find(|spec| { matches!(spec, KernelSpecification::PythonEnv(_)) && spec.has_ipykernel() && language_at_cursor.matches_kernel_language(spec.language().as_ref()) }) .cloned(); if python_env.is_some() { return python_env; } // Fall back to legacy name-based and language-based matching self.kernelspec_legacy_by_lang_only(worktree_id, language_at_cursor, cx) } fn kernelspec_legacy_by_lang_only( &self, worktree_id: WorktreeId, language_at_cursor: Arc<Language>, cx: &App, ) -> Option<KernelSpecification> { let settings = JupyterSettings::get_global(cx); let selected_kernel = settings .kernel_selections .get(language_at_cursor.code_fence_block_name().as_ref()); let found_by_name = self .kernel_specifications_for_worktree(worktree_id) .find(|runtime_specification| { if let (Some(selected), KernelSpecification::Jupyter(runtime_specification)) = (selected_kernel, runtime_specification) { return runtime_specification.name.to_lowercase() == selected.to_lowercase(); } false }) .cloned(); if let Some(found_by_name) = found_by_name { return Some(found_by_name); } self.kernel_specifications_for_worktree(worktree_id) .find(|spec| { spec.has_ipykernel() && language_at_cursor.matches_kernel_language(spec.language().as_ref()) }) .cloned() } pub fn get_session(&self, entity_id: EntityId) -> Option<&Entity<Session>> { self.sessions.get(&entity_id) } pub fn insert_session(&mut self, entity_id: EntityId, session: Entity<Session>) { self.sessions.insert(entity_id, session); } pub fn remove_session(&mut self, entity_id: EntityId) { self.sessions.remove(&entity_id); } fn shutdown_all_sessions( &mut self, cx: &mut Context<Self>, ) -> impl Future<Output = ()> + use<> { for session in self.sessions.values() { session.update(cx, |session, _cx| { if let Kernel::RunningKernel(mut kernel) = std::mem::replace(&mut session.kernel, Kernel::Shutdown) { kernel.kill(); } }); } self.sessions.clear(); futures::future::ready(()) } #[cfg(test)] pub fn set_kernel_specs_for_testing( &mut self, specs: Vec<KernelSpecification>, cx: &mut Context<Self>, ) { self.kernel_specifications = specs; cx.notify(); }}
use std::future::Future;
use std::sync::Arc;
use anyhow::{Context as _, Result};
use collections::{HashMap, HashSet};
use command_palette_hooks::CommandPaletteFilter;
use gpui::{App, Context, Entity, EntityId, Global, SharedString, Subscription, Task, prelude::*};
use jupyter_websocket_client::RemoteServer;
use language::{Language, LanguageName};
use project::{Fs, Project, ProjectPath, WorktreeId};
use remote::RemoteConnectionOptions;
use settings::{Settings, SettingsStore};
use util::rel_path::RelPath;
use crate::kernels::{
Kernel, list_remote_kernelspecs, local_kernel_specifications, python_env_kernel_specifications,
wsl_kernel_specifications,
};
use crate::{JupyterSettings, KernelSpecification, Session};
struct GlobalReplStore(Entity<ReplStore>);
impl Global for GlobalReplStore {}
pub struct ReplStore {
fs: Arc<dyn Fs>,
enabled: bool,
sessions: HashMap<EntityId, Entity<Session>>,
kernel_specifications: Vec<KernelSpecification>,
kernelspecs_initialized: bool,
selected_kernel_for_worktree: HashMap<WorktreeId, KernelSpecification>,
kernel_specifications_for_worktree: HashMap<WorktreeId, Vec<KernelSpecification>>,
active_python_toolchain_for_worktree: HashMap<WorktreeId, SharedString>,
remote_worktrees: HashSet<WorktreeId>,
fetching_python_kernelspecs: HashSet<WorktreeId>,
_subscriptions: Vec<Subscription>,
}
impl ReplStore {
const NAMESPACE: &'static str = "repl";
pub(crate) fn init(fs: Arc<dyn Fs>, cx: &mut App) {
let store = cx.new(move |cx| Self::new(fs, cx));
cx.set_global(GlobalReplStore(store))
}
pub fn global(cx: &App) -> Entity<Self> {
cx.global::<GlobalReplStore>().0.clone()
}
pub fn new(fs: Arc<dyn Fs>, cx: &mut Context<Self>) -> Self {
let subscriptions = vec![
cx.observe_global::<SettingsStore>(move |this, cx| {
this.set_enabled(JupyterSettings::enabled(cx), cx);
}),
cx.on_app_quit(Self::shutdown_all_sessions),
];
let this = Self {
fs,
enabled: JupyterSettings::enabled(cx),
sessions: HashMap::default(),
kernel_specifications: Vec::new(),
kernelspecs_initialized: false,
_subscriptions: subscriptions,
kernel_specifications_for_worktree: HashMap::default(),
selected_kernel_for_worktree: HashMap::default(),
active_python_toolchain_for_worktree: HashMap::default(),
remote_worktrees: HashSet::default(),
fetching_python_kernelspecs: HashSet::default(),
};
this.on_enabled_changed(cx);
this
}
pub fn fs(&self) -> &Arc<dyn Fs> {
&self.fs
}
pub fn is_enabled(&self) -> bool {
self.enabled
}
pub fn has_python_kernelspecs(&self, worktree_id: WorktreeId) -> bool {
self.kernel_specifications_for_worktree
.contains_key(&worktree_id)
}
pub fn kernel_specifications_for_worktree(
&self,
worktree_id: WorktreeId,
) -> impl Iterator<Item = &KernelSpecification> {
let global_specs = if self.remote_worktrees.contains(&worktree_id) {
None
} else {
Some(self.kernel_specifications.iter())
};
self.kernel_specifications_for_worktree
.get(&worktree_id)
.into_iter()
.flat_map(|specs| specs.iter())
.chain(global_specs.into_iter().flatten())
}
pub fn pure_jupyter_kernel_specifications(&self) -> impl Iterator<Item = &KernelSpecification> {
self.kernel_specifications.iter()
}
pub fn sessions(&self) -> impl Iterator<Item = &Entity<Session>> {
self.sessions.values()
}
fn set_enabled(&mut self, enabled: bool, cx: &mut Context<Self>) {
if self.enabled == enabled {
return;
}
self.enabled = enabled;
self.on_enabled_changed(cx);
}
fn on_enabled_changed(&self, cx: &mut Context<Self>) {
if !self.enabled {
CommandPaletteFilter::update_global(cx, |filter, _cx| {
filter.hide_namespace(Self::NAMESPACE);
});
return;
}
CommandPaletteFilter::update_global(cx, |filter, _cx| {
filter.show_namespace(Self::NAMESPACE);
});
cx.notify();
}
pub fn refresh_python_kernelspecs(
&mut self,
worktree_id: WorktreeId,
project: &Entity<Project>,
cx: &mut Context<Self>,
) -> Task<Result<()>> {
if !self.fetching_python_kernelspecs.insert(worktree_id) {
return Task::ready(Ok(()));
}
let is_remote = project.read(cx).is_remote();
// WSL does require access to global kernel specs, so we only exclude remote worktrees that aren't WSL.
// TODO: a better way to handle WSL vs SSH/remote projects,
let is_wsl_remote = project
.read(cx)
.remote_connection_options(cx)
.map_or(false, |opts| {
matches!(opts, RemoteConnectionOptions::Wsl(_))
});
let kernel_specifications_task = python_env_kernel_specifications(project, worktree_id, cx);
let active_toolchain = project.read(cx).active_toolchain(
ProjectPath {
worktree_id,
path: RelPath::empty().into(),
},
LanguageName::new_static("Python"),
cx,
);
cx.spawn(async move |this, cx| {
let kernel_specifications_res = kernel_specifications_task.await;
this.update(cx, |this, _cx| {
this.fetching_python_kernelspecs.remove(&worktree_id);
})
.ok();
let kernel_specifications =
kernel_specifications_res.context("getting python kernelspecs")?;
let active_toolchain_path = active_toolchain.await.map(|toolchain| toolchain.path);
this.update(cx, |this, cx| {
this.kernel_specifications_for_worktree
.insert(worktree_id, kernel_specifications);
if let Some(path) = active_toolchain_path {
this.active_python_toolchain_for_worktree
.insert(worktree_id, path);
}
if is_remote && !is_wsl_remote {
this.remote_worktrees.insert(worktree_id);
} else {
this.remote_worktrees.remove(&worktree_id);
}
cx.notify();
})
})
}
fn get_remote_kernel_specifications(
&self,
cx: &mut Context<Self>,
) -> Option<Task<Result<Vec<KernelSpecification>>>> {
match (
std::env::var("JUPYTER_SERVER"),
std::env::var("JUPYTER_TOKEN"),
) {
(Ok(server), Ok(token)) => {
let remote_server = RemoteServer {
base_url: server,
token,
};
let http_client = cx.http_client();
Some(cx.spawn(async move |_, _| {
list_remote_kernelspecs(remote_server, http_client)
.await
.map(|specs| {
specs
.into_iter()
.map(KernelSpecification::JupyterServer)
.collect()
})
}))
}
_ => None,
}
}
pub fn ensure_kernelspecs(&mut self, cx: &mut Context<Self>) {
if self.kernelspecs_initialized {
return;
}
self.kernelspecs_initialized = true;
self.refresh_kernelspecs(cx).detach_and_log_err(cx);
}
pub fn refresh_kernelspecs(&mut self, cx: &mut Context<Self>) -> Task<Result<()>> {
let local_kernel_specifications = local_kernel_specifications(self.fs.clone());
let wsl_kernel_specifications = wsl_kernel_specifications(cx.background_executor().clone());
let remote_kernel_specifications = self.get_remote_kernel_specifications(cx);
let all_specs = cx.background_spawn(async move {
let mut all_specs = local_kernel_specifications
.await?
.into_iter()
.map(KernelSpecification::Jupyter)
.collect::<Vec<_>>();
if let Ok(wsl_specs) = wsl_kernel_specifications.await {
all_specs.extend(wsl_specs);
}
if let Some(remote_task) = remote_kernel_specifications
&& let Ok(remote_specs) = remote_task.await
{
all_specs.extend(remote_specs);
}
anyhow::Ok(all_specs)
});
cx.spawn(async move |this, cx| {
let all_specs = all_specs.await;
if let Ok(specs) = all_specs {
this.update(cx, |this, cx| {
this.kernel_specifications = specs;
cx.notify();
})
.ok();
}
anyhow::Ok(())
})
}
pub fn set_active_kernelspec(
&mut self,
worktree_id: WorktreeId,
kernelspec: KernelSpecification,
_cx: &mut Context<Self>,
) {
self.selected_kernel_for_worktree
.insert(worktree_id, kernelspec);
}
pub fn active_python_toolchain_path(&self, worktree_id: WorktreeId) -> Option<&SharedString> {
self.active_python_toolchain_for_worktree.get(&worktree_id)
}
pub fn selected_kernel(&self, worktree_id: WorktreeId) -> Option<&KernelSpecification> {
self.selected_kernel_for_worktree.get(&worktree_id)
}
pub fn is_recommended_kernel(
&self,
worktree_id: WorktreeId,
spec: &KernelSpecification,
) -> bool {
if let Some(active_path) = self.active_python_toolchain_path(worktree_id) {
spec.path().as_ref() == active_path.as_ref()
} else {
false
}
}
pub fn active_kernelspec(
&self,
worktree_id: WorktreeId,
language_at_cursor: Option<Arc<Language>>,
cx: &App,
) -> Option<KernelSpecification> {
if let Some(selected) = self.selected_kernel_for_worktree.get(&worktree_id).cloned() {
return Some(selected);
}
let language_at_cursor = language_at_cursor?;
// Prefer the recommended (active toolchain) kernel if it has ipykernel
if let Some(active_path) = self.active_python_toolchain_path(worktree_id) {
let recommended = self
.kernel_specifications_for_worktree(worktree_id)
.find(|spec| {
spec.has_ipykernel()
&& language_at_cursor.matches_kernel_language(spec.language().as_ref())
&& spec.path().as_ref() == active_path.as_ref()
})
.cloned();
if recommended.is_some() {
return recommended;
}
}
// Then try the first PythonEnv with ipykernel matching the language
let python_env = self
.kernel_specifications_for_worktree(worktree_id)
.find(|spec| {
matches!(spec, KernelSpecification::PythonEnv(_))
&& spec.has_ipykernel()
&& language_at_cursor.matches_kernel_language(spec.language().as_ref())
})
.cloned();
if python_env.is_some() {
return python_env;
}
// Fall back to legacy name-based and language-based matching
self.kernelspec_legacy_by_lang_only(worktree_id, language_at_cursor, cx)
}
fn kernelspec_legacy_by_lang_only(
&self,
worktree_id: WorktreeId,
language_at_cursor: Arc<Language>,
cx: &App,
) -> Option<KernelSpecification> {
let settings = JupyterSettings::get_global(cx);
let selected_kernel = settings
.kernel_selections
.get(language_at_cursor.code_fence_block_name().as_ref());
let found_by_name = self
.kernel_specifications_for_worktree(worktree_id)
.find(|runtime_specification| {
if let (Some(selected), KernelSpecification::Jupyter(runtime_specification)) =
(selected_kernel, runtime_specification)
{
return runtime_specification.name.to_lowercase() == selected.to_lowercase();
}
false
})
.cloned();
if let Some(found_by_name) = found_by_name {
return Some(found_by_name);
}
self.kernel_specifications_for_worktree(worktree_id)
.find(|spec| {
spec.has_ipykernel()
&& language_at_cursor.matches_kernel_language(spec.language().as_ref())
})
.cloned()
}
pub fn get_session(&self, entity_id: EntityId) -> Option<&Entity<Session>> {
self.sessions.get(&entity_id)
}
pub fn insert_session(&mut self, entity_id: EntityId, session: Entity<Session>) {
self.sessions.insert(entity_id, session);
}
pub fn remove_session(&mut self, entity_id: EntityId) {
self.sessions.remove(&entity_id);
}
fn shutdown_all_sessions(
&mut self,
cx: &mut Context<Self>,
) -> impl Future<Output = ()> + use<> {
for session in self.sessions.values() {
session.update(cx, |session, _cx| {
if let Kernel::RunningKernel(mut kernel) =
std::mem::replace(&mut session.kernel, Kernel::Shutdown)
{
kernel.kill();
}
});
}
self.sessions.clear();
futures::future::ready(())
}
#[cfg(test)]
pub fn set_kernel_specs_for_testing(
&mut self,
specs: Vec<KernelSpecification>,
cx: &mut Context<Self>,
) {
self.kernel_specifications = specs;
cx.notify();
}
}