Discover available python environments with Jupyter kernel support (#20467)

![image](https://github.com/user-attachments/assets/7c042bc9-88be-4d7b-b63d-e5b555d54b18)

Closes #18291
Closes #16757
Closes #15563

Release Notes:

- Added support for kernelspecs based on python environments

crates/quick_action_bar/src/repl_menu.rs

@@ -4,8 +4,8 @@ use gpui::{percentage, Animation, AnimationExt, AnyElement, Transformation, View
 use picker::Picker;
 use repl::{
     components::{KernelPickerDelegate, KernelSelector},
-    ExecutionState, JupyterSettings, Kernel, KernelSpecification, KernelStatus, Session,
+    worktree_id_for_editor, ExecutionState, JupyterSettings, Kernel, KernelSpecification,
-    SessionSupport,
+    KernelStatus, Session, SessionSupport,
 };
 use ui::{
     prelude::*, ButtonLike, ContextMenu, IconWithIndicator, Indicator, IntoElement, PopoverMenu,
@@ -30,9 +30,6 @@ struct ReplMenuState {
     status: KernelStatus,
     kernel_name: SharedString,
     kernel_language: SharedString,
-    // TODO: Persist rotation state so the
-    // icon doesn't reset on every state change
-    // current_delta: Duration,
 }
 
 impl QuickActionBar {
@@ -178,12 +175,6 @@ impl QuickActionBar {
                         },
                     )
                     .separator()
-                    .link(
-                        "Change Kernel",
-                        Box::new(zed_actions::OpenBrowser {
-                            url: format!("{}#change-kernel", ZED_REPL_DOCUMENTATION),
-                        }),
-                    )
                     .custom_entry(
                         move |_cx| {
                             Label::new("Shut Down Kernel")
@@ -290,7 +281,10 @@ impl QuickActionBar {
         let editor = if let Some(editor) = self.active_editor() {
             editor
         } else {
-            // todo!()
+            return div().into_any_element();
+        };
+
+        let Some(worktree_id) = worktree_id_for_editor(editor.downgrade(), cx) else {
             return div().into_any_element();
         };
 
@@ -313,7 +307,7 @@ impl QuickActionBar {
                     repl::assign_kernelspec(kernelspec, editor.downgrade(), cx).ok();
                 })
             },
-            current_kernelspec.clone(),
+            worktree_id,
             ButtonLike::new("kernel-selector")
                 .style(ButtonStyle::Subtle)
                 .child(

crates/repl/src/components/kernel_options.rs

@@ -4,8 +4,10 @@ use crate::KERNEL_DOCS_URL;
 
 use gpui::DismissEvent;
 
+use gpui::FontWeight;
 use picker::Picker;
 use picker::PickerDelegate;
+use project::WorktreeId;
 
 use std::sync::Arc;
 use ui::ListItemSpacing;
@@ -22,7 +24,7 @@ pub struct KernelSelector<T: PopoverTrigger> {
     on_select: OnSelect,
     trigger: T,
     info_text: Option<SharedString>,
-    current_kernelspec: Option<KernelSpecification>,
+    worktree_id: WorktreeId,
 }
 
 pub struct KernelPickerDelegate {
@@ -33,17 +35,13 @@ pub struct KernelPickerDelegate {
 }
 
 impl<T: PopoverTrigger> KernelSelector<T> {
-    pub fn new(
+    pub fn new(on_select: OnSelect, worktree_id: WorktreeId, trigger: T) -> Self {
-        on_select: OnSelect,
-        current_kernelspec: Option<KernelSpecification>,
-        trigger: T,
-    ) -> Self {
         KernelSelector {
             on_select,
             handle: None,
             trigger,
             info_text: None,
-            current_kernelspec,
+            worktree_id,
         }
     }
 
@@ -130,24 +128,34 @@ impl PickerDelegate for KernelPickerDelegate {
                 .spacing(ListItemSpacing::Sparse)
                 .selected(selected)
                 .child(
-                    h_flex().w_full().justify_between().min_w(px(200.)).child(
+                    v_flex()
-                        h_flex()
+                        .min_w(px(600.))
-                            .gap_1p5()
+                        .w_full()
-                            .child(Label::new(kernelspec.name()))
+                        .gap_0p5()
-                            .child(
+                        .child(
-                                Label::new(kernelspec.type_name())
+                            h_flex()
-                                    .size(LabelSize::XSmall)
+                                .w_full()
-                                    .color(Color::Muted),
+                                .gap_1()
-                            ),
+                                .child(Label::new(kernelspec.name()).weight(FontWeight::MEDIUM))
-                    ),
+                                .child(
+                                    Label::new(kernelspec.language())
+                                        .size(LabelSize::Small)
+                                        .color(Color::Muted),
+                                ),
+                        )
+                        .child(
+                            Label::new(kernelspec.path())
+                                .size(LabelSize::XSmall)
+                                .color(Color::Muted),
+                        ),
                 )
-                .end_slot(div().when(is_selected, |this| {
+                .when(is_selected, |item| {
-                    this.child(
+                    item.end_slot(
                         Icon::new(IconName::Check)
                             .color(Color::Accent)
                             .size(IconSize::Small),
                     )
-                })),
+                }),
         )
     }
 
@@ -175,10 +183,13 @@ impl PickerDelegate for KernelPickerDelegate {
 impl<T: PopoverTrigger> RenderOnce for KernelSelector<T> {
     fn render(self, cx: &mut WindowContext) -> impl IntoElement {
         let store = ReplStore::global(cx).read(cx);
-        let all_kernels: Vec<KernelSpecification> =
-            store.kernel_specifications().cloned().collect();
 
-        let selected_kernelspec = self.current_kernelspec;
+        let all_kernels: Vec<KernelSpecification> = store
+            .kernel_specifications_for_worktree(self.worktree_id)
+            .cloned()
+            .collect();
+
+        let selected_kernelspec = store.active_kernelspec(self.worktree_id, None, cx);
 
         let delegate = KernelPickerDelegate {
             on_select: self.on_select,

crates/repl/src/kernels.rs

@@ -5,8 +5,9 @@ use futures::{
     stream::{self, SelectAll, StreamExt},
     SinkExt as _,
 };
-use gpui::{AppContext, EntityId, Task};
+use gpui::{AppContext, EntityId, Model, Task};
-use project::Fs;
+use language::LanguageName;
+use project::{Fs, Project, WorktreeId};
 use runtimelib::{
     dirs, ConnectionInfo, ExecutionState, JupyterKernelspec, JupyterMessage, JupyterMessageContent,
     KernelInfoReply,
@@ -15,6 +16,7 @@ use smol::{net::TcpListener, process::Command};
 use std::{
     env,
     fmt::Debug,
+    future::Future,
     net::{IpAddr, Ipv4Addr, SocketAddr},
     path::PathBuf,
     sync::Arc,
@@ -465,6 +467,72 @@ async fn read_kernels_dir(path: PathBuf, fs: &dyn Fs) -> Result<Vec<LocalKernelS
     Ok(valid_kernelspecs)
 }
 
+pub fn python_env_kernel_specifications(
+    project: &Model<Project>,
+    worktree_id: WorktreeId,
+    cx: &mut AppContext,
+) -> impl Future<Output = Result<Vec<KernelSpecification>>> {
+    let python_language = LanguageName::new("Python");
+    let toolchains = project
+        .read(cx)
+        .available_toolchains(worktree_id, python_language, cx);
+    let background_executor = cx.background_executor().clone();
+
+    async move {
+        let toolchains = if let Some(toolchains) = toolchains.await {
+            toolchains
+        } else {
+            return Ok(Vec::new());
+        };
+
+        let kernelspecs = toolchains.toolchains.into_iter().map(|toolchain| {
+            background_executor.spawn(async move {
+                let python_path = toolchain.path.to_string();
+
+                // Check if ipykernel is installed
+                let ipykernel_check = Command::new(&python_path)
+                    .args(&["-c", "import ipykernel"])
+                    .output()
+                    .await;
+
+                if ipykernel_check.is_ok() && ipykernel_check.unwrap().status.success() {
+                    // Create a default kernelspec for this environment
+                    let default_kernelspec = JupyterKernelspec {
+                        argv: vec![
+                            python_path.clone(),
+                            "-m".to_string(),
+                            "ipykernel_launcher".to_string(),
+                            "-f".to_string(),
+                            "{connection_file}".to_string(),
+                        ],
+                        display_name: toolchain.name.to_string(),
+                        language: "python".to_string(),
+                        interrupt_mode: None,
+                        metadata: None,
+                        env: None,
+                    };
+
+                    Some(KernelSpecification::PythonEnv(LocalKernelSpecification {
+                        name: toolchain.name.to_string(),
+                        path: PathBuf::from(&python_path),
+                        kernelspec: default_kernelspec,
+                    }))
+                } else {
+                    None
+                }
+            })
+        });
+
+        let kernel_specs = futures::future::join_all(kernelspecs)
+            .await
+            .into_iter()
+            .flatten()
+            .collect();
+
+        anyhow::Ok(kernel_specs)
+    }
+}
+
 pub async fn local_kernel_specifications(fs: Arc<dyn Fs>) -> Result<Vec<LocalKernelSpecification>> {
     let mut data_dirs = dirs::data_dirs();
 

crates/repl/src/repl_editor.rs

@@ -7,6 +7,7 @@ use anyhow::{Context, Result};
 use editor::Editor;
 use gpui::{prelude::*, Entity, View, WeakView, WindowContext};
 use language::{BufferSnapshot, Language, LanguageName, Point};
+use project::{Item as _, WorktreeId};
 
 use crate::repl_store::ReplStore;
 use crate::session::SessionEvent;
@@ -24,6 +25,13 @@ pub fn assign_kernelspec(
         return Ok(());
     }
 
+    let worktree_id = crate::repl_editor::worktree_id_for_editor(weak_editor.clone(), cx)
+        .context("editor is not in a worktree")?;
+
+    store.update(cx, |store, cx| {
+        store.set_active_kernelspec(worktree_id, kernel_specification.clone(), cx);
+    });
+
     let fs = store.read(cx).fs().clone();
     let telemetry = store.read(cx).telemetry().clone();
 
@@ -79,6 +87,10 @@ pub fn run(editor: WeakView<Editor>, move_down: bool, cx: &mut WindowContext) ->
         return Ok(());
     };
 
+    let Some(project_path) = buffer.read(cx).project_path(cx) else {
+        return Ok(());
+    };
+
     let (runnable_ranges, next_cell_point) =
         runnable_ranges(&buffer.read(cx).snapshot(), selected_range);
 
@@ -87,11 +99,10 @@ pub fn run(editor: WeakView<Editor>, move_down: bool, cx: &mut WindowContext) ->
             continue;
         };
 
-        let kernel_specification = store.update(cx, |store, cx| {
+        let kernel_specification = store
-            store
+            .read(cx)
-                .kernelspec(language.code_fence_block_name().as_ref(), cx)
+            .active_kernelspec(project_path.worktree_id, Some(language.clone()), cx)
-                .with_context(|| format!("No kernel found for language: {}", language.name()))
+            .ok_or_else(|| anyhow::anyhow!("No kernel found for language: {}", language.name()))?;
-        })?;
 
         let fs = store.read(cx).fs().clone();
         let telemetry = store.read(cx).telemetry().clone();
@@ -156,6 +167,22 @@ pub enum SessionSupport {
     Unsupported,
 }
 
+pub fn worktree_id_for_editor(
+    editor: WeakView<Editor>,
+    cx: &mut WindowContext,
+) -> Option<WorktreeId> {
+    editor.upgrade().and_then(|editor| {
+        editor
+            .read(cx)
+            .buffer()
+            .read(cx)
+            .as_singleton()?
+            .read(cx)
+            .project_path(cx)
+            .map(|path| path.worktree_id)
+    })
+}
+
 pub fn session(editor: WeakView<Editor>, cx: &mut WindowContext) -> SessionSupport {
     let store = ReplStore::global(cx);
     let entity_id = editor.entity_id();
@@ -164,17 +191,24 @@ pub fn session(editor: WeakView<Editor>, cx: &mut WindowContext) -> SessionSuppo
         return SessionSupport::ActiveSession(session);
     };
 
-    let Some(language) = get_language(editor, cx) else {
+    let Some(language) = get_language(editor.clone(), cx) else {
         return SessionSupport::Unsupported;
     };
-    let kernelspec = store.update(cx, |store, cx| {
+
-        store.kernelspec(language.code_fence_block_name().as_ref(), cx)
+    let worktree_id = worktree_id_for_editor(editor.clone(), cx);
-    });
+
+    let Some(worktree_id) = worktree_id else {
+        return SessionSupport::Unsupported;
+    };
+
+    let kernelspec = store
+        .read(cx)
+        .active_kernelspec(worktree_id, Some(language.clone()), cx);
 
     match kernelspec {
         Some(kernelspec) => SessionSupport::Inactive(kernelspec),
         None => {
-            if language_supported(&language) {
+            if language_supported(&language.clone()) {
                 SessionSupport::RequiresSetup(language.name())
             } else {
                 SessionSupport::Unsupported

crates/repl/src/repl_sessions_ui.rs

@@ -1,10 +1,10 @@
 use editor::Editor;
 use gpui::{
     actions, prelude::*, AnyElement, AppContext, EventEmitter, FocusHandle, FocusableView,
-    FontWeight, Subscription, View,
+    Subscription, View,
 };
-use std::collections::HashMap;
+use project::Item as _;
-use ui::{prelude::*, ButtonLike, ElevationIndex, KeyBinding, ListItem, Tooltip};
+use ui::{prelude::*, ButtonLike, ElevationIndex, KeyBinding};
 use util::ResultExt as _;
 use workspace::item::ItemEvent;
 use workspace::WorkspaceId;
@@ -12,7 +12,6 @@ use workspace::{item::Item, Workspace};
 
 use crate::jupyter_settings::JupyterSettings;
 use crate::repl_store::ReplStore;
-use crate::{KernelSpecification, KERNEL_DOCS_URL};
 
 actions!(
     repl,
@@ -63,17 +62,34 @@ pub fn init(cx: &mut AppContext) {
 
         cx.defer(|editor, cx| {
             let workspace = Workspace::for_window(cx);
+            let project = workspace.map(|workspace| workspace.read(cx).project().clone());
 
-            let is_local_project = workspace
+            let is_local_project = project
-                .map(|workspace| workspace.read(cx).project().read(cx).is_local())
+                .as_ref()
+                .map(|project| project.read(cx).is_local())
                 .unwrap_or(false);
 
             if !is_local_project {
                 return;
             }
 
+            let project_path = editor
+                .buffer()
+                .read(cx)
+                .as_singleton()
+                .and_then(|buffer| buffer.read(cx).project_path(cx));
+
             let editor_handle = cx.view().downgrade();
 
+            if let (Some(project_path), Some(project)) = (project_path, project) {
+                let store = ReplStore::global(cx);
+                store.update(cx, |store, cx| {
+                    store
+                        .refresh_python_kernelspecs(project_path.worktree_id, &project, cx)
+                        .detach_and_log_err(cx);
+                });
+            }
+
             editor
                 .register_action({
                     let editor_handle = editor_handle.clone();
@@ -169,7 +185,10 @@ impl Render for ReplSessionsPage {
 
         let (kernel_specifications, sessions) = store.update(cx, |store, _cx| {
             (
-                store.kernel_specifications().cloned().collect::<Vec<_>>(),
+                store
+                    .pure_jupyter_kernel_specifications()
+                    .cloned()
+                    .collect::<Vec<_>>(),
                 store.sessions().cloned().collect::<Vec<_>>(),
             )
         });
@@ -198,97 +217,18 @@ impl Render for ReplSessionsPage {
                 );
         }
 
-        let mut kernels_by_language: HashMap<SharedString, Vec<&KernelSpecification>> =
-            kernel_specifications
-                .iter()
-                .map(|spec| (spec.language(), spec))
-                .fold(HashMap::new(), |mut acc, (language, spec)| {
-                    acc.entry(language).or_default().push(spec);
-                    acc
-                });
-
-        for kernels in kernels_by_language.values_mut() {
-            kernels.sort_by_key(|a| a.name())
-        }
-
-        // Convert to a sorted Vec of tuples
-        let mut sorted_kernels: Vec<(SharedString, Vec<&KernelSpecification>)> =
-            kernels_by_language.into_iter().collect();
-        sorted_kernels.sort_by(|a, b| a.0.cmp(&b.0));
-
-        let kernels_available = v_flex()
-            .child(Label::new("Kernels available").size(LabelSize::Large))
-            .gap_2()
-            .child(
-                h_flex()
-                    .child(Label::new(
-                        "Defaults indicated with a checkmark. Learn how to change your default kernel in the ",
-                    ))
-                    .child(
-                        ButtonLike::new("configure-kernels")
-                            .style(ButtonStyle::Filled)
-                            // .size(ButtonSize::Compact)
-                            .layer(ElevationIndex::Surface)
-                            .child(Label::new("REPL documentation"))
-                            .child(Icon::new(IconName::Link))
-                            .on_click(move |_, cx| {
-                                cx.open_url(KERNEL_DOCS_URL)
-                            }),
-                    ),
-            )
-            .children(sorted_kernels.into_iter().map(|(language, specs)| {
-                let chosen_kernel = store.read(cx).kernelspec(&language, cx);
-
-                v_flex()
-                    .gap_1()
-                    .child(Label::new(language.clone()).weight(FontWeight::BOLD))
-                    .children(specs.into_iter().map(|spec| {
-                        let is_choice = if let Some(chosen_kernel) = &chosen_kernel {
-                            chosen_kernel == spec
-                        } else {
-                            false
-                        };
-
-                        let path = spec.path();
-
-                        ListItem::new(path.clone())
-                            .selectable(false)
-                            .tooltip({
-                                let path = path.clone();
-                                move |cx| Tooltip::text(path.clone(), cx)})
-                            .child(
-                                h_flex()
-                                    .gap_1()
-                                    .child(div().id(path.clone()).child(Label::new(spec.name())))
-                                    .when(is_choice, |el| {
-
-                                        let language = language.clone();
-
-                                        el.child(
-
-                                        div().id("check").tooltip(move |cx| Tooltip::text(format!("Default Kernel for {language}"), cx))
-                                            .child(Icon::new(IconName::Check)))}),
-                            )
-
-                    }))
-            }));
-
         // When there are no sessions, show the command to run code in an editor
         if sessions.is_empty() {
             let instructions = "To run code in a Jupyter kernel, select some code and use the 'repl::Run' command.";
 
-            return ReplSessionsContainer::new("No Jupyter Kernel Sessions")
+            return ReplSessionsContainer::new("No Jupyter Kernel Sessions").child(
-                .child(
+                v_flex()
-                    v_flex()
+                    .child(Label::new(instructions))
-                        .child(Label::new(instructions))
+                    .children(KeyBinding::for_action(&Run, cx)),
-                        .children(KeyBinding::for_action(&Run, cx)),
+            );
-                )
-                .child(div().pt_3().child(kernels_available));
         }
 
-        ReplSessionsContainer::new("Jupyter Kernel Sessions")
+        ReplSessionsContainer::new("Jupyter Kernel Sessions").children(sessions)
-            .children(sessions)
-            .child(kernels_available)
     }
 }
 

crates/repl/src/repl_store.rs

@@ -7,10 +7,11 @@ use command_palette_hooks::CommandPaletteFilter;
 use gpui::{
     prelude::*, AppContext, EntityId, Global, Model, ModelContext, Subscription, Task, View,
 };
-use project::Fs;
+use language::Language;
+use project::{Fs, Project, WorktreeId};
 use settings::{Settings, SettingsStore};
 
-use crate::kernels::local_kernel_specifications;
+use crate::kernels::{local_kernel_specifications, python_env_kernel_specifications};
 use crate::{JupyterSettings, KernelSpecification, Session};
 
 struct GlobalReplStore(Model<ReplStore>);
@@ -22,6 +23,8 @@ pub struct ReplStore {
     enabled: bool,
     sessions: HashMap<EntityId, View<Session>>,
     kernel_specifications: Vec<KernelSpecification>,
+    selected_kernel_for_worktree: HashMap<WorktreeId, KernelSpecification>,
+    kernel_specifications_for_worktree: HashMap<WorktreeId, Vec<KernelSpecification>>,
     telemetry: Arc<Telemetry>,
     _subscriptions: Vec<Subscription>,
 }
@@ -55,6 +58,8 @@ impl ReplStore {
             sessions: HashMap::default(),
             kernel_specifications: Vec::new(),
             _subscriptions: subscriptions,
+            kernel_specifications_for_worktree: HashMap::default(),
+            selected_kernel_for_worktree: HashMap::default(),
         };
         this.on_enabled_changed(cx);
         this
@@ -72,7 +77,18 @@ impl ReplStore {
         self.enabled
     }
 
-    pub fn kernel_specifications(&self) -> impl Iterator<Item = &KernelSpecification> {
+    pub fn kernel_specifications_for_worktree(
+        &self,
+        worktree_id: WorktreeId,
+    ) -> impl Iterator<Item = &KernelSpecification> {
+        self.kernel_specifications_for_worktree
+            .get(&worktree_id)
+            .into_iter()
+            .flat_map(|specs| specs.iter())
+            .chain(self.kernel_specifications.iter())
+    }
+
+    pub fn pure_jupyter_kernel_specifications(&self) -> impl Iterator<Item = &KernelSpecification> {
         self.kernel_specifications.iter()
     }
 
@@ -105,8 +121,29 @@ impl ReplStore {
         cx.notify();
     }
 
+    pub fn refresh_python_kernelspecs(
+        &mut self,
+        worktree_id: WorktreeId,
+        project: &Model<Project>,
+        cx: &mut ModelContext<Self>,
+    ) -> Task<Result<()>> {
+        let kernel_specifications = python_env_kernel_specifications(project, worktree_id, cx);
+        cx.spawn(move |this, mut cx| async move {
+            let kernel_specifications = kernel_specifications
+                .await
+                .map_err(|e| anyhow::anyhow!("Failed to get python kernelspecs: {:?}", e))?;
+
+            this.update(&mut cx, |this, cx| {
+                this.kernel_specifications_for_worktree
+                    .insert(worktree_id, kernel_specifications);
+                cx.notify();
+            })
+        })
+    }
+
     pub fn refresh_kernelspecs(&mut self, cx: &mut ModelContext<Self>) -> Task<Result<()>> {
         let local_kernel_specifications = local_kernel_specifications(self.fs.clone());
+
         cx.spawn(|this, mut cx| async move {
             let local_kernel_specifications = local_kernel_specifications.await?;
 
@@ -122,9 +159,41 @@ impl ReplStore {
         })
     }
 
-    pub fn kernelspec(&self, language_name: &str, cx: &AppContext) -> Option<KernelSpecification> {
+    pub fn set_active_kernelspec(
+        &mut self,
+        worktree_id: WorktreeId,
+        kernelspec: KernelSpecification,
+        _cx: &mut ModelContext<Self>,
+    ) {
+        self.selected_kernel_for_worktree
+            .insert(worktree_id, kernelspec);
+    }
+
+    pub fn active_kernelspec(
+        &self,
+        worktree_id: WorktreeId,
+        language_at_cursor: Option<Arc<Language>>,
+        cx: &AppContext,
+    ) -> Option<KernelSpecification> {
+        let selected_kernelspec = self.selected_kernel_for_worktree.get(&worktree_id).cloned();
+
+        if let Some(language_at_cursor) = language_at_cursor {
+            selected_kernelspec
+                .or_else(|| self.kernelspec_legacy_by_lang_only(language_at_cursor, cx))
+        } else {
+            selected_kernelspec
+        }
+    }
+
+    fn kernelspec_legacy_by_lang_only(
+        &self,
+        language_at_cursor: Arc<Language>,
+        cx: &AppContext,
+    ) -> Option<KernelSpecification> {
         let settings = JupyterSettings::get_global(cx);
-        let selected_kernel = settings.kernel_selections.get(language_name);
+        let selected_kernel = settings
+            .kernel_selections
+            .get(language_at_cursor.code_fence_block_name().as_ref());
 
         let found_by_name = self
             .kernel_specifications
@@ -149,10 +218,15 @@ impl ReplStore {
             .find(|kernel_option| match kernel_option {
                 KernelSpecification::Jupyter(runtime_specification) => {
                     runtime_specification.kernelspec.language.to_lowercase()
-                        == language_name.to_lowercase()
+                        == language_at_cursor.code_fence_block_name().to_lowercase()
+                }
+                KernelSpecification::PythonEnv(runtime_specification) => {
+                    runtime_specification.kernelspec.language.to_lowercase()
+                        == language_at_cursor.code_fence_block_name().to_lowercase()
+                }
+                KernelSpecification::Remote(_) => {
+                    unimplemented!()
                 }
-                // todo!()
-                _ => false,
             })
             .cloned()
     }