1mod native_kernel;
2use std::{fmt::Debug, future::Future, path::PathBuf};
3
4use futures::{
5 channel::mpsc::{self, Receiver},
6 future::Shared,
7 stream,
8};
9use gpui::{App, Entity, Task, Window};
10use language::LanguageName;
11pub use native_kernel::*;
12
13mod remote_kernels;
14use project::{Project, ProjectPath, Toolchains, WorktreeId};
15pub use remote_kernels::*;
16
17use anyhow::Result;
18use gpui::Context;
19use jupyter_protocol::JupyterKernelspec;
20use runtimelib::{ExecutionState, JupyterMessage, KernelInfoReply};
21use ui::{Icon, IconName, SharedString};
22use util::rel_path::RelPath;
23
24pub trait KernelSession: Sized {
25 fn route(&mut self, message: &JupyterMessage, window: &mut Window, cx: &mut Context<Self>);
26 fn kernel_errored(&mut self, error_message: String, cx: &mut Context<Self>);
27}
28
29pub type JupyterMessageChannel = stream::SelectAll<Receiver<JupyterMessage>>;
30
31#[derive(Debug, Clone, PartialEq, Eq)]
32pub enum KernelSpecification {
33 Remote(RemoteKernelSpecification),
34 Jupyter(LocalKernelSpecification),
35 PythonEnv(LocalKernelSpecification),
36}
37
38impl KernelSpecification {
39 pub fn name(&self) -> SharedString {
40 match self {
41 Self::Jupyter(spec) => spec.name.clone().into(),
42 Self::PythonEnv(spec) => spec.name.clone().into(),
43 Self::Remote(spec) => spec.name.clone().into(),
44 }
45 }
46
47 pub fn type_name(&self) -> SharedString {
48 match self {
49 Self::Jupyter(_) => "Jupyter".into(),
50 Self::PythonEnv(_) => "Python Environment".into(),
51 Self::Remote(_) => "Remote".into(),
52 }
53 }
54
55 pub fn path(&self) -> SharedString {
56 SharedString::from(match self {
57 Self::Jupyter(spec) => spec.path.to_string_lossy().into_owned(),
58 Self::PythonEnv(spec) => spec.path.to_string_lossy().into_owned(),
59 Self::Remote(spec) => spec.url.to_string(),
60 })
61 }
62
63 pub fn language(&self) -> SharedString {
64 SharedString::from(match self {
65 Self::Jupyter(spec) => spec.kernelspec.language.clone(),
66 Self::PythonEnv(spec) => spec.kernelspec.language.clone(),
67 Self::Remote(spec) => spec.kernelspec.language.clone(),
68 })
69 }
70
71 pub fn icon(&self, cx: &App) -> Icon {
72 let lang_name = match self {
73 Self::Jupyter(spec) => spec.kernelspec.language.clone(),
74 Self::PythonEnv(spec) => spec.kernelspec.language.clone(),
75 Self::Remote(spec) => spec.kernelspec.language.clone(),
76 };
77
78 file_icons::FileIcons::get(cx)
79 .get_icon_for_type(&lang_name.to_lowercase(), cx)
80 .map(Icon::from_path)
81 .unwrap_or(Icon::new(IconName::ReplNeutral))
82 }
83}
84
85pub fn python_env_kernel_specifications(
86 project: &Entity<Project>,
87 worktree_id: WorktreeId,
88 cx: &mut App,
89) -> impl Future<Output = Result<Vec<KernelSpecification>>> + use<> {
90 let python_language = LanguageName::new_static("Python");
91 let toolchains = project.read(cx).available_toolchains(
92 ProjectPath {
93 worktree_id,
94 path: RelPath::empty().into(),
95 },
96 python_language,
97 cx,
98 );
99 let background_executor = cx.background_executor().clone();
100
101 async move {
102 let (toolchains, user_toolchains) = if let Some(Toolchains {
103 toolchains,
104 root_path: _,
105 user_toolchains,
106 }) = toolchains.await
107 {
108 (toolchains, user_toolchains)
109 } else {
110 return Ok(Vec::new());
111 };
112
113 let kernelspecs = user_toolchains
114 .into_values()
115 .flatten()
116 .chain(toolchains.toolchains)
117 .map(|toolchain| {
118 background_executor.spawn(async move {
119 let python_path = toolchain.path.to_string();
120
121 // Check if ipykernel is installed
122 let ipykernel_check = util::command::new_smol_command(&python_path)
123 .args(&["-c", "import ipykernel"])
124 .output()
125 .await;
126
127 if ipykernel_check.is_ok() && ipykernel_check.unwrap().status.success() {
128 // Create a default kernelspec for this environment
129 let default_kernelspec = JupyterKernelspec {
130 argv: vec![
131 python_path.clone(),
132 "-m".to_string(),
133 "ipykernel_launcher".to_string(),
134 "-f".to_string(),
135 "{connection_file}".to_string(),
136 ],
137 display_name: toolchain.name.to_string(),
138 language: "python".to_string(),
139 interrupt_mode: None,
140 metadata: None,
141 env: None,
142 };
143
144 Some(KernelSpecification::PythonEnv(LocalKernelSpecification {
145 name: toolchain.name.to_string(),
146 path: PathBuf::from(&python_path),
147 kernelspec: default_kernelspec,
148 }))
149 } else {
150 None
151 }
152 })
153 });
154
155 let kernel_specs = futures::future::join_all(kernelspecs)
156 .await
157 .into_iter()
158 .flatten()
159 .collect();
160
161 anyhow::Ok(kernel_specs)
162 }
163}
164
165pub trait RunningKernel: Send + Debug {
166 fn request_tx(&self) -> mpsc::Sender<JupyterMessage>;
167 fn working_directory(&self) -> &PathBuf;
168 fn execution_state(&self) -> &ExecutionState;
169 fn set_execution_state(&mut self, state: ExecutionState);
170 fn kernel_info(&self) -> Option<&KernelInfoReply>;
171 fn set_kernel_info(&mut self, info: KernelInfoReply);
172 fn force_shutdown(&mut self, window: &mut Window, cx: &mut App) -> Task<anyhow::Result<()>>;
173 fn kill(&mut self);
174}
175
176#[derive(Debug, Clone)]
177pub enum KernelStatus {
178 Idle,
179 Busy,
180 Starting,
181 Error,
182 ShuttingDown,
183 Shutdown,
184 Restarting,
185}
186
187impl KernelStatus {
188 pub fn is_connected(&self) -> bool {
189 matches!(self, KernelStatus::Idle | KernelStatus::Busy)
190 }
191}
192
193impl ToString for KernelStatus {
194 fn to_string(&self) -> String {
195 match self {
196 KernelStatus::Idle => "Idle".to_string(),
197 KernelStatus::Busy => "Busy".to_string(),
198 KernelStatus::Starting => "Starting".to_string(),
199 KernelStatus::Error => "Error".to_string(),
200 KernelStatus::ShuttingDown => "Shutting Down".to_string(),
201 KernelStatus::Shutdown => "Shutdown".to_string(),
202 KernelStatus::Restarting => "Restarting".to_string(),
203 }
204 }
205}
206
207#[derive(Debug)]
208pub enum Kernel {
209 RunningKernel(Box<dyn RunningKernel>),
210 StartingKernel(Shared<Task<()>>),
211 ErroredLaunch(String),
212 ShuttingDown,
213 Shutdown,
214 Restarting,
215}
216
217impl From<&Kernel> for KernelStatus {
218 fn from(kernel: &Kernel) -> Self {
219 match kernel {
220 Kernel::RunningKernel(kernel) => match kernel.execution_state() {
221 ExecutionState::Idle => KernelStatus::Idle,
222 ExecutionState::Busy => KernelStatus::Busy,
223 ExecutionState::Unknown => KernelStatus::Error,
224 ExecutionState::Starting => KernelStatus::Starting,
225 ExecutionState::Restarting => KernelStatus::Restarting,
226 ExecutionState::Terminating => KernelStatus::ShuttingDown,
227 ExecutionState::AutoRestarting => KernelStatus::Restarting,
228 ExecutionState::Dead => KernelStatus::Error,
229 ExecutionState::Other(_) => KernelStatus::Error,
230 },
231 Kernel::StartingKernel(_) => KernelStatus::Starting,
232 Kernel::ErroredLaunch(_) => KernelStatus::Error,
233 Kernel::ShuttingDown => KernelStatus::ShuttingDown,
234 Kernel::Shutdown => KernelStatus::Shutdown,
235 Kernel::Restarting => KernelStatus::Restarting,
236 }
237 }
238}
239
240impl Kernel {
241 pub fn status(&self) -> KernelStatus {
242 self.into()
243 }
244
245 pub fn set_execution_state(&mut self, status: &ExecutionState) {
246 if let Kernel::RunningKernel(running_kernel) = self {
247 running_kernel.set_execution_state(status.clone());
248 }
249 }
250
251 pub fn set_kernel_info(&mut self, kernel_info: &KernelInfoReply) {
252 if let Kernel::RunningKernel(running_kernel) = self {
253 running_kernel.set_kernel_info(kernel_info.clone());
254 }
255 }
256
257 pub fn is_shutting_down(&self) -> bool {
258 match self {
259 Kernel::Restarting | Kernel::ShuttingDown => true,
260 Kernel::RunningKernel(_)
261 | Kernel::StartingKernel(_)
262 | Kernel::ErroredLaunch(_)
263 | Kernel::Shutdown => false,
264 }
265 }
266}