1mod native_kernel;
2use std::{fmt::Debug, future::Future, path::PathBuf};
3
4use futures::{channel::mpsc, future::Shared};
5use gpui::{App, Entity, Task, Window};
6use language::LanguageName;
7pub use native_kernel::*;
8
9mod remote_kernels;
10use project::{Project, ProjectPath, Toolchains, WorktreeId};
11pub use remote_kernels::*;
12
13use anyhow::Result;
14use gpui::Context;
15use jupyter_protocol::JupyterKernelspec;
16use runtimelib::{ExecutionState, JupyterMessage, KernelInfoReply};
17use ui::{Icon, IconName, SharedString};
18use util::rel_path::RelPath;
19
20pub trait KernelSession: Sized {
21 fn route(&mut self, message: &JupyterMessage, window: &mut Window, cx: &mut Context<Self>);
22 fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>);
23}
24
25#[derive(Debug, Clone)]
26pub struct PythonEnvKernelSpecification {
27 pub name: String,
28 pub path: PathBuf,
29 pub kernelspec: JupyterKernelspec,
30 pub has_ipykernel: bool,
31 /// Display label for the environment type: "venv", "Conda", "Pyenv", etc.
32 pub environment_kind: Option<String>,
33}
34
35impl PartialEq for PythonEnvKernelSpecification {
36 fn eq(&self, other: &Self) -> bool {
37 self.name == other.name && self.path == other.path
38 }
39}
40
41impl Eq for PythonEnvKernelSpecification {}
42
43impl PythonEnvKernelSpecification {
44 pub fn as_local_spec(&self) -> LocalKernelSpecification {
45 LocalKernelSpecification {
46 name: self.name.clone(),
47 path: self.path.clone(),
48 kernelspec: self.kernelspec.clone(),
49 }
50 }
51}
52
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub enum KernelSpecification {
55 Remote(RemoteKernelSpecification),
56 Jupyter(LocalKernelSpecification),
57 PythonEnv(PythonEnvKernelSpecification),
58}
59
60impl KernelSpecification {
61 pub fn name(&self) -> SharedString {
62 match self {
63 Self::Jupyter(spec) => spec.name.clone().into(),
64 Self::PythonEnv(spec) => spec.name.clone().into(),
65 Self::Remote(spec) => spec.name.clone().into(),
66 }
67 }
68
69 pub fn type_name(&self) -> SharedString {
70 match self {
71 Self::Jupyter(_) => "Jupyter".into(),
72 Self::PythonEnv(spec) => SharedString::from(
73 spec.environment_kind
74 .clone()
75 .unwrap_or_else(|| "Python Environment".to_string()),
76 ),
77 Self::Remote(_) => "Remote".into(),
78 }
79 }
80
81 pub fn path(&self) -> SharedString {
82 SharedString::from(match self {
83 Self::Jupyter(spec) => spec.path.to_string_lossy().into_owned(),
84 Self::PythonEnv(spec) => spec.path.to_string_lossy().into_owned(),
85 Self::Remote(spec) => spec.url.to_string(),
86 })
87 }
88
89 pub fn language(&self) -> SharedString {
90 SharedString::from(match self {
91 Self::Jupyter(spec) => spec.kernelspec.language.clone(),
92 Self::PythonEnv(spec) => spec.kernelspec.language.clone(),
93 Self::Remote(spec) => spec.kernelspec.language.clone(),
94 })
95 }
96
97 pub fn has_ipykernel(&self) -> bool {
98 match self {
99 Self::Jupyter(_) | Self::Remote(_) => true,
100 Self::PythonEnv(spec) => spec.has_ipykernel,
101 }
102 }
103
104 pub fn environment_kind_label(&self) -> Option<SharedString> {
105 match self {
106 Self::PythonEnv(spec) => spec
107 .environment_kind
108 .as_ref()
109 .map(|kind| SharedString::from(kind.clone())),
110 Self::Jupyter(_) => Some("Jupyter".into()),
111 Self::Remote(_) => Some("Remote".into()),
112 }
113 }
114
115 pub fn icon(&self, cx: &App) -> Icon {
116 let lang_name = match self {
117 Self::Jupyter(spec) => spec.kernelspec.language.clone(),
118 Self::PythonEnv(spec) => spec.kernelspec.language.clone(),
119 Self::Remote(spec) => spec.kernelspec.language.clone(),
120 };
121
122 file_icons::FileIcons::get(cx)
123 .get_icon_for_type(&lang_name.to_lowercase(), cx)
124 .map(Icon::from_path)
125 .unwrap_or(Icon::new(IconName::ReplNeutral))
126 }
127}
128
129fn extract_environment_kind(toolchain_json: &serde_json::Value) -> Option<String> {
130 let kind_str = toolchain_json.get("kind")?.as_str()?;
131 let label = match kind_str {
132 "Conda" => "Conda",
133 "Pixi" => "pixi",
134 "Homebrew" => "Homebrew",
135 "Pyenv" => "global (Pyenv)",
136 "GlobalPaths" => "global",
137 "PyenvVirtualEnv" => "Pyenv",
138 "Pipenv" => "Pipenv",
139 "Poetry" => "Poetry",
140 "MacPythonOrg" => "global (Python.org)",
141 "MacCommandLineTools" => "global (Command Line Tools for Xcode)",
142 "LinuxGlobal" => "global",
143 "MacXCode" => "global (Xcode)",
144 "Venv" => "venv",
145 "VirtualEnv" => "virtualenv",
146 "VirtualEnvWrapper" => "virtualenvwrapper",
147 "WindowsStore" => "global (Windows Store)",
148 "WindowsRegistry" => "global (Windows Registry)",
149 "Uv" => "uv",
150 "UvWorkspace" => "uv (Workspace)",
151 _ => kind_str,
152 };
153 Some(label.to_string())
154}
155
156pub fn python_env_kernel_specifications(
157 project: &Entity<Project>,
158 worktree_id: WorktreeId,
159 cx: &mut App,
160) -> impl Future<Output = Result<Vec<KernelSpecification>>> + use<> {
161 let python_language = LanguageName::new_static("Python");
162 let toolchains = project.read(cx).available_toolchains(
163 ProjectPath {
164 worktree_id,
165 path: RelPath::empty().into(),
166 },
167 python_language,
168 cx,
169 );
170 let background_executor = cx.background_executor().clone();
171
172 async move {
173 let (toolchains, user_toolchains) = if let Some(Toolchains {
174 toolchains,
175 root_path: _,
176 user_toolchains,
177 }) = toolchains.await
178 {
179 (toolchains, user_toolchains)
180 } else {
181 return Ok(Vec::new());
182 };
183
184 let kernelspecs = user_toolchains
185 .into_values()
186 .flatten()
187 .chain(toolchains.toolchains)
188 .map(|toolchain| {
189 background_executor.spawn(async move {
190 let python_path = toolchain.path.to_string();
191 let environment_kind = extract_environment_kind(&toolchain.as_json);
192
193 let has_ipykernel = util::command::new_command(&python_path)
194 .args(&["-c", "import ipykernel"])
195 .output()
196 .await
197 .map(|output| output.status.success())
198 .unwrap_or(false);
199
200 let kernelspec = JupyterKernelspec {
201 argv: vec![
202 python_path.clone(),
203 "-m".to_string(),
204 "ipykernel_launcher".to_string(),
205 "-f".to_string(),
206 "{connection_file}".to_string(),
207 ],
208 display_name: toolchain.name.to_string(),
209 language: "python".to_string(),
210 interrupt_mode: None,
211 metadata: None,
212 env: None,
213 };
214
215 KernelSpecification::PythonEnv(PythonEnvKernelSpecification {
216 name: toolchain.name.to_string(),
217 path: PathBuf::from(&python_path),
218 kernelspec,
219 has_ipykernel,
220 environment_kind,
221 })
222 })
223 });
224
225 let kernel_specs = futures::future::join_all(kernelspecs).await;
226
227 anyhow::Ok(kernel_specs)
228 }
229}
230
231pub trait RunningKernel: Send + Debug {
232 fn request_tx(&self) -> mpsc::Sender<JupyterMessage>;
233 fn stdin_tx(&self) -> mpsc::Sender<JupyterMessage>;
234 fn working_directory(&self) -> &PathBuf;
235 fn execution_state(&self) -> &ExecutionState;
236 fn set_execution_state(&mut self, state: ExecutionState);
237 fn kernel_info(&self) -> Option<&KernelInfoReply>;
238 fn set_kernel_info(&mut self, info: KernelInfoReply);
239 fn force_shutdown(&mut self, window: &mut Window, cx: &mut App) -> Task<anyhow::Result<()>>;
240 fn kill(&mut self);
241}
242
243#[derive(Debug, Clone)]
244pub enum KernelStatus {
245 Idle,
246 Busy,
247 Starting,
248 Error,
249 ShuttingDown,
250 Shutdown,
251 Restarting,
252}
253
254impl KernelStatus {
255 pub fn is_connected(&self) -> bool {
256 matches!(self, KernelStatus::Idle | KernelStatus::Busy)
257 }
258}
259
260impl ToString for KernelStatus {
261 fn to_string(&self) -> String {
262 match self {
263 KernelStatus::Idle => "Idle".to_string(),
264 KernelStatus::Busy => "Busy".to_string(),
265 KernelStatus::Starting => "Starting".to_string(),
266 KernelStatus::Error => "Error".to_string(),
267 KernelStatus::ShuttingDown => "Shutting Down".to_string(),
268 KernelStatus::Shutdown => "Shutdown".to_string(),
269 KernelStatus::Restarting => "Restarting".to_string(),
270 }
271 }
272}
273
274#[derive(Debug)]
275pub enum Kernel {
276 RunningKernel(Box<dyn RunningKernel>),
277 StartingKernel(Shared<Task<()>>),
278 ErroredLaunch(String),
279 ShuttingDown,
280 Shutdown,
281 Restarting,
282}
283
284impl From<&Kernel> for KernelStatus {
285 fn from(kernel: &Kernel) -> Self {
286 match kernel {
287 Kernel::RunningKernel(kernel) => match kernel.execution_state() {
288 ExecutionState::Idle => KernelStatus::Idle,
289 ExecutionState::Busy => KernelStatus::Busy,
290 ExecutionState::Unknown => KernelStatus::Error,
291 ExecutionState::Starting => KernelStatus::Starting,
292 ExecutionState::Restarting => KernelStatus::Restarting,
293 ExecutionState::Terminating => KernelStatus::ShuttingDown,
294 ExecutionState::AutoRestarting => KernelStatus::Restarting,
295 ExecutionState::Dead => KernelStatus::Error,
296 ExecutionState::Other(_) => KernelStatus::Error,
297 },
298 Kernel::StartingKernel(_) => KernelStatus::Starting,
299 Kernel::ErroredLaunch(_) => KernelStatus::Error,
300 Kernel::ShuttingDown => KernelStatus::ShuttingDown,
301 Kernel::Shutdown => KernelStatus::Shutdown,
302 Kernel::Restarting => KernelStatus::Restarting,
303 }
304 }
305}
306
307impl Kernel {
308 pub fn status(&self) -> KernelStatus {
309 self.into()
310 }
311
312 pub fn set_execution_state(&mut self, status: &ExecutionState) {
313 if let Kernel::RunningKernel(running_kernel) = self {
314 running_kernel.set_execution_state(status.clone());
315 }
316 }
317
318 pub fn set_kernel_info(&mut self, kernel_info: &KernelInfoReply) {
319 if let Kernel::RunningKernel(running_kernel) = self {
320 running_kernel.set_kernel_info(kernel_info.clone());
321 }
322 }
323
324 pub fn is_shutting_down(&self) -> bool {
325 match self {
326 Kernel::Restarting | Kernel::ShuttingDown => true,
327 Kernel::RunningKernel(_)
328 | Kernel::StartingKernel(_)
329 | Kernel::ErroredLaunch(_)
330 | Kernel::Shutdown => false,
331 }
332 }
333}