http-client: add safe route-aware request pool

codex-rs/http-client/src/default_client.rs

@@ -1,4 +1,4 @@
-use http::Error as HttpError;
+use http::Error as HttpRequestBuildError;
 use http::HeaderMap;
 use http::HeaderName;
 use http::HeaderValue;
@@ -6,13 +6,15 @@ use opentelemetry::global;
 use opentelemetry::propagation::Injector;
 use reqwest::IntoUrl;
 use reqwest::Method;
-use reqwest::Response;
 use serde::Serialize;
 use std::fmt::Display;
 use std::time::Duration;
 use tracing::Span;
 use tracing_opentelemetry::OpenTelemetrySpanExt;
 
+pub type HttpError = reqwest::Error;
+pub type HttpResponse = reqwest::Response;
+
 #[derive(Clone, Debug)]
 pub struct HttpClient {
     inner: reqwest::Client,
@@ -64,6 +66,43 @@ impl HttpClient {
             self.request_logging,
         )
     }
+
+    pub(crate) async fn execute(
+        &self,
+        mut request: reqwest::Request,
+    ) -> Result<reqwest::Response, reqwest::Error> {
+        request.headers_mut().extend(trace_headers());
+        let method = request.method().clone();
+        let url = request.url().to_string();
+
+        match self.inner.execute(request).await {
+            Ok(response) => {
+                if self.request_logging == RequestLogging::Enabled {
+                    tracing::debug!(
+                        method = %method,
+                        url = %url,
+                        status = %response.status(),
+                        headers = ?response.headers(),
+                        version = ?response.version(),
+                        "Request completed"
+                    );
+                }
+                Ok(response)
+            }
+            Err(error) => {
+                if self.request_logging == RequestLogging::Enabled {
+                    tracing::debug!(
+                        method = %method,
+                        url = %url,
+                        status = error.status().map(|status| status.as_u16()),
+                        error = %error,
+                        "Request failed"
+                    );
+                }
+                Err(error)
+            }
+        }
+    }
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -112,9 +151,9 @@ impl RequestBuilder {
     pub fn header<K, V>(self, key: K, value: V) -> Self
     where
         HeaderName: TryFrom<K>,
-        <HeaderName as TryFrom<K>>::Error: Into<HttpError>,
+        <HeaderName as TryFrom<K>>::Error: Into<HttpRequestBuildError>,
         HeaderValue: TryFrom<V>,
-        <HeaderValue as TryFrom<V>>::Error: Into<HttpError>,
+        <HeaderValue as TryFrom<V>>::Error: Into<HttpRequestBuildError>,
     {
         self.map(|builder| builder.header(key, value))
     }
@@ -144,7 +183,7 @@ impl RequestBuilder {
         self.map(|builder| builder.body(body))
     }
 
-    pub async fn send(self) -> Result<Response, reqwest::Error> {
+    pub async fn send(self) -> Result<HttpResponse, HttpError> {
         let headers = trace_headers();
 
         match self.builder.headers(headers).send().await {
@@ -192,7 +231,7 @@ impl<'a> Injector for HeaderMapInjector<'a> {
     }
 }
 
-fn trace_headers() -> HeaderMap {
+pub(crate) fn trace_headers() -> HeaderMap {
     let mut headers = HeaderMap::new();
     global::get_text_map_propagator(|prop| {
         prop.inject_context(

codex-rs/http-client/src/lib.rs

@@ -5,6 +5,7 @@ mod default_client;
 mod error;
 mod outbound_proxy;
 mod request;
+mod route_aware_client_pool;
 mod transport;
 
 pub use crate::chatgpt_cloudflare_cookies::with_chatgpt_cloudflare_cookie_store;
@@ -21,6 +22,8 @@ pub use crate::custom_ca::build_reqwest_client_with_custom_ca;
 pub use crate::custom_ca::build_rustls_client_config_with_custom_ca;
 pub use crate::custom_ca::maybe_build_rustls_client_config_with_custom_ca;
 pub use crate::default_client::HttpClient;
+pub use crate::default_client::HttpError;
+pub use crate::default_client::HttpResponse;
 pub use crate::default_client::RequestBuilder;
 pub use crate::error::StreamError;
 pub use crate::error::TransportError;
@@ -36,6 +39,10 @@ pub use crate::request::Request;
 pub use crate::request::RequestBody;
 pub use crate::request::RequestCompression;
 pub use crate::request::Response;
+pub use crate::route_aware_client_pool::RouteAwareClientPool;
+pub use crate::route_aware_client_pool::RouteAwareClientPoolError;
+pub use crate::route_aware_client_pool::RouteAwareRequestBuilder;
+pub use crate::route_aware_client_pool::RouteAwareRequestError;
 pub use crate::transport::ByteStream;
 pub use crate::transport::HttpTransport;
 pub use crate::transport::ReqwestTransport;

codex-rs/http-client/src/outbound_proxy.rs

@@ -253,6 +253,16 @@ impl HttpClientFactory {
             self.outbound_proxy_policy,
         )
     }
+
+    pub(crate) fn build_reqwest_client_for_resolved_route(
+        &self,
+        builder: reqwest::ClientBuilder,
+        route_class: ClientRouteClass,
+        route: &OutboundProxyRoute,
+    ) -> Result<reqwest::Client, BuildRouteAwareHttpClientError> {
+        let builder = configure_builder_for_resolved_route(builder, route_class, route)?;
+        build_reqwest_client_with_custom_ca(builder).map_err(Into::into)
+    }
 }
 
 fn resolve_proxy_route(

codex-rs/http-client/src/route_aware_client_pool.rs

@@ -0,0 +1,399 @@
+use std::collections::HashMap;
+use std::fmt;
+use std::future::Future;
+use std::io;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::time::Duration;
+
+use http::HeaderMap;
+use http::HeaderName;
+use http::HeaderValue;
+use http::Method;
+use http::StatusCode;
+use http::header::AUTHORIZATION;
+use http::header::CONTENT_TYPE;
+use reqwest::IntoUrl;
+use serde::Serialize;
+
+use crate::BuildRouteAwareHttpClientError;
+use crate::ClientRouteClass;
+use crate::HttpClient;
+use crate::HttpClientFactory;
+use crate::OutboundProxyPolicy;
+use crate::OutboundProxyRoute;
+use crate::with_chatgpt_cloudflare_cookie_store;
+
+const MAX_CACHED_ROUTES: usize = 16;
+
+/// Reuses transport clients by resolved route while selecting a route for every request URL.
+///
+/// Request creation stays on the pool so the URL used for PAC or system-proxy resolution cannot
+/// differ from the URL that is sent.
+#[derive(Clone)]
+pub struct RouteAwareClientPool {
+    http_client_factory: HttpClientFactory,
+    route_class: ClientRouteClass,
+    builder_factory: Arc<dyn Fn() -> reqwest::ClientBuilder + Send + Sync>,
+    request_logging: PoolRequestLogging,
+    clients: Arc<Mutex<HashMap<OutboundProxyRoute, HttpClient>>>,
+}
+
+impl fmt::Debug for RouteAwareClientPool {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter
+            .debug_struct("RouteAwareClientPool")
+            .field("http_client_factory", &self.http_client_factory)
+            .field("route_class", &self.route_class)
+            .field("request_logging", &self.request_logging)
+            .finish_non_exhaustive()
+    }
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+enum PoolRequestLogging {
+    Enabled,
+    Disabled,
+}
+
+/// Error returned when selecting a route or constructing its pooled HTTP client.
+#[derive(Debug, thiserror::Error)]
+pub enum RouteAwareClientPoolError {
+    #[error("failed to resolve the outbound proxy route: {0}")]
+    Resolve(#[source] io::Error),
+    #[error(transparent)]
+    Build(#[from] BuildRouteAwareHttpClientError),
+}
+
+/// Error returned while building, routing, or sending a route-aware request.
+#[derive(Debug, thiserror::Error)]
+pub enum RouteAwareRequestError {
+    #[error(transparent)]
+    Request(#[from] reqwest::Error),
+    #[error(transparent)]
+    Route(#[from] RouteAwareClientPoolError),
+    #[error("failed to build route-aware request: {0}")]
+    Build(String),
+}
+
+impl RouteAwareRequestError {
+    pub fn status(&self) -> Option<StatusCode> {
+        match self {
+            Self::Request(error) => error.status(),
+            Self::Route(_) | Self::Build(_) => None,
+        }
+    }
+
+    pub fn is_timeout(&self) -> bool {
+        matches!(self, Self::Request(error) if error.is_timeout())
+    }
+
+    pub fn is_connect(&self) -> bool {
+        matches!(self, Self::Request(error) if error.is_connect())
+    }
+}
+
+#[must_use = "requests are not sent unless `send` is awaited"]
+pub struct RouteAwareRequestBuilder {
+    pool: RouteAwareClientPool,
+    request: Result<reqwest::Request, RouteAwareRequestError>,
+}
+
+impl fmt::Debug for RouteAwareRequestBuilder {
+    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+        formatter
+            .debug_struct("RouteAwareRequestBuilder")
+            .field("pool", &self.pool)
+            .field(
+                "url",
+                &self.request.as_ref().ok().map(reqwest::Request::url),
+            )
+            .finish_non_exhaustive()
+    }
+}
+
+impl RouteAwareRequestBuilder {
+    fn new<U>(pool: RouteAwareClientPool, method: Method, url: U) -> Self
+    where
+        U: IntoUrl,
+    {
+        let request = url
+            .into_url()
+            .map(|url| reqwest::Request::new(method, url))
+            .map_err(RouteAwareRequestError::Request);
+        Self { pool, request }
+    }
+
+    pub fn headers(mut self, headers: HeaderMap) -> Self {
+        if let Ok(request) = &mut self.request {
+            request.headers_mut().extend(headers);
+        }
+        self
+    }
+
+    pub fn header<K, V>(mut self, key: K, value: V) -> Self
+    where
+        HeaderName: TryFrom<K>,
+        <HeaderName as TryFrom<K>>::Error: Into<http::Error>,
+        HeaderValue: TryFrom<V>,
+        <HeaderValue as TryFrom<V>>::Error: Into<http::Error>,
+    {
+        if let Ok(request) = &mut self.request {
+            let header = HeaderName::try_from(key)
+                .map_err(Into::into)
+                .and_then(|key| {
+                    HeaderValue::try_from(value)
+                        .map(|value| (key, value))
+                        .map_err(Into::into)
+                });
+            match header {
+                Ok((key, value)) => {
+                    request.headers_mut().append(key, value);
+                }
+                Err(error) => {
+                    self.request = Err(RouteAwareRequestError::Build(error.to_string()));
+                }
+            }
+        }
+        self
+    }
+
+    pub fn bearer_auth<T>(mut self, token: T) -> Self
+    where
+        T: fmt::Display,
+    {
+        let value = HeaderValue::from_str(&format!("Bearer {token}"));
+        match (&mut self.request, value) {
+            (Ok(request), Ok(mut value)) => {
+                value.set_sensitive(true);
+                request.headers_mut().append(AUTHORIZATION, value);
+            }
+            (Ok(_), Err(error)) => {
+                self.request = Err(RouteAwareRequestError::Build(error.to_string()));
+            }
+            (Err(_), _) => {}
+        }
+        self
+    }
+
+    pub fn timeout(mut self, timeout: Duration) -> Self {
+        if let Ok(request) = &mut self.request {
+            *request.timeout_mut() = Some(timeout);
+        }
+        self
+    }
+
+    pub fn json<T>(mut self, value: &T) -> Self
+    where
+        T: ?Sized + Serialize,
+    {
+        if let Ok(request) = &mut self.request {
+            match serde_json::to_vec(value) {
+                Ok(body) => {
+                    if !request.headers().contains_key(CONTENT_TYPE) {
+                        request
+                            .headers_mut()
+                            .insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));
+                    }
+                    *request.body_mut() = Some(body.into());
+                }
+                Err(error) => {
+                    self.request = Err(RouteAwareRequestError::Build(error.to_string()));
+                }
+            }
+        }
+        self
+    }
+
+    pub fn body<B>(mut self, body: B) -> Self
+    where
+        B: Into<reqwest::Body>,
+    {
+        if let Ok(request) = &mut self.request {
+            *request.body_mut() = Some(body.into());
+        }
+        self
+    }
+
+    pub async fn send(self) -> Result<reqwest::Response, RouteAwareRequestError> {
+        self.pool.send(self.request?).await
+    }
+}
+
+impl RouteAwareClientPool {
+    pub fn outbound_proxy_policy(&self) -> OutboundProxyPolicy {
+        self.http_client_factory.outbound_proxy_policy()
+    }
+
+    /// Creates a pool with the shared default HTTP transport settings.
+    pub fn new(http_client_factory: HttpClientFactory, route_class: ClientRouteClass) -> Self {
+        Self::with_builder_factory(
+            http_client_factory,
+            route_class,
+            reqwest::Client::builder,
+            PoolRequestLogging::Enabled,
+        )
+    }
+
+    /// Creates a pool with the shared defaults but without URL or response-header diagnostics.
+    pub fn new_without_request_logging(
+        http_client_factory: HttpClientFactory,
+        route_class: ClientRouteClass,
+    ) -> Self {
+        Self::with_builder_factory(
+            http_client_factory,
+            route_class,
+            reqwest::Client::builder,
+            PoolRequestLogging::Disabled,
+        )
+    }
+
+    /// Creates a pool that retains the Cloudflare cookies required by ChatGPT endpoints.
+    pub fn with_chatgpt_cloudflare_cookies(
+        http_client_factory: HttpClientFactory,
+        route_class: ClientRouteClass,
+    ) -> Self {
+        Self::with_builder_factory(
+            http_client_factory,
+            route_class,
+            || with_chatgpt_cloudflare_cookie_store(reqwest::Client::builder()),
+            PoolRequestLogging::Enabled,
+        )
+    }
+
+    /// Creates a ChatGPT Cloudflare-cookie pool without URL or response-header diagnostics.
+    pub fn with_chatgpt_cloudflare_cookies_without_request_logging(
+        http_client_factory: HttpClientFactory,
+        route_class: ClientRouteClass,
+    ) -> Self {
+        Self::with_builder_factory(
+            http_client_factory,
+            route_class,
+            || with_chatgpt_cloudflare_cookie_store(reqwest::Client::builder()),
+            PoolRequestLogging::Disabled,
+        )
+    }
+
+    pub fn get<U>(&self, url: U) -> RouteAwareRequestBuilder
+    where
+        U: IntoUrl,
+    {
+        self.request(Method::GET, url)
+    }
+
+    pub fn post<U>(&self, url: U) -> RouteAwareRequestBuilder
+    where
+        U: IntoUrl,
+    {
+        self.request(Method::POST, url)
+    }
+
+    pub fn put<U>(&self, url: U) -> RouteAwareRequestBuilder
+    where
+        U: IntoUrl,
+    {
+        self.request(Method::PUT, url)
+    }
+
+    pub fn delete<U>(&self, url: U) -> RouteAwareRequestBuilder
+    where
+        U: IntoUrl,
+    {
+        self.request(Method::DELETE, url)
+    }
+
+    pub fn request<U>(&self, method: Method, url: U) -> RouteAwareRequestBuilder
+    where
+        U: IntoUrl,
+    {
+        RouteAwareRequestBuilder::new(self.clone(), method, url)
+    }
+
+    fn with_builder_factory(
+        http_client_factory: HttpClientFactory,
+        route_class: ClientRouteClass,
+        builder_factory: impl Fn() -> reqwest::ClientBuilder + Send + Sync + 'static,
+        request_logging: PoolRequestLogging,
+    ) -> Self {
+        Self {
+            http_client_factory,
+            route_class,
+            builder_factory: Arc::new(builder_factory),
+            request_logging,
+            clients: Arc::new(Mutex::new(HashMap::new())),
+        }
+    }
+
+    async fn send(
+        &self,
+        request: reqwest::Request,
+    ) -> Result<reqwest::Response, RouteAwareRequestError> {
+        let client = self.client_for_url(request.url().as_str()).await?;
+        Ok(client.execute(request).await?)
+    }
+
+    async fn client_for_url(
+        &self,
+        request_url: &str,
+    ) -> Result<HttpClient, RouteAwareClientPoolError> {
+        let http_client_factory = self.http_client_factory.clone();
+        self.client_for_url_with_resolver(request_url, move |request_url| async move {
+            http_client_factory
+                .resolve_proxy_route_async(request_url)
+                .await
+        })
+        .await
+    }
+
+    async fn client_for_url_with_resolver<F, Fut>(
+        &self,
+        request_url: &str,
+        resolve_route: F,
+    ) -> Result<HttpClient, RouteAwareClientPoolError>
+    where
+        F: FnOnce(String) -> Fut,
+        Fut: Future<Output = io::Result<OutboundProxyRoute>>,
+    {
+        let route = resolve_route(request_url.to_string())
+            .await
+            .map_err(RouteAwareClientPoolError::Resolve)?;
+        let clients = match self.clients.lock() {
+            Ok(clients) => clients,
+            Err(error) => panic!("route-aware client cache lock should not be poisoned: {error}"),
+        };
+        if let Some(client) = clients.get(&route) {
+            return Ok(client.clone());
+        }
+        drop(clients);
+
+        let client = self
+            .http_client_factory
+            .build_reqwest_client_for_resolved_route(
+                (self.builder_factory)().redirect(reqwest::redirect::Policy::none()),
+                self.route_class,
+                &route,
+            )?;
+        let client = match self.request_logging {
+            PoolRequestLogging::Enabled => HttpClient::new(client),
+            PoolRequestLogging::Disabled => HttpClient::new_without_request_logging(client),
+        };
+        let mut clients = match self.clients.lock() {
+            Ok(clients) => clients,
+            Err(error) => panic!("route-aware client cache lock should not be poisoned: {error}"),
+        };
+        if let Some(existing_client) = clients.get(&route) {
+            return Ok(existing_client.clone());
+        }
+        if clients.len() >= MAX_CACHED_ROUTES
+            && let Some(route_to_evict) = clients.keys().next().cloned()
+        {
+            clients.remove(&route_to_evict);
+        }
+        clients.insert(route, client.clone());
+        Ok(client)
+    }
+}
+
+#[cfg(test)]
+#[path = "route_aware_client_pool_tests.rs"]
+mod tests;

codex-rs/http-client/src/route_aware_client_pool_tests.rs

@@ -0,0 +1,158 @@
+use std::collections::HashMap;
+use std::io;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+
+use pretty_assertions::assert_eq;
+
+use super::*;
+use crate::OutboundProxyPolicy;
+
+#[tokio::test]
+async fn forwards_exact_urls_and_reuses_clients_by_resolved_route() {
+    let builder_count = Arc::new(AtomicUsize::new(0));
+    let observed_builder_count = Arc::clone(&builder_count);
+    let pool = RouteAwareClientPool::with_builder_factory(
+        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+        ClientRouteClass::Api,
+        move || {
+            observed_builder_count.fetch_add(1, Ordering::SeqCst);
+            reqwest::Client::builder()
+        },
+        PoolRequestLogging::Enabled,
+    );
+
+    let direct_url = "https://example.com/first?target=direct";
+    let same_route_url = "https://example.com/second?target=direct%202";
+    let proxy_url = "https://example.com/third?target=proxy";
+    let resolver = FakeRouteResolver::new(HashMap::from([
+        (direct_url.to_string(), OutboundProxyRoute::Direct),
+        (same_route_url.to_string(), OutboundProxyRoute::Direct),
+        (
+            proxy_url.to_string(),
+            OutboundProxyRoute::Proxy {
+                url: "http://proxy.example".to_string(),
+                no_proxy: None,
+            },
+        ),
+    ]));
+
+    resolve_with(&pool, &resolver, direct_url)
+        .await
+        .expect("first client should build");
+    resolve_with(&pool, &resolver, same_route_url)
+        .await
+        .expect("second client should reuse the route");
+    resolve_with(&pool, &resolver, proxy_url)
+        .await
+        .expect("proxy client should build separately");
+
+    assert_eq!(builder_count.load(Ordering::SeqCst), 2);
+    assert_eq!(
+        resolver.observed_urls(),
+        vec![
+            direct_url.to_string(),
+            same_route_url.to_string(),
+            proxy_url.to_string(),
+        ]
+    );
+}
+
+#[tokio::test]
+async fn bounds_cached_routes_and_rebuilds_an_evicted_route() {
+    let builder_count = Arc::new(AtomicUsize::new(0));
+    let observed_builder_count = Arc::clone(&builder_count);
+    let pool = RouteAwareClientPool::with_builder_factory(
+        HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
+        ClientRouteClass::Api,
+        move || {
+            observed_builder_count.fetch_add(1, Ordering::SeqCst);
+            reqwest::Client::builder()
+        },
+        PoolRequestLogging::Enabled,
+    );
+    let routes = (0..=MAX_CACHED_ROUTES)
+        .map(|index| {
+            (
+                format!("https://target-{index}.example"),
+                OutboundProxyRoute::Proxy {
+                    url: format!("http://proxy-{index}.example"),
+                    no_proxy: None,
+                },
+            )
+        })
+        .collect::<HashMap<_, _>>();
+    let resolver = FakeRouteResolver::new(routes.clone());
+
+    for request_url in routes.keys() {
+        resolve_with(&pool, &resolver, request_url)
+            .await
+            .expect("client should build");
+    }
+    let evicted_route = {
+        let clients = pool.clients.lock().expect("client cache lock");
+        assert_eq!(clients.len(), MAX_CACHED_ROUTES);
+        routes
+            .iter()
+            .find(|(_, route)| !clients.contains_key(*route))
+            .map(|(request_url, _)| request_url.clone())
+            .expect("one route should have been evicted")
+    };
+
+    resolve_with(&pool, &resolver, &evicted_route)
+        .await
+        .expect("evicted client should rebuild");
+
+    assert_eq!(builder_count.load(Ordering::SeqCst), MAX_CACHED_ROUTES + 2);
+    assert_eq!(
+        pool.clients.lock().expect("client cache lock").len(),
+        MAX_CACHED_ROUTES
+    );
+}
+
+#[derive(Clone)]
+struct FakeRouteResolver {
+    routes: Arc<HashMap<String, OutboundProxyRoute>>,
+    observed_urls: Arc<Mutex<Vec<String>>>,
+}
+
+impl FakeRouteResolver {
+    fn new(routes: HashMap<String, OutboundProxyRoute>) -> Self {
+        Self {
+            routes: Arc::new(routes),
+            observed_urls: Arc::new(Mutex::new(Vec::new())),
+        }
+    }
+
+    async fn resolve(&self, request_url: String) -> io::Result<OutboundProxyRoute> {
+        self.observed_urls
+            .lock()
+            .expect("observed URL lock")
+            .push(request_url.clone());
+        self.routes
+            .get(&request_url)
+            .cloned()
+            .ok_or_else(|| io::Error::other(format!("no route for {request_url}")))
+    }
+
+    fn observed_urls(&self) -> Vec<String> {
+        self.observed_urls
+            .lock()
+            .expect("observed URL lock")
+            .clone()
+    }
+}
+
+async fn resolve_with(
+    pool: &RouteAwareClientPool,
+    resolver: &FakeRouteResolver,
+    request_url: &str,
+) -> Result<HttpClient, RouteAwareClientPoolError> {
+    let resolver = resolver.clone();
+    pool.client_for_url_with_resolver(request_url, move |request_url| async move {
+        resolver.resolve(request_url).await
+    })
+    .await
+}