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