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

2
codex-rs/Cargo.lock generated
View File

@@ -3308,8 +3308,8 @@ name = "codex-lmstudio"
version = "0.0.0"
dependencies = [
"codex-core",
"codex-http-client",
"codex-model-provider-info",
"reqwest 0.12.28",
"serde_json",
"tokio",
"tracing",

View File

@@ -243,7 +243,6 @@ deny = [
"codex-app-server",
"codex-core",
"codex-exec-server",
"codex-lmstudio",
"codex-ollama",
"codex-otel",
"codex-protocol",

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(),
)
}

View File

@@ -12,8 +12,8 @@ doctest = false
[dependencies]
codex-core = { path = "../core" }
codex-http-client = { workspace = true }
codex-model-provider-info = { path = "../model-provider-info" }
reqwest = { version = "0.12", features = ["json", "stream"] }
serde_json = "1"
tokio = { version = "1", features = ["rt"] }
tracing = { version = "0.1.44", features = ["log"] }

View File

@@ -1,15 +1,19 @@
use codex_core::config::Config;
use codex_http_client::ClientRouteClass;
use codex_http_client::RouteAwareClientPool;
use codex_model_provider_info::LMSTUDIO_OSS_PROVIDER_ID;
use std::io;
use std::path::Path;
use std::time::Duration;
#[derive(Clone)]
pub struct LMStudioClient {
client: reqwest::Client,
client: RouteAwareClientPool,
base_url: String,
}
const LMSTUDIO_CONNECTION_ERROR: &str = "LM Studio is not responding. Install from https://lmstudio.ai/download and run 'lms server start'.";
const LMSTUDIO_CONNECTION_TIMEOUT: Duration = Duration::from_secs(5);
impl LMStudioClient {
pub async fn try_from_provider(config: &Config) -> std::io::Result<Self> {
@@ -29,10 +33,11 @@ impl LMStudioClient {
)
})?;
let client = reqwest::Client::builder()
.connect_timeout(std::time::Duration::from_secs(5))
.build()
.unwrap_or_else(|_| reqwest::Client::new());
let client = RouteAwareClientPool::with_connect_timeout(
config.http_client_factory(),
ClientRouteClass::Other,
LMSTUDIO_CONNECTION_TIMEOUT,
);
let client = LMStudioClient {
client,
@@ -188,19 +193,6 @@ impl LMStudioClient {
tracing::info!("Successfully downloaded model '{model}'");
Ok(())
}
/// Low-level constructor given a raw host root, e.g. "http://localhost:1234".
#[cfg(test)]
fn from_host_root(host_root: impl Into<String>) -> Self {
let client = reqwest::Client::builder()
.connect_timeout(std::time::Duration::from_secs(5))
.build()
.unwrap_or_else(|_| reqwest::Client::new());
Self {
client,
base_url: host_root.into(),
}
}
}
#[cfg(test)]
@@ -208,6 +200,23 @@ mod tests {
#![allow(clippy::expect_used, clippy::unwrap_used)]
use super::*;
fn client_from_host_root(
host_root: impl Into<String>,
connection_timeout: Duration,
) -> LMStudioClient {
let client = RouteAwareClientPool::with_connect_timeout(
codex_http_client::HttpClientFactory::new(
codex_http_client::OutboundProxyPolicy::ReqwestDefault,
),
ClientRouteClass::Other,
connection_timeout,
);
LMStudioClient {
client,
base_url: host_root.into(),
}
}
#[tokio::test]
async fn test_fetch_models_happy_path() {
if std::env::var(codex_core::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() {
@@ -235,7 +244,7 @@ mod tests {
.mount(&server)
.await;
let client = LMStudioClient::from_host_root(server.uri());
let client = client_from_host_root(server.uri(), LMSTUDIO_CONNECTION_TIMEOUT);
let models = client.fetch_models().await.expect("fetch models");
assert!(models.contains(&"openai/gpt-oss-20b".to_string()));
}
@@ -260,7 +269,7 @@ mod tests {
.mount(&server)
.await;
let client = LMStudioClient::from_host_root(server.uri());
let client = client_from_host_root(server.uri(), LMSTUDIO_CONNECTION_TIMEOUT);
let result = client.fetch_models().await;
assert!(result.is_err());
assert!(
@@ -288,7 +297,7 @@ mod tests {
.mount(&server)
.await;
let client = LMStudioClient::from_host_root(server.uri());
let client = client_from_host_root(server.uri(), LMSTUDIO_CONNECTION_TIMEOUT);
let result = client.fetch_models().await;
assert!(result.is_err());
assert!(
@@ -316,13 +325,40 @@ mod tests {
.mount(&server)
.await;
let client = LMStudioClient::from_host_root(server.uri());
let client = client_from_host_root(server.uri(), LMSTUDIO_CONNECTION_TIMEOUT);
client
.check_server()
.await
.expect("server check should pass");
}
#[tokio::test]
async fn test_check_server_allows_slow_response_after_connect() {
if std::env::var(codex_core::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() {
tracing::info!(
"{} is set; skipping test_check_server_allows_slow_response_after_connect",
codex_core::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR
);
return;
}
let server = wiremock::MockServer::start().await;
wiremock::Mock::given(wiremock::matchers::method("GET"))
.and(wiremock::matchers::path("/models"))
.respond_with(
wiremock::ResponseTemplate::new(200).set_delay(Duration::from_millis(250)),
)
.mount(&server)
.await;
let client = client_from_host_root(server.uri(), Duration::from_millis(100));
client
.check_server()
.await
.expect("server check should allow a slow response after connecting");
}
#[tokio::test]
async fn test_check_server_error() {
if std::env::var(codex_core::spawn::CODEX_SANDBOX_NETWORK_DISABLED_ENV_VAR).is_ok() {
@@ -340,7 +376,7 @@ mod tests {
.mount(&server)
.await;
let client = LMStudioClient::from_host_root(server.uri());
let client = client_from_host_root(server.uri(), LMSTUDIO_CONNECTION_TIMEOUT);
let result = client.check_server().await;
assert!(result.is_err());
assert!(
@@ -385,13 +421,4 @@ mod tests {
}
}
}
#[test]
fn test_from_host_root() {
let client = LMStudioClient::from_host_root("http://localhost:1234");
assert_eq!(client.base_url, "http://localhost:1234");
let client = LMStudioClient::from_host_root("https://example.com:8080/api");
assert_eq!(client.base_url, "https://example.com:8080/api");
}
}