Skip to main content

limitless/ws/
client.rs

1use crate::prelude::*;
2
3/// Manages a WebSocket connection lifecycle.
4///
5/// Wraps a `WebSocketStream` and provides connect/disconnect with
6/// proper WebSocket Close frame semantics.
7pub struct WsClient {
8    stream: WebSocketStream<MaybeTlsStream<TcpStream>>,
9}
10
11impl WsClient {
12    /// Create a new `WsClient` from an already-established `WebSocketStream`.
13    pub fn new(stream: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
14        Self { stream }
15    }
16
17    /// Get a mutable reference to the underlying stream.
18    pub fn stream(&mut self) -> &mut WebSocketStream<MaybeTlsStream<TcpStream>> {
19        &mut self.stream
20    }
21
22    /// Send a WebSocket Close frame and consume the connection.
23    pub async fn disconnect(&mut self) -> Result<(), LimitlessError> {
24        self.stream
25            .close(None)
26            .await
27            .map_err(|e| LimitlessError::Base(format!("Error closing WebSocket: {}", e)))?;
28        Ok(())
29    }
30}