diff --git a/src/backend.rs b/src/backend.rs index eacdf6c..c270b27 100644 --- a/src/backend.rs +++ b/src/backend.rs @@ -5,18 +5,19 @@ use crate::{ use codex_api::{ ApiError, Compression, ReqwestTransport, ResponsesApiRequest, ResponsesClient, ResponsesOptions, }; -use codex_http_client::{ClientRouteClass, HttpClientFactory}; +use codex_http_client::{ClientRouteClass, HttpClient, HttpClientFactory}; use codex_login::{AuthManager, default_client::create_client_for_route}; use codex_model_provider::SharedModelProvider; use codex_protocol::protocol::SessionSource; use futures::StreamExt; -use std::sync::{Arc, OnceLock}; +use std::sync::{Arc, Mutex, OnceLock}; use tokio::sync::{mpsc, oneshot}; pub struct Backend { pub provider: SharedModelProvider, pub auth: Arc, pub factory: HttpClientFactory, + pub responses_http: Mutex>, pub chatgpt_base_url: String, pub platform_base_url: String, pub auth_base_url: String, @@ -339,13 +340,26 @@ impl Backend { if *auth_changes.borrow() != revision { continue; } - let http = create_client_for_route( - &self.factory, - &resolved.url_for_path("/responses"), - ClientRouteClass::Api, - codex_login::default_client::ClientRedirectPolicy::Default, - ) - .map_err(GatewayError::internal)?; + let url = resolved.url_for_path("/responses"); + let http = { + let mut cached = self + .responses_http + .lock() + .map_err(|_| GatewayError::internal("responses HTTP client lock poisoned"))?; + if let Some((_, http)) = cached.as_ref().filter(|(u, _)| u == &url) { + http.clone() + } else { + let http = create_client_for_route( + &self.factory, + &url, + ClientRouteClass::Api, + codex_login::default_client::ClientRedirectPolicy::Default, + ) + .map_err(GatewayError::internal)?; + *cached = Some((url, http.clone())); + http + } + }; let (sender, receiver) = mpsc::channel(8); let response_headers = Arc::new(OnceLock::new()); let transport = TapTransport { diff --git a/src/settings.rs b/src/settings.rs index 9c9a6ab..61be429 100644 --- a/src/settings.rs +++ b/src/settings.rs @@ -109,6 +109,7 @@ async fn load_backend_inner(codex_home: PathBuf, device_login: bool) -> anyhow:: provider, auth, factory, + responses_http: Default::default(), chatgpt_base_url, platform_base_url: "https://api.openai.com/v1".into(), auth_base_url: "https://auth.openai.com".into(), diff --git a/tests/catalog.rs b/tests/catalog.rs index e50623a..a1b808d 100644 --- a/tests/catalog.rs +++ b/tests/catalog.rs @@ -23,6 +23,7 @@ async fn backend(app: Router) -> (Backend, tokio::task::JoinHandle<()>) { ), auth, factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + responses_http: Default::default(), chatgpt_base_url: url.clone(), platform_base_url: url.clone(), auth_base_url: url, diff --git a/tests/gateway.rs b/tests/gateway.rs index 7ba381f..eac9be0 100644 --- a/tests/gateway.rs +++ b/tests/gateway.rs @@ -303,6 +303,7 @@ impl Harness { ), auth, factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + responses_http: Default::default(), chatgpt_base_url: upstream_url.clone(), platform_base_url: upstream_url.clone(), auth_base_url: upstream_url.clone(), @@ -378,6 +379,32 @@ fn parse_sse(text: &str) -> Vec { .collect() } +#[tokio::test] +async fn responses_reuse_upstream_connection() { + let expected = events(); + let h = Harness::new( + expected.clone(), + StatusCode::OK, + false, + Duration::from_secs(5), + ) + .await; + let client = reqwest::Client::new(); + for path in ["/v1/responses", "/codex/responses", "/v1/responses"] { + let response = client + .post(format!("{}{path}", h.url)) + .bearer_auth("client-key") + .json(&json!({"model":h.model,"input":"hello","stream":true})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(parse_sse(&response.text().await.unwrap()), expected); + } + assert_eq!(h.fake.received.lock().unwrap().len(), 3); + assert_eq!(h.capture.connection_count(), 1); +} + #[tokio::test] async fn stream_preserves_full_events_and_uses_codex_request_and_headers() { let expected = events(); diff --git a/tests/support/wire.rs b/tests/support/wire.rs index ecbb907..81dac7b 100644 --- a/tests/support/wire.rs +++ b/tests/support/wire.rs @@ -20,6 +20,10 @@ impl Drop for Capture { } } impl Capture { + pub fn connection_count(&self) -> usize { + self.connections.lock().unwrap().len() + } + pub async fn start(destination: String) -> Self { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let url = format!("http://{}", listener.local_addr().unwrap());