mcp_proxy/
ws_transport.rs

1//! WebSocket client transport for connecting to WebSocket-based MCP backends.
2//!
3//! Implements tower-mcp's [`ClientTransport`](tower_mcp::client::ClientTransport) trait over a WebSocket connection
4//! using `tokio-tungstenite`. Messages are sent and received as text frames
5//! containing JSON-RPC payloads.
6//!
7//! # Example
8//!
9//! ```rust,no_run
10//! use mcp_proxy::ws_transport::WebSocketClientTransport;
11//!
12//! # async fn example() -> anyhow::Result<()> {
13//! let transport = WebSocketClientTransport::connect("ws://localhost:8080/ws").await?;
14//! // Pass to McpProxy::builder().backend("name", transport).await
15//! # Ok(())
16//! # }
17//! ```
18
19use std::sync::Arc;
20use std::sync::atomic::{AtomicBool, Ordering};
21
22use async_trait::async_trait;
23use futures_util::{SinkExt, StreamExt};
24use tokio::sync::Mutex;
25use tokio_tungstenite::tungstenite::Message;
26
27type WsStream =
28    tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>;
29
30/// A WebSocket client transport for MCP backend connections.
31///
32/// Connects to a WebSocket endpoint and exchanges JSON-RPC messages
33/// as text frames. Supports `ws://` and `wss://` (TLS) URLs.
34pub struct WebSocketClientTransport {
35    sink: Arc<Mutex<futures_util::stream::SplitSink<WsStream, Message>>>,
36    stream: Arc<Mutex<futures_util::stream::SplitStream<WsStream>>>,
37    connected: Arc<AtomicBool>,
38}
39
40impl WebSocketClientTransport {
41    /// Connect to a WebSocket endpoint.
42    ///
43    /// # Errors
44    ///
45    /// Returns an error if the WebSocket handshake fails or the URL is invalid.
46    pub async fn connect(url: &str) -> anyhow::Result<Self> {
47        let (ws_stream, _response) = tokio_tungstenite::connect_async(url)
48            .await
49            .map_err(|e| anyhow::anyhow!("WebSocket connection failed: {e}"))?;
50
51        let (sink, stream) = ws_stream.split();
52
53        Ok(Self {
54            sink: Arc::new(Mutex::new(sink)),
55            stream: Arc::new(Mutex::new(stream)),
56            connected: Arc::new(AtomicBool::new(true)),
57        })
58    }
59
60    /// Connect to a WebSocket endpoint with a bearer token for authentication.
61    ///
62    /// The token is sent in the `Authorization` header during the handshake.
63    pub async fn connect_with_bearer_token(url: &str, token: &str) -> anyhow::Result<Self> {
64        Self::connect_with_headers(url, Some(token), &Default::default()).await
65    }
66
67    /// Connect with custom handshake headers. An explicit Authorization header
68    /// overrides the optional bearer token, regardless of header name casing.
69    pub async fn connect_with_headers(
70        url: &str,
71        token: Option<&str>,
72        headers: &std::collections::HashMap<String, String>,
73    ) -> anyhow::Result<Self> {
74        use tokio_tungstenite::tungstenite::client::IntoClientRequest;
75        use tokio_tungstenite::tungstenite::http::{HeaderName, HeaderValue};
76        let mut request = url.into_client_request()?;
77        if let Some(token) = token {
78            request.headers_mut().insert(
79                "authorization",
80                HeaderValue::try_from(format!("Bearer {token}"))?,
81            );
82        }
83        for (name, value) in headers {
84            request.headers_mut().insert(
85                HeaderName::try_from(name.as_str())?,
86                HeaderValue::try_from(value.as_str())?,
87            );
88        }
89
90        let (ws_stream, _response) = tokio_tungstenite::connect_async(request)
91            .await
92            .map_err(|e| anyhow::anyhow!("WebSocket connection failed: {e}"))?;
93
94        let (sink, stream) = ws_stream.split();
95
96        Ok(Self {
97            sink: Arc::new(Mutex::new(sink)),
98            stream: Arc::new(Mutex::new(stream)),
99            connected: Arc::new(AtomicBool::new(true)),
100        })
101    }
102}
103
104#[async_trait]
105impl tower_mcp::client::ClientTransport for WebSocketClientTransport {
106    async fn send(&mut self, message: &str) -> tower_mcp::error::Result<()> {
107        let mut sink = self.sink.lock().await;
108        sink.send(Message::Text(message.into()))
109            .await
110            .map_err(|e| tower_mcp::error::Error::Transport(e.to_string()))?;
111        Ok(())
112    }
113
114    async fn recv(&mut self) -> tower_mcp::error::Result<Option<String>> {
115        let mut stream = self.stream.lock().await;
116        loop {
117            match stream.next().await {
118                Some(Ok(Message::Text(text))) => return Ok(Some(text.as_str().to_owned())),
119                Some(Ok(Message::Close(_))) | None => {
120                    self.connected.store(false, Ordering::SeqCst);
121                    return Ok(None);
122                }
123                Some(Ok(Message::Ping(_) | Message::Pong(_) | Message::Frame(_))) => {
124                    // Pong is handled automatically by tungstenite; skip control frames
125                    continue;
126                }
127                Some(Ok(Message::Binary(data))) => {
128                    // Try to interpret binary as UTF-8 text
129                    let text = std::str::from_utf8(&data)
130                        .map_err(|e| tower_mcp::error::Error::Transport(e.to_string()))?;
131                    return Ok(Some(text.to_string()));
132                }
133                Some(Err(e)) => {
134                    self.connected.store(false, Ordering::SeqCst);
135                    return Err(tower_mcp::error::Error::Transport(e.to_string()));
136                }
137            }
138        }
139    }
140
141    fn is_connected(&self) -> bool {
142        self.connected.load(Ordering::SeqCst)
143    }
144
145    async fn close(&mut self) -> tower_mcp::error::Result<()> {
146        self.connected.store(false, Ordering::SeqCst);
147        let mut sink = self.sink.lock().await;
148        let _ = sink.send(Message::Close(None)).await;
149        Ok(())
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    #[tokio::test]
158    async fn connect_fails_with_invalid_url() {
159        let result = WebSocketClientTransport::connect("ws://127.0.0.1:1").await;
160        let err = result.err().expect("should fail").to_string();
161        assert!(
162            err.contains("WebSocket connection failed"),
163            "unexpected error: {err}"
164        );
165    }
166
167    #[tokio::test]
168    async fn connect_with_bearer_token_fails_with_invalid_url() {
169        let result =
170            WebSocketClientTransport::connect_with_bearer_token("ws://127.0.0.1:1", "tok").await;
171        let err = result.err().expect("should fail").to_string();
172        assert!(
173            err.contains("WebSocket connection failed"),
174            "unexpected error: {err}"
175        );
176    }
177}
mcp_proxy/ws_transport.rs

mcp_proxy/
ws_transport.rs
