Route LM Studio requests through the shared HTTP client (#34678)

## What changed

- Use the configured route-aware HTTP client pool for LM Studio server requests.
- Add connection-timeout support to `HttpClientBuilder` and route-aware pools, and keep LM Studio's five-second limit scoped to connection establishment.
- Verify that LM Studio accepts a response that arrives after the connection timeout has elapsed once the connection is established.

GitOrigin-RevId: c4300f4b5d37c4418822783ab09cb50d506ee423
This commit is contained in:
Celia Chen
2026-07-22 04:12:23 +00:00
committed by copyberry
parent 4f3852107e
commit 21db216db0
7 changed files with 161 additions and 118 deletions

View File

@@ -6,12 +6,14 @@
//! paths.
use http::HeaderMap;
use std::time::Duration;
use crate::BuildCustomCaTransportError;
use crate::BuildRouteAwareHttpClientError;
use crate::ClientRouteClass;
use crate::HttpClient;
use crate::HttpClientFactory;
use crate::OutboundProxyRoute;
use crate::client::RequestLogging;
use crate::custom_ca::build_reqwest_client_with_custom_ca;
use crate::with_chatgpt_cloudflare_cookie_store;
@@ -25,6 +27,7 @@ use crate::with_chatgpt_cloudflare_cookie_store;
pub struct HttpClientBuilder {
default_headers: Option<HeaderMap>,
follow_redirects: bool,
connect_timeout: Option<Duration>,
chatgpt_cloudflare_cookie_store: bool,
request_logging: RequestLogging,
}
@@ -75,6 +78,12 @@ impl HttpClientBuilder {
self
}
/// Limits only connection establishment, not the request as a whole.
pub fn connect_timeout(mut self, timeout: Duration) -> Self {
self.connect_timeout = Some(timeout);
self
}
pub fn with_chatgpt_cloudflare_cookie_store(mut self) -> Self {
self.chatgpt_cloudflare_cookie_store = true;
self
@@ -102,6 +111,22 @@ impl HttpClientBuilder {
Ok(HttpClient::from_parts(inner, request_logging))
}
/// Builds a client for a route that was already resolved by a route-aware caller.
pub(crate) fn build_for_resolved_route(
self,
http_client_factory: &HttpClientFactory,
route_class: ClientRouteClass,
route: &OutboundProxyRoute,
) -> Result<HttpClient, BuildRouteAwareHttpClientError> {
let (builder, request_logging) = self.into_reqwest_parts();
let inner = http_client_factory.build_reqwest_client_for_resolved_route(
builder,
route_class,
route,
)?;
Ok(HttpClient::from_parts(inner, request_logging))
}
/// Builds a client using the transport's default proxy behavior.
///
/// # Legacy compatibility only
@@ -219,6 +244,9 @@ impl HttpClientBuilder {
if !self.follow_redirects {
builder = builder.redirect(reqwest::redirect::Policy::none());
}
if let Some(connect_timeout) = self.connect_timeout {
builder = builder.connect_timeout(connect_timeout);
}
if self.chatgpt_cloudflare_cookie_store {
builder = with_chatgpt_cloudflare_cookie_store(builder);
}
@@ -231,6 +259,7 @@ impl Default for HttpClientBuilder {
Self {
default_headers: None,
follow_redirects: true,
connect_timeout: None,
chatgpt_cloudflare_cookie_store: false,
request_logging: RequestLogging::Enabled,
}

View File

@@ -19,6 +19,7 @@ use serde::Serialize;
use crate::BuildRouteAwareHttpClientError;
use crate::ClientRouteClass;
use crate::HttpClient;
use crate::HttpClientBuilder;
use crate::HttpClientFactory;
use crate::OutboundProxyPolicy;
use crate::OutboundProxyRoute;
@@ -28,7 +29,6 @@ use crate::route_aware_redirect::is_redirect;
use crate::route_aware_redirect::redirect_request;
use crate::route_aware_redirect::redirect_url;
use crate::route_aware_redirect::remove_sensitive_headers;
use crate::with_chatgpt_cloudflare_cookie_store;
const MAX_CACHED_ROUTES: usize = 16;
@@ -41,8 +41,7 @@ const MAX_CACHED_ROUTES: usize = 16;
pub struct RouteAwareClientPool {
http_client_factory: HttpClientFactory,
route_class: ClientRouteClass,
builder_factory: Arc<dyn Fn() -> reqwest::ClientBuilder + Send + Sync>,
request_logging: PoolRequestLogging,
client_builder: HttpClientBuilder,
clients: Arc<Mutex<HashMap<OutboundProxyRoute, HttpClient>>>,
}
@@ -52,17 +51,10 @@ impl fmt::Debug for RouteAwareClientPool {
.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 {
@@ -174,6 +166,12 @@ impl RouteAwareRequestBuilder {
self
}
/// Sets a timeout for the request as a whole.
///
/// The budget starts before outbound-route resolution and covers selecting or constructing a
/// pooled client, establishing a connection, sending the request, and awaiting the response.
/// Use [`HttpClientBuilder::connect_timeout`] when only connection establishment should be
/// bounded.
pub fn timeout(mut self, timeout: Duration) -> Self {
if let Ok(request) = &mut self.request {
*request.timeout_mut() = Some(timeout);
@@ -225,24 +223,46 @@ impl RouteAwareClientPool {
/// 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(
Self::with_builder(http_client_factory, route_class, HttpClientBuilder::new())
}
/// Creates a pool whose clients limit only connection establishment.
///
/// The timeout applies to every client built for a resolved route, including redirect hops.
pub fn with_connect_timeout(
http_client_factory: HttpClientFactory,
route_class: ClientRouteClass,
connect_timeout: Duration,
) -> Self {
Self::with_builder(
http_client_factory,
route_class,
reqwest::Client::builder,
PoolRequestLogging::Enabled,
HttpClientBuilder::new().connect_timeout(connect_timeout),
)
}
fn with_builder(
http_client_factory: HttpClientFactory,
route_class: ClientRouteClass,
client_builder: HttpClientBuilder,
) -> Self {
Self {
http_client_factory,
route_class,
client_builder,
clients: Arc::new(Mutex::new(HashMap::new())),
}
}
/// 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(
Self::with_builder(
http_client_factory,
route_class,
reqwest::Client::builder,
PoolRequestLogging::Disabled,
HttpClientBuilder::new().without_request_logging(),
)
}
@@ -251,11 +271,10 @@ impl RouteAwareClientPool {
http_client_factory: HttpClientFactory,
route_class: ClientRouteClass,
) -> Self {
Self::with_builder_factory(
Self::with_builder(
http_client_factory,
route_class,
|| with_chatgpt_cloudflare_cookie_store(reqwest::Client::builder()),
PoolRequestLogging::Enabled,
HttpClientBuilder::new().with_chatgpt_cloudflare_cookie_store(),
)
}
@@ -264,11 +283,12 @@ impl RouteAwareClientPool {
http_client_factory: HttpClientFactory,
route_class: ClientRouteClass,
) -> Self {
Self::with_builder_factory(
Self::with_builder(
http_client_factory,
route_class,
|| with_chatgpt_cloudflare_cookie_store(reqwest::Client::builder()),
PoolRequestLogging::Disabled,
HttpClientBuilder::new()
.with_chatgpt_cloudflare_cookie_store()
.without_request_logging(),
)
}
@@ -307,21 +327,6 @@ impl RouteAwareClientPool {
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,
@@ -473,20 +478,17 @@ impl RouteAwareClientPool {
}
drop(clients);
let builder = (self.builder_factory)();
let builder = match self.http_client_factory.outbound_proxy_policy() {
OutboundProxyPolicy::ReqwestDefault => builder,
let client_builder = match self.http_client_factory.outbound_proxy_policy() {
OutboundProxyPolicy::ReqwestDefault => self.client_builder.clone(),
OutboundProxyPolicy::RespectSystemProxy => {
builder.redirect(reqwest::redirect::Policy::none())
self.client_builder.clone().without_redirects()
}
};
let client = self
.http_client_factory
.build_reqwest_client_for_resolved_route(builder, self.route_class, &route)?;
let client = match self.request_logging {
PoolRequestLogging::Enabled => HttpClient::new(client),
PoolRequestLogging::Disabled => HttpClient::new_without_request_logging(client),
};
let client = client_builder.build_for_resolved_route(
&self.http_client_factory,
self.route_class,
&route,
)?;
let mut clients = match self.clients.lock() {
Ok(clients) => clients,
Err(error) => panic!("route-aware client cache lock should not be poisoned: {error}"),

View File

@@ -32,24 +32,18 @@ fn request_builder_debug_redacts_url_secrets() {
concat!(
"RouteAwareRequestBuilder { pool: RouteAwareClientPool { ",
"http_client_factory: HttpClientFactory { outbound_proxy_policy: ReqwestDefault }, ",
"route_class: Api, request_logging: Enabled, .. }, method: Some(GET), ",
"route_class: Api, .. }, method: Some(GET), ",
"url: Some(\"<redacted>\"), .. }"
)
);
}
#[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(
async fn forwards_exact_urls_and_caches_clients_by_resolved_route() {
let pool = RouteAwareClientPool::with_builder(
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
ClientRouteClass::Api,
move || {
observed_builder_count.fetch_add(1, Ordering::SeqCst);
reqwest::Client::builder()
},
PoolRequestLogging::Enabled,
HttpClientBuilder::new(),
);
let direct_url = "https://example.com/first?target=direct";
@@ -77,7 +71,7 @@ async fn forwards_exact_urls_and_reuses_clients_by_resolved_route() {
.await
.expect("proxy client should build separately");
assert_eq!(builder_count.load(Ordering::SeqCst), 2);
assert_eq!(pool.clients.lock().expect("client cache lock").len(), 2);
assert_eq!(
resolver.observed_urls(),
vec![
@@ -135,17 +129,19 @@ async fn reqwest_default_route_preserves_transport_redirects() {
}
request_lines
});
let pool = RouteAwareClientPool::with_builder_factory(
let pool = RouteAwareClientPool::with_builder(
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
ClientRouteClass::Api,
|| reqwest::Client::builder().no_proxy(),
PoolRequestLogging::Enabled,
HttpClientBuilder::new(),
);
let initial_url = format!("http://{address}/start");
let request = reqwest::Request::new(
Method::GET,
reqwest::Url::parse(&initial_url).expect("request URL should parse"),
);
let response = pool
.get(initial_url)
.send()
.send_with_resolver(request, |_| async { Ok(OutboundProxyRoute::Direct) })
.await
.expect("default-routed request should follow redirect");
@@ -162,16 +158,10 @@ async fn reqwest_default_route_preserves_transport_redirects() {
#[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(
let pool = RouteAwareClientPool::with_builder(
HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault),
ClientRouteClass::Api,
move || {
observed_builder_count.fetch_add(1, Ordering::SeqCst);
reqwest::Client::builder()
},
PoolRequestLogging::Enabled,
HttpClientBuilder::new(),
);
let routes = (0..=MAX_CACHED_ROUTES)
.map(|index| {
@@ -205,11 +195,9 @@ async fn bounds_cached_routes_and_rebuilds_an_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
);
let clients = pool.clients.lock().expect("client cache lock");
assert_eq!(clients.len(), MAX_CACHED_ROUTES);
assert!(clients.contains_key(&routes[&evicted_route]));
}
#[tokio::test]
@@ -342,11 +330,10 @@ async fn disabled_pool_logging_does_not_expose_request_or_response_data() {
"HTTP/1.1 200 OK\r\nx-sensitive-response: response-secret-value\r\nContent-Length: 2\r\nConnection: close\r\n\r\nok"
.to_string(),
]);
let pool = RouteAwareClientPool::with_builder_factory(
let pool = RouteAwareClientPool::with_builder(
HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy),
ClientRouteClass::Api,
|| reqwest::Client::builder().no_proxy(),
PoolRequestLogging::Disabled,
HttpClientBuilder::new().without_request_logging(),
);
let buffer = Arc::new(Mutex::new(Vec::new()));
let subscriber = tracing_subscriber::registry().with(
@@ -467,11 +454,10 @@ async fn resolve_with(
}
fn manual_redirect_pool() -> RouteAwareClientPool {
RouteAwareClientPool::with_builder_factory(
RouteAwareClientPool::with_builder(
HttpClientFactory::new(OutboundProxyPolicy::RespectSystemProxy),
ClientRouteClass::Api,
|| reqwest::Client::builder().no_proxy(),
PoolRequestLogging::Enabled,
HttpClientBuilder::new(),
)
}