rig_tungstenite/
lib.rs

1#![cfg_attr(docsrs, feature(doc_cfg))]
2#![cfg_attr(
3    test,
4    allow(
5        clippy::expect_used,
6        clippy::indexing_slicing,
7        clippy::panic,
8        clippy::unwrap_used
9    )
10)]
11//! The bundled native websocket backend for Rig: a
12//! [`rig_http::ws_client::WebSocketClientExt`] implementation over tokio-tungstenite. rig-core's `tungstenite` feature
13//! opens Responses WebSocket sessions over it with no backend named.
14//!
15//! The caller owns and polls every socket, in the current Tokio runtime's
16//! context or the lazy fallback's ([`rig_http::runtime`]), from any executor.
17//!
18//! ```
19//! let backend = rig_tungstenite::TungsteniteClient::new();
20//! ```
21
22// The native dependency is target-gated; diagnose unsupported WASM use before
23// unresolved imports obscure the required browser backend.
24#[cfg(target_family = "wasm")]
25compile_error!(
26    "rig-tungstenite is a native websocket backend (tokio-tungstenite). On wasm, implement \
27     `rig_http::ws_client::WebSocketClientExt` over `web_sys::WebSocket` and open sessions with \
28     `connect_with(..)`."
29);
30
31#[cfg(not(target_family = "wasm"))]
32pub use tokio_tungstenite;
33
34#[cfg(not(target_family = "wasm"))]
35mod connection;
36
37#[cfg(not(target_family = "wasm"))]
38use connection::Connection;
39#[cfg(not(target_family = "wasm"))]
40use rig_http::http_client::{Error, NoBody, Request, Result};
41#[cfg(not(target_family = "wasm"))]
42use rig_http::ws_client::{BoxedWebSocketConnection, ConnectOptions, WebSocketClientExt};
43#[cfg(not(target_family = "wasm"))]
44use std::time::Duration;
45#[cfg(not(target_family = "wasm"))]
46use tokio_tungstenite::tungstenite::{self, client::IntoClientRequest};
47
48#[cfg(not(target_family = "wasm"))]
49/// A stateless websocket backend that takes handshake configuration from each
50/// request.
51#[derive(Clone, Copy, Debug, Default)]
52pub struct TungsteniteClient;
53
54#[cfg(not(target_family = "wasm"))]
55impl TungsteniteClient {
56    /// The bundled backend.
57    #[must_use]
58    pub fn new() -> Self {
59        Self
60    }
61}
62
63#[cfg(not(target_family = "wasm"))]
64impl WebSocketClientExt for TungsteniteClient {
65    async fn connect(
66        &self,
67        request: Request<NoBody>,
68        options: ConnectOptions,
69    ) -> Result<BoxedWebSocketConnection> {
70        let request = client_request(request)?;
71        let socket =
72            rig_http::runtime::bind(Box::pin(handshake(request, options.timeout)))?.await?;
73        Ok(Box::new(Connection(socket)))
74    }
75}
76
77#[cfg(not(target_family = "wasm"))]
78/// Build a tungstenite handshake request, returning an error for an invalid URI.
79/// Caller headers, including authentication, override generated handshake headers.
80fn client_request(request: Request<NoBody>) -> Result<tungstenite::handshake::client::Request> {
81    let (parts, _) = request.into_parts();
82    let mut request = parts
83        .uri
84        .to_string()
85        .into_client_request()
86        .map_err(from_tungstenite)?;
87    for (name, value) in &parts.headers {
88        request.headers_mut().insert(name, value.clone());
89    }
90    Ok(request)
91}
92
93#[cfg(not(target_family = "wasm"))]
94/// A handshake that did not complete in time.
95#[derive(Debug, thiserror::Error)]
96#[error("timed out connecting the websocket after {0:?}")]
97struct ConnectTimeout(Duration);
98
99#[cfg(not(target_family = "wasm"))]
100async fn handshake(
101    request: tungstenite::handshake::client::Request,
102    timeout: Option<Duration>,
103) -> Result<connection::Socket> {
104    let connect = async {
105        tokio_tungstenite::connect_async(request)
106            .await
107            .map(|(socket, _)| socket)
108            .map_err(from_tungstenite)
109    };
110
111    let Some(timeout) = timeout else {
112        return connect.await;
113    };
114
115    match rig_http::wasm_compat::timeout(timeout, connect).await {
116        Ok(result) => result,
117        Err(_) => Err(Error::instance(ConnectTimeout(timeout))),
118    }
119}
120
121#[cfg(not(target_family = "wasm"))]
122/// Convert a tungstenite failure to a transport error, preserving the status,
123/// headers, and body of a rejected upgrade.
124/// Other failures become [`Error::Instance`].
125fn from_tungstenite(error: tungstenite::Error) -> Error {
126    let tungstenite::Error::Http(response) = error else {
127        return Error::instance(error);
128    };
129
130    let (parts, body) = (*response).into_parts();
131    let body = body
132        .map(|body| String::from_utf8_lossy(&body).into_owned())
133        .unwrap_or_default();
134
135    Error::non_success_with_details(parts.status, parts.headers, body)
136}
137
138#[cfg(all(test, not(target_family = "wasm")))]
139mod tests;
rig_tungstenite/lib.rs

rig_tungstenite/
lib.rs
