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