lean_ctx/proxy/
chatgpt_ws.rs

1//! WebSocket passthrough for ChatGPT's `/backend-api` rail (#597).
2//!
3//! When the Codex ChatGPT subscription opt-in is enabled, Codex's
4//! `chatgpt_base_url` points at the proxy, so *every* ChatGPT backend call —
5//! including Codex Desktop's **remote-control pairing**, which opens a
6//! WebSocket to chatgpt.com — flows through the proxy. The HTTP/SSE
7//! [`super::chatgpt::backend_api_handler`] cannot carry that: it strips the
8//! `Upgrade`/`Connection` headers and never speaks the WS protocol, so pairing
9//! never completed and remote control stayed broken.
10//!
11//! This module makes the proxy a transparent WebSocket tunnel for those calls:
12//! it accepts the client upgrade, opens an upstream `wss://chatgpt.com` socket
13//! (replaying the client's auth + the shared Cloudflare clearance), and relays
14//! every frame verbatim in both directions. The model-turn rail
15//! (`/backend-api/codex/responses`) keeps its own dedicated handlers and is
16//! never reached here.
17
18use axum::body::Body;
19use axum::extract::FromRequestParts;
20use axum::extract::ws::{
21    CloseFrame as AxumCloseFrame, Message as AxumMessage, WebSocket, WebSocketUpgrade,
22};
23use axum::http::{
24    HeaderMap, HeaderName, HeaderValue, Request, StatusCode, header, uri::PathAndQuery,
25};
26use axum::response::{IntoResponse, Response};
27use futures::{SinkExt, StreamExt};
28use tokio::net::TcpStream;
29use tokio_tungstenite::tungstenite::Message as TMessage;
30use tokio_tungstenite::tungstenite::client::IntoClientRequest;
31use tokio_tungstenite::tungstenite::protocol::CloseFrame as TCloseFrame;
32use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
33
34use super::ProxyState;
35
36/// True when `headers` describe a WebSocket upgrade (`Connection: Upgrade` +
37/// `Upgrade: websocket`, both case-insensitive). Lets the `/backend-api`
38/// handler branch to the tunnel without consuming the request body.
39pub(super) fn is_websocket_upgrade(headers: &HeaderMap) -> bool {
40    let connection_upgrade = headers
41        .get(header::CONNECTION)
42        .and_then(|v| v.to_str().ok())
43        .is_some_and(|v| {
44            v.split(',')
45                .any(|t| t.trim().eq_ignore_ascii_case("upgrade"))
46        });
47    let upgrade_websocket = headers
48        .get(header::UPGRADE)
49        .and_then(|v| v.to_str().ok())
50        .is_some_and(|v| v.eq_ignore_ascii_case("websocket"));
51    connection_upgrade && upgrade_websocket
52}
53
54/// Handshake headers tungstenite regenerates itself, plus the body framing
55/// headers a WS upgrade never carries. Everything else (auth, cookies,
56/// user-agent, the `x-openai-*`/`x-codex-*` identity set, subprotocol) is
57/// forwarded so chatgpt.com sees the same request Codex would have sent direct.
58fn is_handshake_header(name: &HeaderName) -> bool {
59    matches!(
60        name.as_str().to_ascii_lowercase().as_str(),
61        "host"
62            | "connection"
63            | "upgrade"
64            | "content-length"
65            | "content-type"
66            | "sec-websocket-key"
67            | "sec-websocket-version"
68            | "sec-websocket-accept"
69            | "sec-websocket-extensions"
70    )
71}
72
73fn capture_forward_headers(headers: &HeaderMap) -> Vec<(HeaderName, HeaderValue)> {
74    headers
75        .iter()
76        .filter(|(name, _)| !is_handshake_header(name))
77        .map(|(name, value)| (name.clone(), value.clone()))
78        .collect()
79}
80
81/// `https://host` → `wss://host{path}`, `http://host` → `ws://host{path}`.
82fn to_ws_url(upstream: &str, path: &str) -> Option<String> {
83    let base = upstream.trim_end_matches('/');
84    let ws_base = if let Some(rest) = base.strip_prefix("https://") {
85        format!("wss://{rest}")
86    } else {
87        let rest = base.strip_prefix("http://")?;
88        format!("ws://{rest}")
89    };
90    Some(format!("{ws_base}{path}"))
91}
92
93/// Merge the proxy's shared Cloudflare clearance into the upstream `Cookie`
94/// header so chatgpt.com does not bounce the handshake.
95fn merge_cookie(headers: &mut HeaderMap, cf_cookie: &str) {
96    let merged = match headers.get(header::COOKIE).and_then(|v| v.to_str().ok()) {
97        Some(existing) if !existing.trim().is_empty() => format!("{existing}; {cf_cookie}"),
98        _ => cf_cookie.to_string(),
99    };
100    if let Ok(value) = HeaderValue::from_str(&merged) {
101        headers.insert(header::COOKIE, value);
102    }
103}
104
105/// Accept the client WebSocket and tunnel it to chatgpt.com's `/backend-api`.
106/// Returns an error response only when the request is not a valid upgrade; the
107/// actual relay runs after the 101 on the upgraded connection.
108pub(super) async fn passthrough(state: ProxyState, req: Request<Body>) -> Response {
109    let (mut parts, _body) = req.into_parts();
110
111    let path = parts
112        .uri
113        .path_and_query()
114        .map_or("/backend-api", PathAndQuery::as_str)
115        .to_string();
116    let Some(ws_url) = to_ws_url(&state.chatgpt_upstream(), &path) else {
117        return (StatusCode::BAD_GATEWAY, "invalid ChatGPT upstream").into_response();
118    };
119
120    let forwarded = capture_forward_headers(&parts.headers);
121    let cf_cookie = state.chatgpt_cookie_header();
122
123    let ws = match WebSocketUpgrade::from_request_parts(&mut parts, &state).await {
124        Ok(ws) => ws,
125        Err(rejection) => return rejection.into_response(),
126    };
127
128    ws.on_upgrade(move |client| async move {
129        if let Err(err) = tunnel(client, ws_url, forwarded, cf_cookie).await {
130            tracing::warn!("lean-ctx proxy: ChatGPT WebSocket passthrough failed: {err}");
131        }
132    })
133}
134
135async fn tunnel(
136    client: WebSocket,
137    ws_url: String,
138    forwarded: Vec<(HeaderName, HeaderValue)>,
139    cf_cookie: Option<String>,
140) -> Result<(), tokio_tungstenite::tungstenite::Error> {
141    let mut request = ws_url.into_client_request()?;
142    {
143        let headers = request.headers_mut();
144        for (name, value) in forwarded {
145            headers.insert(name, value);
146        }
147        if let Some(cf) = cf_cookie {
148            merge_cookie(headers, &cf);
149        }
150    }
151
152    let (upstream, _response) = tokio_tungstenite::connect_async(request).await?;
153    relay(client, upstream).await;
154    Ok(())
155}
156
157async fn relay(client: WebSocket, upstream: WebSocketStream<MaybeTlsStream<TcpStream>>) {
158    let (mut client_tx, mut client_rx) = client.split();
159    let (mut upstream_tx, mut upstream_rx) = upstream.split();
160
161    let client_to_upstream = async {
162        while let Some(Ok(msg)) = client_rx.next().await {
163            let closing = matches!(msg, AxumMessage::Close(_));
164            if upstream_tx.send(axum_to_tungstenite(msg)).await.is_err() {
165                break;
166            }
167            if closing {
168                break;
169            }
170        }
171    };
172
173    let upstream_to_client = async {
174        while let Some(Ok(msg)) = upstream_rx.next().await {
175            let Some(msg) = tungstenite_to_axum(msg) else {
176                continue;
177            };
178            let closing = matches!(msg, AxumMessage::Close(_));
179            if client_tx.send(msg).await.is_err() {
180                break;
181            }
182            if closing {
183                break;
184            }
185        }
186    };
187
188    // Either side closing tears down the other: dropping the unfinished future
189    // releases its socket half, which closes the connection.
190    tokio::select! {
191        () = client_to_upstream => {},
192        () = upstream_to_client => {},
193    }
194}
195
196fn axum_to_tungstenite(msg: AxumMessage) -> TMessage {
197    match msg {
198        AxumMessage::Text(text) => TMessage::Text(text.as_str().into()),
199        AxumMessage::Binary(data) => TMessage::Binary(data),
200        AxumMessage::Ping(data) => TMessage::Ping(data),
201        AxumMessage::Pong(data) => TMessage::Pong(data),
202        AxumMessage::Close(None) => TMessage::Close(None),
203        AxumMessage::Close(Some(frame)) => TMessage::Close(Some(TCloseFrame {
204            code: frame.code.into(),
205            reason: frame.reason.as_str().into(),
206        })),
207    }
208}
209
210fn tungstenite_to_axum(msg: TMessage) -> Option<AxumMessage> {
211    match msg {
212        TMessage::Text(text) => Some(AxumMessage::Text(text.as_str().into())),
213        TMessage::Binary(data) => Some(AxumMessage::Binary(data)),
214        TMessage::Ping(data) => Some(AxumMessage::Ping(data)),
215        TMessage::Pong(data) => Some(AxumMessage::Pong(data)),
216        TMessage::Close(None) => Some(AxumMessage::Close(None)),
217        TMessage::Close(Some(frame)) => Some(AxumMessage::Close(Some(AxumCloseFrame {
218            code: frame.code.into(),
219            reason: frame.reason.as_str().into(),
220        }))),
221        // Raw frames never surface from a high-level `next()` read.
222        TMessage::Frame(_) => None,
223    }
224}
225
226#[cfg(test)]
227mod tests {
228    use super::*;
229
230    #[test]
231    fn detects_websocket_upgrade() {
232        let mut headers = HeaderMap::new();
233        headers.insert(header::CONNECTION, HeaderValue::from_static("Upgrade"));
234        headers.insert(header::UPGRADE, HeaderValue::from_static("websocket"));
235        assert!(is_websocket_upgrade(&headers));
236    }
237
238    #[test]
239    fn detects_websocket_upgrade_case_and_list_insensitive() {
240        let mut headers = HeaderMap::new();
241        headers.insert(
242            header::CONNECTION,
243            HeaderValue::from_static("keep-alive, Upgrade"),
244        );
245        headers.insert(header::UPGRADE, HeaderValue::from_static("WebSocket"));
246        assert!(is_websocket_upgrade(&headers));
247    }
248
249    #[test]
250    fn plain_request_is_not_an_upgrade() {
251        let mut headers = HeaderMap::new();
252        headers.insert(header::CONNECTION, HeaderValue::from_static("keep-alive"));
253        assert!(!is_websocket_upgrade(&headers));
254
255        let empty = HeaderMap::new();
256        assert!(!is_websocket_upgrade(&empty));
257    }
258
259    #[test]
260    fn to_ws_url_rewrites_scheme_and_keeps_path() {
261        assert_eq!(
262            to_ws_url("https://chatgpt.com", "/backend-api/wham/connect?x=1"),
263            Some("wss://chatgpt.com/backend-api/wham/connect?x=1".to_string())
264        );
265        assert_eq!(
266            to_ws_url("http://127.0.0.1:4444/", "/backend-api/ws"),
267            Some("ws://127.0.0.1:4444/backend-api/ws".to_string())
268        );
269        assert_eq!(to_ws_url("ftp://nope", "/x"), None);
270    }
271
272    #[test]
273    fn handshake_headers_are_dropped_but_auth_is_forwarded() {
274        let mut headers = HeaderMap::new();
275        headers.insert(header::HOST, HeaderValue::from_static("127.0.0.1:4444"));
276        headers.insert(header::CONNECTION, HeaderValue::from_static("Upgrade"));
277        headers.insert(header::UPGRADE, HeaderValue::from_static("websocket"));
278        headers.insert(
279            "sec-websocket-key",
280            HeaderValue::from_static("dGhlIHNhbXBsZQ=="),
281        );
282        headers.insert(
283            header::AUTHORIZATION,
284            HeaderValue::from_static("Bearer chatgpt-token"),
285        );
286        headers.insert("x-codex-installation-id", HeaderValue::from_static("abc"));
287
288        let forwarded = capture_forward_headers(&headers);
289        let names: Vec<String> = forwarded
290            .iter()
291            .map(|(n, _)| n.as_str().to_string())
292            .collect();
293
294        assert!(names.contains(&"authorization".to_string()));
295        assert!(names.contains(&"x-codex-installation-id".to_string()));
296        assert!(!names.iter().any(|n| n == "host"));
297        assert!(!names.iter().any(|n| n == "connection"));
298        assert!(!names.iter().any(|n| n == "upgrade"));
299        assert!(!names.iter().any(|n| n == "sec-websocket-key"));
300    }
301
302    #[test]
303    fn merge_cookie_appends_to_existing() {
304        let mut headers = HeaderMap::new();
305        headers.insert(header::COOKIE, HeaderValue::from_static("session=abc"));
306        merge_cookie(&mut headers, "cf_clearance=xyz");
307        assert_eq!(
308            headers.get(header::COOKIE).unwrap().to_str().unwrap(),
309            "session=abc; cf_clearance=xyz"
310        );
311    }
312
313    #[test]
314    fn merge_cookie_sets_when_absent() {
315        let mut headers = HeaderMap::new();
316        merge_cookie(&mut headers, "cf_clearance=xyz");
317        assert_eq!(
318            headers.get(header::COOKIE).unwrap().to_str().unwrap(),
319            "cf_clearance=xyz"
320        );
321    }
322
323    #[test]
324    fn message_conversion_round_trips() {
325        let original = AxumMessage::Text("ping".into());
326        let back = tungstenite_to_axum(axum_to_tungstenite(original)).unwrap();
327        assert!(matches!(back, AxumMessage::Text(t) if t.as_str() == "ping"));
328
329        let binary = AxumMessage::Binary(vec![1, 2, 3].into());
330        let back = tungstenite_to_axum(axum_to_tungstenite(binary)).unwrap();
331        assert!(matches!(back, AxumMessage::Binary(b) if b.as_ref() == [1, 2, 3]));
332    }
333
334    /// End-to-end: a WebSocket client → proxy `/backend-api` → upstream echo
335    /// server. Proves the handshake is tunnelled and frames relay both ways,
336    /// which is exactly what Codex Desktop remote-control pairing needs (#597).
337    #[tokio::test]
338    async fn tunnels_websocket_through_backend_api_to_upstream() {
339        use std::sync::Arc;
340        use std::time::Duration;
341
342        use axum::Router;
343        use axum::routing::any;
344        use tokio::net::TcpListener;
345        use tokio_tungstenite::tungstenite::Message;
346
347        // Upstream echo WS server, addressed exactly like chatgpt.com would be.
348        let upstream_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
349        let upstream_addr = upstream_listener.local_addr().unwrap();
350        tokio::spawn(async move {
351            while let Ok((stream, _)) = upstream_listener.accept().await {
352                tokio::spawn(async move {
353                    let mut ws = tokio_tungstenite::accept_async(stream).await.unwrap();
354                    while let Some(Ok(msg)) = ws.next().await {
355                        match msg {
356                            Message::Text(_) | Message::Binary(_) => {
357                                if ws.send(msg).await.is_err() {
358                                    break;
359                                }
360                            }
361                            Message::Close(_) => break,
362                            _ => {}
363                        }
364                    }
365                });
366            }
367        });
368
369        // Proxy app: just the `/backend-api` rail, pointed at the echo upstream.
370        let (_tx, rx) = tokio::sync::watch::channel(Arc::new(crate::core::config::Upstreams {
371            anthropic: "https://api.anthropic.com".into(),
372            openai: "https://api.openai.com".into(),
373            chatgpt: format!("http://{upstream_addr}"),
374            gemini: "https://generativelanguage.googleapis.com".into(),
375            providers: Vec::new(),
376        }));
377        let state = ProxyState {
378            client: reqwest::Client::new(),
379            port: 0,
380            stats: Arc::new(crate::proxy::ProxyStats::default()),
381            break_even: Arc::new(crate::proxy::break_even::BreakEvenCalculator::new(1500)),
382            introspect: Arc::new(crate::proxy::introspect::IntrospectState::default()),
383            ocla_cache: None,
384            upstreams: rx,
385            chatgpt_cookies: crate::proxy::chatgpt_cookies::shared_chatgpt_cloudflare_cookie_store(
386            ),
387            mcp_servers: Arc::new(Vec::new()),
388            web_app_tracker: Arc::new(std::sync::Mutex::new(
389                crate::proxy::web_app::conversation_tracker::ConversationTracker::default(),
390            )),
391        };
392        let app = Router::new()
393            .route(
394                "/backend-api/{*rest}",
395                any(crate::proxy::chatgpt::backend_api_handler),
396            )
397            .with_state(state);
398        let proxy_listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
399        let proxy_addr = proxy_listener.local_addr().unwrap();
400        tokio::spawn(async move {
401            axum::serve(proxy_listener, app).await.unwrap();
402        });
403
404        // Client connects to the proxy and round-trips a frame each way.
405        let url = format!("ws://{proxy_addr}/backend-api/wham/connect");
406        let (mut client, _resp) = tokio::time::timeout(
407            Duration::from_secs(3),
408            tokio_tungstenite::connect_async(url),
409        )
410        .await
411        .expect("handshake must complete")
412        .expect("proxy must tunnel the upgrade to the upstream");
413
414        client
415            .send(Message::Text("remote-control".into()))
416            .await
417            .unwrap();
418        let echoed = tokio::time::timeout(Duration::from_secs(3), client.next())
419            .await
420            .expect("echo must arrive")
421            .expect("stream open")
422            .expect("valid frame");
423        assert_eq!(echoed, Message::Text("remote-control".into()));
424
425        let binary = Message::Binary(vec![9, 8, 7].into());
426        client.send(binary.clone()).await.unwrap();
427        let echoed = tokio::time::timeout(Duration::from_secs(3), client.next())
428            .await
429            .expect("binary echo must arrive")
430            .expect("stream open")
431            .expect("valid frame");
432        assert_eq!(echoed, binary);
433    }
434}
lean_ctx/proxy/chatgpt_ws.rs

lean_ctx/proxy/
chatgpt_ws.rs
