diff --git a/codex-rs/Cargo.lock b/codex-rs/Cargo.lock index d7df4b89ad..bd0a291eaa 100644 --- a/codex-rs/Cargo.lock +++ b/codex-rs/Cargo.lock @@ -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", diff --git a/codex-rs/deny.toml b/codex-rs/deny.toml index 23f4468280..5730b3c878 100644 --- a/codex-rs/deny.toml +++ b/codex-rs/deny.toml @@ -243,7 +243,6 @@ deny = [ "codex-app-server", "codex-core", "codex-exec-server", - "codex-lmstudio", "codex-ollama", "codex-otel", "codex-protocol", diff --git a/codex-rs/http-client/src/client_builder.rs b/codex-rs/http-client/src/client_builder.rs index 22609b3d1d..34ff54e59e 100644 --- a/codex-rs/http-client/src/client_builder.rs +++ b/codex-rs/http-client/src/client_builder.rs @@ -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, follow_redirects: bool, + connect_timeout: Option, 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 { + 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, } diff --git a/codex-rs/http-client/src/route_aware_client_pool.rs b/codex-rs/http-client/src/route_aware_client_pool.rs index 0536533f4e..261ee1d2a2 100644 --- a/codex-rs/http-client/src/route_aware_client_pool.rs +++ b/codex-rs/http-client/src/route_aware_client_pool.rs @@ -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 reqwest::ClientBuilder + Send + Sync>, - request_logging: PoolRequestLogging, + client_builder: HttpClientBuilder, clients: Arc>>, } @@ -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}"), diff --git a/codex-rs/http-client/src/route_aware_client_pool_tests.rs b/codex-rs/http-client/src/route_aware_client_pool_tests.rs index fe042aa372..f7e98add3b 100644 --- a/codex-rs/http-client/src/route_aware_client_pool_tests.rs +++ b/codex-rs/http-client/src/route_aware_client_pool_tests.rs @@ -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(\"\"), .. }" ) ); } #[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(), ) } diff --git a/codex-rs/lmstudio/Cargo.toml b/codex-rs/lmstudio/Cargo.toml index e43d0b3bbe..e84bcfed51 100644 --- a/codex-rs/lmstudio/Cargo.toml +++ b/codex-rs/lmstudio/Cargo.toml @@ -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"] } diff --git a/codex-rs/lmstudio/src/client.rs b/codex-rs/lmstudio/src/client.rs index baad560115..8b0a89a932 100644 --- a/codex-rs/lmstudio/src/client.rs +++ b/codex-rs/lmstudio/src/client.rs @@ -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 { @@ -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) -> 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, + 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"); - } }