codex_websocket_client/
lib.rs

1//! Proxy-aware WebSocket connection setup shared by Codex API clients.
2
3mod dialer;
4
5use std::pin::Pin;
6use std::sync::Arc;
7use std::task::Context;
8use std::task::Poll;
9
10use codex_http_client::BuildCustomCaTransportError;
11use codex_http_client::HttpClientFactory;
12use codex_http_client::build_rustls_client_config_with_custom_ca;
13use futures::Sink;
14use futures::Stream;
15use rustls::ClientConfig;
16use tokio::io::AsyncRead;
17use tokio::io::AsyncWrite;
18use tokio::net::TcpStream;
19use tokio_tungstenite::MaybeTlsStream;
20use tokio_tungstenite::WebSocketStream as TungsteniteStream;
21use tokio_tungstenite::tungstenite::Error as WebSocketError;
22use tokio_tungstenite::tungstenite::Message;
23use tokio_tungstenite::tungstenite::handshake::client::Request;
24use tokio_tungstenite::tungstenite::handshake::client::Response;
25use tokio_tungstenite::tungstenite::protocol::WebSocketConfig;
26
27/// Connects WebSockets using the outbound proxy policy resolved by application configuration.
28///
29/// Construct this from the effective [`HttpClientFactory`] rather than selecting proxy behavior at
30/// individual call sites. Each connection resolves its destination through that factory before
31/// opening a socket.
32#[derive(Clone)]
33pub struct WebSocketConnector {
34    http_client_factory: HttpClientFactory,
35    tls_config: Arc<ClientConfig>,
36}
37
38impl WebSocketConnector {
39    /// Creates a connector using native roots and any configured Codex custom CA bundle.
40    pub fn new(
41        http_client_factory: &HttpClientFactory,
42    ) -> Result<Self, BuildCustomCaTransportError> {
43        Ok(Self {
44            http_client_factory: http_client_factory.clone(),
45            tls_config: build_rustls_client_config_with_custom_ca()?,
46        })
47    }
48
49    /// Connects a WebSocket after resolving the request destination through the configured proxy
50    /// policy.
51    pub async fn connect(
52        &self,
53        request: Request,
54        config: WebSocketConfig,
55    ) -> Result<(WebSocketConnection, Response), WebSocketError> {
56        let proxy_route = self
57            .http_client_factory
58            .resolve_proxy_route(&request.uri().to_string());
59        let (inner, response) =
60            dialer::connect(request, config, Arc::clone(&self.tls_config), proxy_route).await?;
61        Ok((WebSocketConnection { inner }, response))
62    }
63}
64
65/// An established WebSocket independent of its direct, proxy, and TLS transport layers.
66///
67/// This implements [`Stream`] and [`Sink`] so protocol clients can process Tungstenite messages
68/// without knowing which concrete network stream route selection produced.
69pub struct WebSocketConnection {
70    inner: ConnectionInner,
71}
72
73impl Stream for WebSocketConnection {
74    type Item = Result<Message, WebSocketError>;
75
76    fn poll_next(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
77        match &mut self.get_mut().inner {
78            ConnectionInner::TransportDefault(stream) => Pin::new(stream).poll_next(context),
79            ConnectionInner::Routed(stream) => Pin::new(stream).poll_next(context),
80        }
81    }
82}
83
84impl Sink<Message> for WebSocketConnection {
85    type Error = WebSocketError;
86
87    fn poll_ready(
88        self: Pin<&mut Self>,
89        context: &mut Context<'_>,
90    ) -> Poll<Result<(), Self::Error>> {
91        match &mut self.get_mut().inner {
92            ConnectionInner::TransportDefault(stream) => Pin::new(stream).poll_ready(context),
93            ConnectionInner::Routed(stream) => Pin::new(stream).poll_ready(context),
94        }
95    }
96
97    fn start_send(self: Pin<&mut Self>, message: Message) -> Result<(), Self::Error> {
98        match &mut self.get_mut().inner {
99            ConnectionInner::TransportDefault(stream) => Pin::new(stream).start_send(message),
100            ConnectionInner::Routed(stream) => Pin::new(stream).start_send(message),
101        }
102    }
103
104    fn poll_flush(
105        self: Pin<&mut Self>,
106        context: &mut Context<'_>,
107    ) -> Poll<Result<(), Self::Error>> {
108        match &mut self.get_mut().inner {
109            ConnectionInner::TransportDefault(stream) => Pin::new(stream).poll_flush(context),
110            ConnectionInner::Routed(stream) => Pin::new(stream).poll_flush(context),
111        }
112    }
113
114    fn poll_close(
115        self: Pin<&mut Self>,
116        context: &mut Context<'_>,
117    ) -> Poll<Result<(), Self::Error>> {
118        match &mut self.get_mut().inner {
119            ConnectionInner::TransportDefault(stream) => Pin::new(stream).poll_close(context),
120            ConnectionInner::Routed(stream) => Pin::new(stream).poll_close(context),
121        }
122    }
123}
124
125pub(crate) enum ConnectionInner {
126    TransportDefault(TungsteniteStream<MaybeTlsStream<TcpStream>>),
127    Routed(TungsteniteStream<MaybeTlsStream<Box<dyn AsyncIo>>>),
128}
129
130/// Async network I/O carried through optional proxy and target TLS handshakes.
131pub(crate) trait AsyncIo: AsyncRead + AsyncWrite + Send + Unpin {}
132
133impl<T> AsyncIo for T where T: AsyncRead + AsyncWrite + Send + Unpin {}
codex_websocket_client/lib.rs

codex_websocket_client/
lib.rs
