1use futures::{SinkExt as _, channel::mpsc};
2use gpui::{App, AppContext as _, Entity, Task, Window};
3use http_client::{AsyncBody, HttpClient, Request};
4use jupyter_protocol::{ExecutionState, JupyterKernelspec, JupyterMessage, KernelInfoReply};
5
6use async_tungstenite::{
7 async_std::connect_async,
8 tungstenite::{client::IntoClientRequest, http::HeaderValue},
9};
10
11use futures::StreamExt;
12use smol::io::AsyncReadExt as _;
13
14use crate::Session;
15
16use super::RunningKernel;
17use anyhow::Result;
18use jupyter_websocket_client::{
19 JupyterWebSocket, JupyterWebSocketReader, JupyterWebSocketWriter, KernelLaunchRequest,
20 KernelSpecsResponse, RemoteServer,
21};
22use std::{fmt::Debug, sync::Arc};
23
24#[derive(Debug, Clone)]
25pub struct RemoteKernelSpecification {
26 pub name: String,
27 pub url: String,
28 pub token: String,
29 pub kernelspec: JupyterKernelspec,
30}
31
32pub async fn launch_remote_kernel(
33 remote_server: &RemoteServer,
34 http_client: Arc<dyn HttpClient>,
35 kernel_name: &str,
36 _path: &str,
37) -> Result<String> {
38 //
39 let kernel_launch_request = KernelLaunchRequest {
40 name: kernel_name.to_string(),
41 // Note: since the path we have locally may not be the same as the one on the remote server,
42 // we don't send it. We'll have to evaluate this decision along the way.
43 path: None,
44 };
45
46 let kernel_launch_request = serde_json::to_string(&kernel_launch_request)?;
47
48 let request = Request::builder()
49 .method("POST")
50 .uri(&remote_server.api_url("/kernels"))
51 .header("Authorization", format!("token {}", remote_server.token))
52 .body(AsyncBody::from(kernel_launch_request))?;
53
54 let response = http_client.send(request).await?;
55
56 if !response.status().is_success() {
57 let mut body = String::new();
58 response.into_body().read_to_string(&mut body).await?;
59 return Err(anyhow::anyhow!("Failed to launch kernel: {}", body));
60 }
61
62 let mut body = String::new();
63 response.into_body().read_to_string(&mut body).await?;
64
65 let response: jupyter_websocket_client::Kernel = serde_json::from_str(&body)?;
66
67 Ok(response.id)
68}
69
70pub async fn list_remote_kernelspecs(
71 remote_server: RemoteServer,
72 http_client: Arc<dyn HttpClient>,
73) -> Result<Vec<RemoteKernelSpecification>> {
74 let url = remote_server.api_url("/kernelspecs");
75
76 let request = Request::builder()
77 .method("GET")
78 .uri(&url)
79 .header("Authorization", format!("token {}", remote_server.token))
80 .body(AsyncBody::default())?;
81
82 let response = http_client.send(request).await?;
83
84 if response.status().is_success() {
85 let mut body = response.into_body();
86
87 let mut body_bytes = Vec::new();
88 body.read_to_end(&mut body_bytes).await?;
89
90 let kernel_specs: KernelSpecsResponse = serde_json::from_slice(&body_bytes)?;
91
92 let remote_kernelspecs = kernel_specs
93 .kernelspecs
94 .into_iter()
95 .map(|(name, spec)| RemoteKernelSpecification {
96 name: name.clone(),
97 url: remote_server.base_url.clone(),
98 token: remote_server.token.clone(),
99 kernelspec: spec.spec,
100 })
101 .collect::<Vec<RemoteKernelSpecification>>();
102
103 if remote_kernelspecs.is_empty() {
104 Err(anyhow::anyhow!("No kernel specs found"))
105 } else {
106 Ok(remote_kernelspecs.clone())
107 }
108 } else {
109 Err(anyhow::anyhow!(
110 "Failed to fetch kernel specs: {}",
111 response.status()
112 ))
113 }
114}
115
116impl PartialEq for RemoteKernelSpecification {
117 fn eq(&self, other: &Self) -> bool {
118 self.name == other.name && self.url == other.url
119 }
120}
121
122impl Eq for RemoteKernelSpecification {}
123
124pub struct RemoteRunningKernel {
125 remote_server: RemoteServer,
126 _receiving_task: Task<Result<()>>,
127 _routing_task: Task<Result<()>>,
128 http_client: Arc<dyn HttpClient>,
129 pub working_directory: std::path::PathBuf,
130 pub request_tx: mpsc::Sender<JupyterMessage>,
131 pub execution_state: ExecutionState,
132 pub kernel_info: Option<KernelInfoReply>,
133 pub kernel_id: String,
134}
135
136impl RemoteRunningKernel {
137 pub fn new(
138 kernelspec: RemoteKernelSpecification,
139 working_directory: std::path::PathBuf,
140 session: Entity<Session>,
141 window: &mut Window,
142 cx: &mut App,
143 ) -> Task<Result<Box<dyn RunningKernel>>> {
144 let remote_server = RemoteServer {
145 base_url: kernelspec.url,
146 token: kernelspec.token,
147 };
148
149 let http_client = cx.http_client();
150
151 window.spawn(cx, async move |cx| {
152 let kernel_id = launch_remote_kernel(
153 &remote_server,
154 http_client.clone(),
155 &kernelspec.name,
156 working_directory.to_str().unwrap_or_default(),
157 )
158 .await?;
159
160 let ws_url = format!(
161 "{}/api/kernels/{}/channels?token={}",
162 remote_server.base_url.replace("http", "ws"),
163 kernel_id,
164 remote_server.token
165 );
166
167 let mut req: Request<()> = ws_url.into_client_request()?;
168 let headers = req.headers_mut();
169
170 headers.insert(
171 "User-Agent",
172 HeaderValue::from_str(&format!(
173 "Zed/{} ({}; {})",
174 "repl",
175 std::env::consts::OS,
176 std::env::consts::ARCH
177 ))?,
178 );
179
180 let response = connect_async(req).await;
181
182 let (ws_stream, _response) = response?;
183
184 let kernel_socket = JupyterWebSocket { inner: ws_stream };
185
186 let (mut w, mut r): (JupyterWebSocketWriter, JupyterWebSocketReader) =
187 kernel_socket.split();
188
189 let (request_tx, mut request_rx) =
190 futures::channel::mpsc::channel::<JupyterMessage>(100);
191
192 let routing_task = cx.background_spawn({
193 async move {
194 while let Some(message) = request_rx.next().await {
195 w.send(message).await.ok();
196 }
197 Ok(())
198 }
199 });
200
201 let receiving_task = cx.spawn({
202 let session = session.clone();
203
204 async move |cx| {
205 while let Some(message) = r.next().await {
206 match message {
207 Ok(message) => {
208 session
209 .update_in(cx, |session, window, cx| {
210 session.route(&message, window, cx);
211 })
212 .ok();
213 }
214 Err(e) => {
215 log::error!("Error receiving message: {:?}", e);
216 }
217 }
218 }
219 Ok(())
220 }
221 });
222
223 anyhow::Ok(Box::new(Self {
224 _routing_task: routing_task,
225 _receiving_task: receiving_task,
226 remote_server,
227 working_directory,
228 request_tx,
229 // todo(kyle): pull this from the kernel API to start with
230 execution_state: ExecutionState::Idle,
231 kernel_info: None,
232 kernel_id,
233 http_client: http_client.clone(),
234 }) as Box<dyn RunningKernel>)
235 })
236 }
237}
238
239impl Debug for RemoteRunningKernel {
240 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
241 f.debug_struct("RemoteRunningKernel")
242 // custom debug that keeps tokens out of logs
243 .field("remote_server url", &self.remote_server.base_url)
244 .field("working_directory", &self.working_directory)
245 .field("request_tx", &self.request_tx)
246 .field("execution_state", &self.execution_state)
247 .field("kernel_info", &self.kernel_info)
248 .finish()
249 }
250}
251
252impl RunningKernel for RemoteRunningKernel {
253 fn request_tx(&self) -> futures::channel::mpsc::Sender<runtimelib::JupyterMessage> {
254 self.request_tx.clone()
255 }
256
257 fn working_directory(&self) -> &std::path::PathBuf {
258 &self.working_directory
259 }
260
261 fn execution_state(&self) -> &runtimelib::ExecutionState {
262 &self.execution_state
263 }
264
265 fn set_execution_state(&mut self, state: runtimelib::ExecutionState) {
266 self.execution_state = state;
267 }
268
269 fn kernel_info(&self) -> Option<&runtimelib::KernelInfoReply> {
270 self.kernel_info.as_ref()
271 }
272
273 fn set_kernel_info(&mut self, info: runtimelib::KernelInfoReply) {
274 self.kernel_info = Some(info);
275 }
276
277 fn force_shutdown(&mut self, window: &mut Window, cx: &mut App) -> Task<anyhow::Result<()>> {
278 let url = self
279 .remote_server
280 .api_url(&format!("/kernels/{}", self.kernel_id));
281 let token = self.remote_server.token.clone();
282 let http_client = self.http_client.clone();
283
284 window.spawn(cx, async move |_| {
285 let request = Request::builder()
286 .method("DELETE")
287 .uri(&url)
288 .header("Authorization", format!("token {}", token))
289 .body(AsyncBody::default())?;
290
291 let response = http_client.send(request).await?;
292
293 if response.status().is_success() {
294 Ok(())
295 } else {
296 Err(anyhow::anyhow!(
297 "Failed to shutdown kernel: {}",
298 response.status()
299 ))
300 }
301 })
302 }
303}