lean_ctx/proxy/
chatgpt.rs1use axum::{
2 body::Body,
3 extract::State,
4 http::{HeaderName, Request, StatusCode},
5 response::Response,
6};
7
8use super::{ProxyState, forward, openai_responses};
9
10pub async fn codex_responses_handler(
14 State(state): State<ProxyState>,
15 mut req: Request<Body>,
16) -> Result<Response, StatusCode> {
17 req.headers_mut().remove(CODEX_RESPONSES_LITE_HEADER);
25 let upstream = state.chatgpt_upstream();
26 forward::forward_request(
27 State(state),
28 req,
29 &upstream,
30 "/backend-api/codex/responses",
31 openai_responses::compress_request_body,
32 "ChatGPT",
33 &[],
34 )
35 .await
36}
37
38const CODEX_RESPONSES_LITE_HEADER: &str = "x-openai-internal-codex-responses-lite";
42
43pub async fn codex_responses_ws_handler(
46 State(_state): State<ProxyState>,
47 _headers: axum::http::HeaderMap,
48 _ws: axum::extract::ws::WebSocketUpgrade,
49) -> Response {
50 chatgpt_responses_ws_fallback_response()
51}
52
53fn chatgpt_responses_ws_fallback_response() -> Response {
54 Response::builder()
55 .status(StatusCode::UPGRADE_REQUIRED)
56 .header("content-type", "application/json")
57 .body(Body::from(
58 r#"{"error":{"type":"unsupported_transport","message":"ChatGPT codex responses use HTTP/SSE; retry without WebSocket."}}"#,
59 ))
60 .expect("static response is valid")
61}
62
63pub async fn backend_api_handler(
66 State(state): State<ProxyState>,
67 req: Request<Body>,
68) -> Result<Response, StatusCode> {
69 let (parts, body) = req.into_parts();
70 let body_bytes = axum::body::to_bytes(body, forward::max_body_bytes())
71 .await
72 .map_err(|_| StatusCode::PAYLOAD_TOO_LARGE)?;
73 let upstream = state.chatgpt_upstream();
74 let path = parts
75 .uri
76 .path_and_query()
77 .map_or("/backend-api", axum::http::uri::PathAndQuery::as_str);
78 let url = format!("{upstream}{path}");
79
80 let mut upstream_req = state.client.request(parts.method.clone(), &url);
81 for (key, value) in &parts.headers {
82 if is_backend_passthrough_request_header(key) {
83 upstream_req = upstream_req.header(key.clone(), value.clone());
84 }
85 }
86
87 let response = upstream_req
88 .body(body_bytes.to_vec())
89 .send()
90 .await
91 .map_err(|e| {
92 tracing::error!("lean-ctx proxy: ChatGPT backend upstream error: {e}");
93 StatusCode::BAD_GATEWAY
94 })?;
95
96 let status = StatusCode::from_u16(response.status().as_u16()).unwrap_or(StatusCode::OK);
97 let headers = response.headers().clone();
98 let is_stream = headers
99 .get("content-type")
100 .and_then(|v| v.to_str().ok())
101 .is_some_and(|ct| ct.contains("text/event-stream"));
102
103 let mut out = Response::builder().status(status);
104 for (key, value) in &headers {
105 if is_backend_passthrough_response_header(key) {
106 out = out.header(key, value);
107 }
108 }
109
110 if is_stream {
111 return out
112 .body(Body::from_stream(response.bytes_stream()))
113 .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR);
114 }
115
116 let bytes = response
117 .bytes()
118 .await
119 .map_err(|_| StatusCode::BAD_GATEWAY)?;
120
121 out.body(Body::from(bytes))
122 .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)
123}
124
125fn is_backend_passthrough_request_header(name: &HeaderName) -> bool {
126 let lower = name.as_str().to_ascii_lowercase();
127 !matches!(
128 lower.as_str(),
129 "host"
130 | "connection"
131 | "content-length"
132 | "transfer-encoding"
133 | "upgrade"
134 | "keep-alive"
135 | "proxy-authenticate"
136 | "proxy-authorization"
137 | "te"
138 | "trailer"
139 | "accept-encoding"
140 )
141}
142
143fn is_backend_passthrough_response_header(name: &HeaderName) -> bool {
144 let lower = name.as_str().to_ascii_lowercase();
145 !matches!(
146 lower.as_str(),
147 "connection"
148 | "content-length"
149 | "transfer-encoding"
150 | "upgrade"
151 | "keep-alive"
152 | "proxy-authenticate"
153 | "proxy-authorization"
154 | "te"
155 | "trailer"
156 )
157}
158
159#[cfg(test)]
160mod tests {
161 use std::sync::Arc;
162 use std::time::Duration;
163
164 use tokio::io::{AsyncReadExt, AsyncWriteExt};
165
166 use super::*;
167 use crate::core::config::Upstreams;
168
169 fn proxy_state(chatgpt_upstream: String) -> ProxyState {
170 let (_tx, rx) = tokio::sync::watch::channel(Arc::new(Upstreams {
171 anthropic: "https://api.anthropic.com".into(),
172 openai: "https://api.openai.com".into(),
173 chatgpt: chatgpt_upstream,
174 gemini: "https://generativelanguage.googleapis.com".into(),
175 }));
176 ProxyState {
177 client: reqwest::Client::new(),
178 port: 0,
179 stats: Arc::new(crate::proxy::ProxyStats::default()),
180 introspect: Arc::new(crate::proxy::introspect::IntrospectState::default()),
181 upstreams: rx,
182 }
183 }
184
185 #[test]
186 fn codex_responses_ws_requests_trigger_http_fallback() {
187 let response = chatgpt_responses_ws_fallback_response();
188 assert_eq!(response.status(), StatusCode::UPGRADE_REQUIRED);
189 assert_eq!(
190 response
191 .headers()
192 .get(axum::http::header::CONTENT_TYPE)
193 .unwrap(),
194 "application/json"
195 );
196 }
197
198 async fn spawn_streaming_upstream() -> (String, tokio::sync::oneshot::Receiver<String>) {
199 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
200 let addr = listener.local_addr().unwrap();
201 let (tx, rx) = tokio::sync::oneshot::channel();
202 tokio::spawn(async move {
203 let (mut socket, _) = listener.accept().await.unwrap();
204 let mut buf = Vec::new();
205 loop {
206 let mut chunk = [0_u8; 1024];
207 let n = socket.read(&mut chunk).await.unwrap();
208 if n == 0 {
209 break;
210 }
211 buf.extend_from_slice(&chunk[..n]);
212 if buf.windows(4).any(|w| w == b"\r\n\r\n") {
213 break;
214 }
215 }
216 let _ = tx.send(String::from_utf8_lossy(&buf).into_owned());
217 socket
218 .write_all(
219 b"HTTP/1.1 200 OK\r\n\
220 content-type: text/event-stream\r\n\
221 mcp-session-id: server-session\r\n\
222 cache-control: no-cache\r\n\
223 x-custom-backend-state: passthrough\r\n\
224 \r\n\
225 event: message\n\
226 data: {\"jsonrpc\":\"2.0\"}\n\n",
227 )
228 .await
229 .unwrap();
230 tokio::time::sleep(Duration::from_secs(2)).await;
231 });
232 (format!("http://{addr}"), rx)
233 }
234
235 #[tokio::test]
236 async fn backend_api_streams_mcp_sse_and_preserves_session_headers() {
237 let (upstream, seen_request) = spawn_streaming_upstream().await;
238 let state = proxy_state(upstream);
239 let req = Request::builder()
240 .method("POST")
241 .uri("/backend-api/ps/mcp?transport=streamable")
242 .header("Authorization", "Bearer codex-token")
243 .header("Mcp-Session-Id", "client-session")
244 .header("Last-Event-ID", "event-7")
245 .header("X-OpenAI-Product-Sku", "codex")
246 .header("X-OpenAI-Internal-Codex-Residency", "us")
247 .header("Originator", "codex_cli_rs")
248 .header("Accept", "application/json, text/event-stream")
249 .body(Body::empty())
250 .unwrap();
251
252 let response = tokio::time::timeout(
253 Duration::from_millis(500),
254 backend_api_handler(State(state), req),
255 )
256 .await
257 .expect("SSE passthrough must return after upstream headers")
258 .expect("backend request should succeed");
259
260 assert_eq!(response.status(), StatusCode::OK);
261 assert_eq!(
262 response.headers().get("mcp-session-id").unwrap(),
263 "server-session"
264 );
265 assert_eq!(
266 response.headers().get("x-custom-backend-state").unwrap(),
267 "passthrough"
268 );
269
270 let request = seen_request.await.unwrap().to_ascii_lowercase();
271 assert!(request.contains("post /backend-api/ps/mcp?transport=streamable http/1.1"));
272 assert!(request.contains("authorization: bearer codex-token"));
273 assert!(request.contains("mcp-session-id: client-session"));
274 assert!(request.contains("last-event-id: event-7"));
275 assert!(request.contains("x-openai-product-sku: codex"));
276 assert!(request.contains("x-openai-internal-codex-residency: us"));
277 assert!(request.contains("originator: codex_cli_rs"));
278 }
279}