jupyter_websocket_client/
client.rs

1use anyhow::{Context, Result};
2use async_tungstenite::{
3    async_std::connect_async,
4    tungstenite::{
5        client::IntoClientRequest,
6        http::{HeaderValue, Request, Response},
7    },
8};
9use serde::{Deserialize, Serialize};
10use url::Url;
11
12use crate::websocket::JupyterWebSocket;
13
14pub struct RemoteServer {
15    pub base_url: String,
16    pub token: String,
17}
18
19// Only `id` is used right now, but other fields will be useful when pulling up a listing later
20#[derive(Debug, Serialize, Deserialize)]
21pub struct Kernel {
22    pub id: String,
23    pub name: String,
24    pub last_activity: String,
25    pub execution_state: String,
26    pub connections: u64,
27}
28
29#[derive(Debug, Serialize, Deserialize)]
30pub struct Session {
31    pub id: String,
32    pub path: String,
33    pub name: String,
34    #[serde(rename = "type")]
35    pub session_type: String,
36    pub kernel: Kernel,
37}
38
39#[derive(Debug, Serialize, Deserialize)]
40pub struct NewSession {
41    pub path: String,
42    pub name: Option<String>,
43}
44
45#[derive(Debug, Serialize, Deserialize)]
46pub struct KernelSpec {
47    pub name: String,
48    pub spec: jupyter_protocol::JupyterKernelspec,
49    pub resources: std::collections::HashMap<String, String>,
50}
51
52#[derive(Debug, Serialize, Deserialize)]
53pub struct KernelLaunchRequest {
54    pub name: String,
55    #[serde(skip_serializing_if = "Option::is_none")]
56    pub path: Option<String>,
57}
58
59#[derive(Debug, Serialize, Deserialize)]
60pub struct HelpLink {
61    pub text: String,
62    pub url: String,
63}
64
65#[derive(Debug, Serialize, Deserialize)]
66pub struct KernelSpecsResponse {
67    pub default: String,
68    pub kernelspecs: std::collections::HashMap<String, KernelSpec>,
69}
70
71fn api_url(base_url: &str, path: &str) -> String {
72    format!(
73        "{}/api/{}",
74        base_url.trim_end_matches('/'),
75        path.trim_start_matches('/')
76    )
77}
78
79impl RemoteServer {
80    pub fn from_url(url: &str) -> Result<Self> {
81        let parsed_url = Url::parse(url).context("Failed to parse Jupyter URL")?;
82        let base_url = format!(
83            "{}://{}{}{}",
84            parsed_url.scheme(),
85            parsed_url.host_str().unwrap_or("localhost"),
86            parsed_url
87                .port()
88                .map(|p| format!(":{}", p))
89                .unwrap_or_default(),
90            parsed_url.path().trim_end_matches("/tree")
91        );
92
93        let token = parsed_url
94            .query_pairs()
95            .find(|(key, _)| key == "token")
96            .map(|(_, value)| value.into_owned())
97            .ok_or_else(|| anyhow::anyhow!("Token not found in URL"))?;
98
99        Ok(Self { base_url, token })
100    }
101
102    pub fn api_url(&self, path: &str) -> String {
103        api_url(&self.base_url, path)
104    }
105
106    /// Connect to a kernel by ID
107    ///
108    /// ```rust
109    /// use jupyter_websocket_client::RemoteServer;
110    ///
111    /// use jupyter_protocol::{KernelInfoRequest, JupyterMessageContent};
112    ///
113    /// // Import the sink and stream extensions to allow splitting the socket into a writer and reader pair
114    /// use futures::{SinkExt as _, StreamExt as _};
115    ///
116    /// pub async fn connect_kernel() -> anyhow::Result<()> {
117    ///     let server = RemoteServer::from_url(
118    ///         "http://127.0.0.1:8888/lab?token=f487535a46268da4a0752c0e162c873b721e33a9e6ec8390"
119    ///     )?;
120    ///
121    ///     // You'll need to launch a kernel and get a kernel ID using your own HTTP
122    ///     // request library
123    ///     let kernel_id = "1057-1057-1057-1057";
124    ///
125    ///     let (kernel_socket, response) = server.connect_to_kernel(kernel_id).await?;
126    ///
127    ///     let (mut w, mut r) = kernel_socket.split();
128    ///
129    ///     w.send(KernelInfoRequest {}.into()).await?;
130    ///
131    ///     while let Some(response) = r.next().await.transpose()? {
132    ///         match response.content {
133    ///             JupyterMessageContent::KernelInfoReply(kernel_info_reply) => {
134    ///                 println!("Received kernel_info_reply");
135    ///                 println!("{:?}", kernel_info_reply);
136    ///                 break;
137    ///             }
138    ///             other => {
139    ///                 println!("Received");
140    ///                 println!("{:?}", other);
141    ///             }
142    ///         }
143    ///     }
144    ///
145    ///     Ok(())
146    /// }
147    /// ```
148    pub async fn connect_to_kernel(
149        &self,
150        kernel_id: &str,
151    ) -> Result<(JupyterWebSocket, Response<Option<Vec<u8>>>)> {
152        let ws_url = format!(
153            "{}?token={}",
154            api_url(&self.base_url, &format!("kernels/{}/channels", kernel_id))
155                .replace("http", "ws"),
156            self.token
157        );
158
159        let mut req: Request<()> = ws_url.into_client_request()?;
160        let headers = req.headers_mut();
161        headers.insert(
162            "User-Agent",
163            HeaderValue::from_str("runtimed/jupyter-websocket-client")?,
164        );
165
166        let response = connect_async(req).await;
167
168        let (ws_stream, response) = response?;
169
170        Ok((JupyterWebSocket { inner: ws_stream }, response))
171    }
172}
jupyter_websocket_client/client.rs

jupyter_websocket_client/
client.rs
