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;
7use log;
8pub use native_kernel::*;
9
10mod remote_kernels;
11use project::{Project, ProjectPath, Toolchains, WorktreeId};
12pub use remote_kernels::*;
13
14mod ssh_kernel;
15pub use ssh_kernel::*;
16
17mod wsl_kernel;
18pub use wsl_kernel::*;
19
20use std::collections::HashMap;
21
22use anyhow::Result;
23use futures::{FutureExt, StreamExt};
24use gpui::{AppContext, AsyncWindowContext, Context};
25use jupyter_protocol::{JupyterKernelspec, JupyterMessageContent};
26use runtimelib::{
27 ClientControlConnection, ClientIoPubConnection, ClientShellConnection, ClientStdinConnection,
28 ExecutionState, JupyterMessage, KernelInfoReply,
29};
30use ui::{Icon, IconName, SharedString};
31use util::rel_path::RelPath;
32
33pub fn start_kernel_tasks<S: KernelSession + 'static>(
34 session: Entity<S>,
35 iopub_socket: ClientIoPubConnection,
36 shell_socket: ClientShellConnection,
37 control_socket: ClientControlConnection,
38 stdin_socket: ClientStdinConnection,
39 cx: &mut AsyncWindowContext,
40) -> (
41 futures::channel::mpsc::Sender<JupyterMessage>,
42 futures::channel::mpsc::Sender<JupyterMessage>,
43) {
44 let (mut shell_send, shell_recv) = shell_socket.split();
45 let (mut control_send, control_recv) = control_socket.split();
46 let (mut stdin_send, stdin_recv) = stdin_socket.split();
47
48 let (request_tx, mut request_rx) = futures::channel::mpsc::channel::<JupyterMessage>(100);
49 let (stdin_tx, mut stdin_rx) = futures::channel::mpsc::channel::<JupyterMessage>(100);
50
51 let recv_task = cx.spawn({
52 let session = session.clone();
53 let mut iopub = iopub_socket;
54 let mut shell = shell_recv;
55 let mut control = control_recv;
56 let mut stdin = stdin_recv;
57
58 async move |cx| -> anyhow::Result<()> {
59 loop {
60 let (channel, result) = futures::select! {
61 msg = iopub.read().fuse() => ("iopub", msg),
62 msg = shell.read().fuse() => ("shell", msg),
63 msg = control.read().fuse() => ("control", msg),
64 msg = stdin.read().fuse() => ("stdin", msg),
65 };
66 match result {
67 Ok(message) => {
68 session
69 .update_in(cx, |session, window, cx| {
70 session.route(&message, window, cx);
71 })
72 .ok();
73 }
74 Err(
75 ref err @ (runtimelib::RuntimeError::ParseError { .. }
76 | runtimelib::RuntimeError::SerdeError(_)),
77 ) => {
78 let error_detail = format!("Kernel issue on {channel} channel\n\n{err}");
79 log::warn!("kernel: {error_detail}");
80 session
81 .update_in(cx, |session, _window, cx| {
82 session.kernel_errored(error_detail, cx);
83 cx.notify();
84 })
85 .ok();
86 }
87 Err(err) => {
88 log::warn!("kernel: error reading from {channel}: {err:?}");
89 anyhow::bail!("{channel} recv: {err}");
90 }
91 }
92 }
93 }
94 });
95
96 let routing_task = cx.background_spawn(async move {
97 while let Some(message) = request_rx.next().await {
98 match message.content {
99 JupyterMessageContent::DebugRequest(_)
100 | JupyterMessageContent::InterruptRequest(_)
101 | JupyterMessageContent::ShutdownRequest(_) => {
102 control_send.send(message).await?;
103 }
104 _ => {
105 shell_send.send(message).await?;
106 }
107 }
108 }
109 anyhow::Ok(())
110 });
111
112 let stdin_routing_task = cx.background_spawn(async move {
113 while let Some(message) = stdin_rx.next().await {
114 stdin_send.send(message).await?;
115 }
116 anyhow::Ok(())
117 });
118
119 cx.spawn({
120 async move |cx| {
121 async fn with_name(
122 name: &'static str,
123 task: Task<Result<()>>,
124 ) -> (&'static str, Result<()>) {
125 (name, task.await)
126 }
127
128 let mut tasks = futures::stream::FuturesUnordered::new();
129 tasks.push(with_name("recv task", recv_task));
130 tasks.push(with_name("routing task", routing_task));
131 tasks.push(with_name("stdin routing task", stdin_routing_task));
132
133 while let Some((name, result)) = tasks.next().await {
134 if let Err(err) = result {
135 session.update(cx, |session, cx| {
136 session.kernel_errored(format!("handling failed for {name}: {err}"), cx);
137 cx.notify();
138 });
139 }
140 }
141 }
142 })
143 .detach();
144
145 (request_tx, stdin_tx)
146}
147
148pub trait KernelSession: Sized {
149 fn route(&mut self, message: &JupyterMessage, window: &mut Window, cx: &mut Context<Self>);
150 fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>);
151}
152
153#[derive(Debug, Clone)]
154pub struct PythonEnvKernelSpecification {
155 pub name: String,
156 pub path: PathBuf,
157 pub kernelspec: JupyterKernelspec,
158 pub has_ipykernel: bool,
159 /// Display label for the environment type: "venv", "Conda", "Pyenv", etc.
160 pub environment_kind: Option<String>,
161}
162
163impl PartialEq for PythonEnvKernelSpecification {
164 fn eq(&self, other: &Self) -> bool {
165 self.name == other.name && self.path == other.path
166 }
167}
168
169impl Eq for PythonEnvKernelSpecification {}
170
171impl PythonEnvKernelSpecification {
172 pub fn as_local_spec(&self) -> LocalKernelSpecification {
173 LocalKernelSpecification {
174 name: self.name.clone(),
175 path: self.path.clone(),
176 kernelspec: self.kernelspec.clone(),
177 }
178 }
179}
180
181#[derive(Debug, Clone, PartialEq, Eq)]
182pub enum KernelSpecification {
183 JupyterServer(RemoteKernelSpecification),
184 Jupyter(LocalKernelSpecification),
185 PythonEnv(PythonEnvKernelSpecification),
186 SshRemote(SshRemoteKernelSpecification),
187 WslRemote(WslKernelSpecification),
188}
189
190#[derive(Debug, Clone)]
191pub struct SshRemoteKernelSpecification {
192 pub name: String,
193 pub path: SharedString,
194 pub kernelspec: JupyterKernelspec,
195}
196
197#[derive(Debug, Clone)]
198pub struct WslKernelSpecification {
199 pub name: String,
200 pub kernelspec: JupyterKernelspec,
201 pub distro: String,
202}
203
204impl PartialEq for SshRemoteKernelSpecification {
205 fn eq(&self, other: &Self) -> bool {
206 self.name == other.name
207 && self.kernelspec.argv == other.kernelspec.argv
208 && self.path == other.path
209 && self.kernelspec.display_name == other.kernelspec.display_name
210 && self.kernelspec.language == other.kernelspec.language
211 && self.kernelspec.interrupt_mode == other.kernelspec.interrupt_mode
212 && self.kernelspec.env == other.kernelspec.env
213 && self.kernelspec.metadata == other.kernelspec.metadata
214 }
215}
216
217impl Eq for SshRemoteKernelSpecification {}
218
219impl PartialEq for WslKernelSpecification {
220 fn eq(&self, other: &Self) -> bool {
221 self.name == other.name
222 && self.kernelspec.argv == other.kernelspec.argv
223 && self.kernelspec.display_name == other.kernelspec.display_name
224 && self.kernelspec.language == other.kernelspec.language
225 && self.kernelspec.interrupt_mode == other.kernelspec.interrupt_mode
226 && self.kernelspec.env == other.kernelspec.env
227 && self.kernelspec.metadata == other.kernelspec.metadata
228 && self.distro == other.distro
229 }
230}
231
232impl Eq for WslKernelSpecification {}
233
234impl KernelSpecification {
235 pub fn name(&self) -> SharedString {
236 match self {
237 Self::Jupyter(spec) => spec.name.clone().into(),
238 Self::PythonEnv(spec) => spec.name.clone().into(),
239 Self::JupyterServer(spec) => spec.name.clone().into(),
240 Self::SshRemote(spec) => spec.name.clone().into(),
241 Self::WslRemote(spec) => spec.name.clone().into(),
242 }
243 }
244
245 pub fn type_name(&self) -> SharedString {
246 match self {
247 Self::Jupyter(_) => "Jupyter".into(),
248 Self::PythonEnv(spec) => SharedString::from(
249 spec.environment_kind
250 .clone()
251 .unwrap_or_else(|| "Python Environment".to_string()),
252 ),
253 Self::JupyterServer(_) => "Jupyter Server".into(),
254 Self::SshRemote(_) => "SSH Remote".into(),
255 Self::WslRemote(_) => "WSL Remote".into(),
256 }
257 }
258
259 pub fn path(&self) -> SharedString {
260 SharedString::from(match self {
261 Self::Jupyter(spec) => spec.path.to_string_lossy().into_owned(),
262 Self::PythonEnv(spec) => spec.path.to_string_lossy().into_owned(),
263 Self::JupyterServer(spec) => spec.url.to_string(),
264 Self::SshRemote(spec) => spec.path.to_string(),
265 Self::WslRemote(_) => "WSL".to_string(),
266 })
267 }
268
269 pub fn language(&self) -> SharedString {
270 SharedString::from(match self {
271 Self::Jupyter(spec) => spec.kernelspec.language.clone(),
272 Self::PythonEnv(spec) => spec.kernelspec.language.clone(),
273 Self::JupyterServer(spec) => spec.kernelspec.language.clone(),
274 Self::SshRemote(spec) => spec.kernelspec.language.clone(),
275 Self::WslRemote(spec) => spec.kernelspec.language.clone(),
276 })
277 }
278
279 pub fn has_ipykernel(&self) -> bool {
280 match self {
281 Self::Jupyter(_) | Self::JupyterServer(_) | Self::SshRemote(_) | Self::WslRemote(_) => {
282 true
283 }
284 Self::PythonEnv(spec) => spec.has_ipykernel,
285 }
286 }
287
288 pub fn environment_kind_label(&self) -> Option<SharedString> {
289 match self {
290 Self::PythonEnv(spec) => spec
291 .environment_kind
292 .as_ref()
293 .map(|kind| SharedString::from(kind.clone())),
294 Self::Jupyter(_) => Some("Jupyter".into()),
295 Self::JupyterServer(_) => Some("Jupyter Server".into()),
296 Self::SshRemote(_) => Some("SSH Remote".into()),
297 Self::WslRemote(_) => Some("WSL Remote".into()),
298 }
299 }
300
301 pub fn icon(&self, cx: &App) -> Icon {
302 let lang_name = match self {
303 Self::Jupyter(spec) => spec.kernelspec.language.clone(),
304 Self::PythonEnv(spec) => spec.kernelspec.language.clone(),
305 Self::JupyterServer(spec) => spec.kernelspec.language.clone(),
306 Self::SshRemote(spec) => spec.kernelspec.language.clone(),
307 Self::WslRemote(spec) => spec.kernelspec.language.clone(),
308 };
309
310 file_icons::FileIcons::get(cx)
311 .get_icon_for_type(&lang_name.to_lowercase(), cx)
312 .map(Icon::from_path)
313 .unwrap_or(Icon::new(IconName::ReplNeutral))
314 }
315}
316
317fn extract_environment_kind(toolchain_json: &serde_json::Value) -> Option<String> {
318 let kind_str = toolchain_json.get("kind")?.as_str()?;
319 let label = match kind_str {
320 "Conda" => "Conda",
321 "Pixi" => "pixi",
322 "Homebrew" => "Homebrew",
323 "Pyenv" => "global (Pyenv)",
324 "GlobalPaths" => "global",
325 "PyenvVirtualEnv" => "Pyenv",
326 "Pipenv" => "Pipenv",
327 "Poetry" => "Poetry",
328 "MacPythonOrg" => "global (Python.org)",
329 "MacCommandLineTools" => "global (Command Line Tools for Xcode)",
330 "LinuxGlobal" => "global",
331 "MacXCode" => "global (Xcode)",
332 "Venv" => "venv",
333 "VirtualEnv" => "virtualenv",
334 "VirtualEnvWrapper" => "virtualenvwrapper",
335 "WindowsStore" => "global (Windows Store)",
336 "WindowsRegistry" => "global (Windows Registry)",
337 "Uv" => "uv",
338 "UvWorkspace" => "uv (Workspace)",
339 _ => kind_str,
340 };
341 Some(label.to_string())
342}
343
344pub fn python_env_kernel_specifications(
345 project: &Entity<Project>,
346 worktree_id: WorktreeId,
347 cx: &mut App,
348) -> impl Future<Output = Result<Vec<KernelSpecification>>> + use<> {
349 let python_language = LanguageName::new_static("Python");
350 let is_remote = project.read(cx).is_remote();
351 log::info!("python_env_kernel_specifications: is_remote: {}", is_remote);
352
353 let toolchains = project.read(cx).available_toolchains(
354 ProjectPath {
355 worktree_id,
356 path: RelPath::empty().into(),
357 },
358 python_language,
359 cx,
360 );
361 #[allow(unused)]
362 let worktree_root_path: Option<std::sync::Arc<std::path::Path>> = project
363 .read(cx)
364 .worktree_for_id(worktree_id, cx)
365 .map(|w| w.read(cx).abs_path());
366
367 let background_executor = cx.background_executor().clone();
368
369 async move {
370 let (toolchains, user_toolchains) = if let Some(Toolchains {
371 toolchains,
372 root_path: _,
373 user_toolchains,
374 }) = toolchains.await
375 {
376 (toolchains, user_toolchains)
377 } else {
378 return Ok(Vec::new());
379 };
380
381 let kernelspecs = user_toolchains
382 .into_values()
383 .flatten()
384 .chain(toolchains.toolchains)
385 .map(|toolchain| {
386 background_executor.spawn(async move {
387 // For remote projects, we assume python is available assuming toolchain is reported.
388 // We can skip the `ipykernel` check or run it remotely.
389 // For MVP, lets trust the toolchain existence or do the check if it's cheap.
390 // `new_smol_command` runs locally. We need to run remotely if `is_remote`.
391
392 if is_remote {
393 log::info!(
394 "python_env_kernel_specifications: returning SshRemote for toolchain {}",
395 toolchain.name
396 );
397 let default_kernelspec = JupyterKernelspec {
398 argv: vec![
399 toolchain.path.to_string(),
400 "-m".to_string(),
401 "ipykernel_launcher".to_string(),
402 "-f".to_string(),
403 "{connection_file}".to_string(),
404 ],
405 display_name: toolchain.name.to_string(),
406 language: "python".to_string(),
407 interrupt_mode: None,
408 metadata: None,
409 env: None,
410 };
411
412 return Some(KernelSpecification::SshRemote(
413 SshRemoteKernelSpecification {
414 name: format!("Remote {}", toolchain.name),
415 path: toolchain.path.clone(),
416 kernelspec: default_kernelspec,
417 },
418 ));
419 }
420
421 let python_path = toolchain.path.to_string();
422 let environment_kind = extract_environment_kind(&toolchain.as_json);
423
424 let has_ipykernel = util::command::new_command(&python_path)
425 .args(&["-c", "import ipykernel"])
426 .output()
427 .await
428 .map(|output| output.status.success())
429 .unwrap_or(false);
430
431 let mut env = HashMap::new();
432 if let Some(python_bin_dir) = PathBuf::from(&python_path).parent() {
433 if let Some(path_var) = std::env::var_os("PATH") {
434 let mut paths = std::env::split_paths(&path_var).collect::<Vec<_>>();
435 paths.insert(0, python_bin_dir.to_path_buf());
436 if let Ok(new_path) = std::env::join_paths(paths) {
437 env.insert("PATH".to_string(), new_path.to_string_lossy().to_string());
438 }
439 }
440
441 if let Some(venv_root) = python_bin_dir.parent() {
442 env.insert("VIRTUAL_ENV".to_string(), venv_root.to_string_lossy().to_string());
443 }
444 }
445
446 log::info!("Preparing Python kernel for toolchain: {}", toolchain.name);
447 log::info!("Python path: {}", python_path);
448 if let Some(path) = env.get("PATH") {
449 log::info!("Kernel PATH: {}", path);
450 } else {
451 log::info!("Kernel PATH not set in env");
452 }
453 if let Some(venv) = env.get("VIRTUAL_ENV") {
454 log::info!("Kernel VIRTUAL_ENV: {}", venv);
455 }
456
457 let kernelspec = JupyterKernelspec {
458 argv: vec![
459 python_path.clone(),
460 "-m".to_string(),
461 "ipykernel_launcher".to_string(),
462 "-f".to_string(),
463 "{connection_file}".to_string(),
464 ],
465 display_name: toolchain.name.to_string(),
466 language: "python".to_string(),
467 interrupt_mode: None,
468 metadata: None,
469 env: Some(env),
470 };
471
472 Some(KernelSpecification::PythonEnv(PythonEnvKernelSpecification {
473 name: toolchain.name.to_string(),
474 path: PathBuf::from(&python_path),
475 kernelspec,
476 has_ipykernel,
477 environment_kind,
478 }))
479 })
480 });
481
482 #[allow(unused_mut)]
483 let mut kernel_specs: Vec<KernelSpecification> = futures::future::join_all(kernelspecs)
484 .await
485 .into_iter()
486 .flatten()
487 .collect();
488
489 #[cfg(target_os = "windows")]
490 if kernel_specs.is_empty() && !is_remote {
491 if let Some(root_path) = worktree_root_path {
492 let root_path_str: std::borrow::Cow<str> = root_path.to_string_lossy();
493 let (distro, internal_path) =
494 if let Some(path_without_prefix) = root_path_str.strip_prefix(r"\\wsl$\") {
495 if let Some((distro, path)) = path_without_prefix.split_once('\\') {
496 let replaced_path: String = path.replace('\\', "/");
497 (Some(distro), Some(format!("/{}", replaced_path)))
498 } else {
499 (Some(path_without_prefix), Some("/".to_string()))
500 }
501 } else if let Some(path_without_prefix) =
502 root_path_str.strip_prefix(r"\\wsl.localhost\")
503 {
504 if let Some((distro, path)) = path_without_prefix.split_once('\\') {
505 let replaced_path: String = path.replace('\\', "/");
506 (Some(distro), Some(format!("/{}", replaced_path)))
507 } else {
508 (Some(path_without_prefix), Some("/".to_string()))
509 }
510 } else {
511 (None, None)
512 };
513
514 if let (Some(distro), Some(internal_path)) = (distro, internal_path) {
515 let python_path = format!("{}/.venv/bin/python", internal_path);
516 let check = util::command::new_command("wsl")
517 .args(&["-d", distro, "test", "-f", &python_path])
518 .output()
519 .await;
520
521 if check.is_ok() && check.unwrap().status.success() {
522 let default_kernelspec = JupyterKernelspec {
523 argv: vec![
524 python_path.clone(),
525 "-m".to_string(),
526 "ipykernel_launcher".to_string(),
527 "-f".to_string(),
528 "{connection_file}".to_string(),
529 ],
530 display_name: format!("WSL: {} (.venv)", distro),
531 language: "python".to_string(),
532 interrupt_mode: None,
533 metadata: None,
534 env: None,
535 };
536
537 kernel_specs.push(KernelSpecification::WslRemote(WslKernelSpecification {
538 name: format!("WSL: {} (.venv)", distro),
539 kernelspec: default_kernelspec,
540 distro: distro.to_string(),
541 }));
542 } else {
543 let check_system = util::command::new_command("wsl")
544 .args(&["-d", distro, "command", "-v", "python3"])
545 .output()
546 .await;
547
548 if check_system.is_ok() && check_system.unwrap().status.success() {
549 let default_kernelspec = JupyterKernelspec {
550 argv: vec![
551 "python3".to_string(),
552 "-m".to_string(),
553 "ipykernel_launcher".to_string(),
554 "-f".to_string(),
555 "{connection_file}".to_string(),
556 ],
557 display_name: format!("WSL: {} (System)", distro),
558 language: "python".to_string(),
559 interrupt_mode: None,
560 metadata: None,
561 env: None,
562 };
563
564 kernel_specs.push(KernelSpecification::WslRemote(
565 WslKernelSpecification {
566 name: format!("WSL: {} (System)", distro),
567 kernelspec: default_kernelspec,
568 distro: distro.to_string(),
569 },
570 ));
571 }
572 }
573 }
574 }
575 }
576
577 anyhow::Ok(kernel_specs)
578 }
579}
580
581pub trait RunningKernel: Send + Debug {
582 fn request_tx(&self) -> mpsc::Sender<JupyterMessage>;
583 fn stdin_tx(&self) -> mpsc::Sender<JupyterMessage>;
584 fn working_directory(&self) -> &PathBuf;
585 fn execution_state(&self) -> &ExecutionState;
586 fn set_execution_state(&mut self, state: ExecutionState);
587 fn kernel_info(&self) -> Option<&KernelInfoReply>;
588 fn set_kernel_info(&mut self, info: KernelInfoReply);
589 fn force_shutdown(&mut self, window: &mut Window, cx: &mut App) -> Task<anyhow::Result<()>>;
590 fn kill(&mut self);
591}
592
593#[derive(Debug, Clone)]
594pub enum KernelStatus {
595 Idle,
596 Busy,
597 Starting,
598 Error,
599 ShuttingDown,
600 Shutdown,
601 Restarting,
602}
603
604impl KernelStatus {
605 pub fn is_connected(&self) -> bool {
606 matches!(self, KernelStatus::Idle | KernelStatus::Busy)
607 }
608}
609
610impl ToString for KernelStatus {
611 fn to_string(&self) -> String {
612 match self {
613 KernelStatus::Idle => "Idle".to_string(),
614 KernelStatus::Busy => "Busy".to_string(),
615 KernelStatus::Starting => "Starting".to_string(),
616 KernelStatus::Error => "Error".to_string(),
617 KernelStatus::ShuttingDown => "Shutting Down".to_string(),
618 KernelStatus::Shutdown => "Shutdown".to_string(),
619 KernelStatus::Restarting => "Restarting".to_string(),
620 }
621 }
622}
623
624#[derive(Debug)]
625pub enum Kernel {
626 RunningKernel(Box<dyn RunningKernel>),
627 StartingKernel(Shared<Task<()>>),
628 ErroredLaunch(String),
629 ShuttingDown,
630 Shutdown,
631 Restarting,
632}
633
634impl From<&Kernel> for KernelStatus {
635 fn from(kernel: &Kernel) -> Self {
636 match kernel {
637 Kernel::RunningKernel(kernel) => match kernel.execution_state() {
638 ExecutionState::Idle => KernelStatus::Idle,
639 ExecutionState::Busy => KernelStatus::Busy,
640 ExecutionState::Unknown => KernelStatus::Error,
641 ExecutionState::Starting => KernelStatus::Starting,
642 ExecutionState::Restarting => KernelStatus::Restarting,
643 ExecutionState::Terminating => KernelStatus::ShuttingDown,
644 ExecutionState::AutoRestarting => KernelStatus::Restarting,
645 ExecutionState::Dead => KernelStatus::Error,
646 ExecutionState::Other(_) => KernelStatus::Error,
647 },
648 Kernel::StartingKernel(_) => KernelStatus::Starting,
649 Kernel::ErroredLaunch(_) => KernelStatus::Error,
650 Kernel::ShuttingDown => KernelStatus::ShuttingDown,
651 Kernel::Shutdown => KernelStatus::Shutdown,
652 Kernel::Restarting => KernelStatus::Restarting,
653 }
654 }
655}
656
657impl Kernel {
658 pub fn status(&self) -> KernelStatus {
659 self.into()
660 }
661
662 pub fn set_execution_state(&mut self, status: &ExecutionState) {
663 if let Kernel::RunningKernel(running_kernel) = self {
664 running_kernel.set_execution_state(status.clone());
665 }
666 }
667
668 pub fn set_kernel_info(&mut self, kernel_info: &KernelInfoReply) {
669 if let Kernel::RunningKernel(running_kernel) = self {
670 running_kernel.set_kernel_info(kernel_info.clone());
671 }
672 }
673
674 pub fn is_shutting_down(&self) -> bool {
675 match self {
676 Kernel::Restarting | Kernel::ShuttingDown => true,
677 Kernel::RunningKernel(_)
678 | Kernel::StartingKernel(_)
679 | Kernel::ErroredLaunch(_)
680 | Kernel::Shutdown => false,
681 }
682 }
683}