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