diff --git a/codex-rs/codex-client/src/lib.rs b/codex-rs/codex-client/src/lib.rs index ef451a53d8..91d9e6799d 100644 --- a/codex-rs/codex-client/src/lib.rs +++ b/codex-rs/codex-client/src/lib.rs @@ -33,6 +33,7 @@ pub use crate::outbound_proxy::RouteFailureClass; pub use crate::outbound_proxy::SystemProxyRouteDecision; pub use crate::outbound_proxy::build_reqwest_client_for_route; pub use crate::outbound_proxy::resolve_system_proxy_for_url; +pub use crate::outbound_proxy::resolve_system_proxy_for_url_async; pub use crate::request::EncodedJsonBody; pub use crate::request::PreparedRequestBody; pub use crate::request::Request; diff --git a/codex-rs/codex-client/src/outbound_proxy.rs b/codex-rs/codex-client/src/outbound_proxy.rs index 3d5c340f69..d0018296c7 100644 --- a/codex-rs/codex-client/src/outbound_proxy.rs +++ b/codex-rs/codex-client/src/outbound_proxy.rs @@ -24,11 +24,14 @@ const SYSTEM_PROXY_SUCCESS_CACHE_TTL: Duration = Duration::from_secs(60); const SYSTEM_PROXY_UNAVAILABLE_CACHE_TTL: Duration = Duration::from_secs(5); const SYSTEM_PROXY_CACHE_MAX_ENTRIES: usize = 256; +mod async_resolver; #[cfg(target_os = "macos")] mod macos; #[cfg(target_os = "windows")] mod windows; +pub use async_resolver::resolve_system_proxy_for_url_async; + /// Coarse semantic bucket for the HTTP or WebSocket client being constructed. /// /// This is not the selected proxy route or a concrete endpoint. It labels the diff --git a/codex-rs/codex-client/src/outbound_proxy/async_resolver.rs b/codex-rs/codex-client/src/outbound_proxy/async_resolver.rs new file mode 100644 index 0000000000..e1ea1ff88c --- /dev/null +++ b/codex-rs/codex-client/src/outbound_proxy/async_resolver.rs @@ -0,0 +1,125 @@ +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::OnceLock; + +use tokio::sync::Mutex; +use tokio::sync::Semaphore; +use tokio::sync::watch; + +use super::RouteFailureClass; +use super::SystemProxyRouteDecision; +use super::cached_system_proxy_decision; +use super::resolve_system_proxy_for_url; +use super::system_proxy_cache_key; + +const MAX_CONCURRENT_SYSTEM_PROXY_RESOLUTIONS: usize = 2; + +type PlatformResolver = fn(&str) -> SystemProxyRouteDecision; +type ResolutionSender = watch::Sender>; + +static ASYNC_SYSTEM_PROXY_RESOLVER: OnceLock> = + OnceLock::new(); + +/// Resolves system proxy/PAC/WPAD routing without blocking an async runtime worker. +pub async fn resolve_system_proxy_for_url_async(request_url: &str) -> SystemProxyRouteDecision { + if let Some(decision) = cached_system_proxy_decision(request_url) { + return decision.into(); + } + + ASYNC_SYSTEM_PROXY_RESOLVER + .get_or_init(|| { + AsyncSystemProxyResolver::new( + MAX_CONCURRENT_SYSTEM_PROXY_RESOLUTIONS, + resolve_system_proxy_for_url as PlatformResolver, + ) + }) + .resolve(request_url) + .await +} + +struct AsyncSystemProxyResolver { + permits: Arc, + in_flight: Arc>>, + resolve: Arc, +} + +impl AsyncSystemProxyResolver +where + R: Fn(&str) -> SystemProxyRouteDecision + Send + Sync + 'static, +{ + fn new(max_concurrent_resolutions: usize, resolve: R) -> Self { + Self { + permits: Arc::new(Semaphore::new(max_concurrent_resolutions)), + in_flight: Arc::new(Mutex::new(HashMap::new())), + resolve: Arc::new(resolve), + } + } + + async fn resolve(&self, request_url: &str) -> SystemProxyRouteDecision { + let cache_key = system_proxy_cache_key(request_url); + let existing = self + .in_flight + .lock() + .await + .get(&cache_key) + .map(ResolutionSender::subscribe); + + let mut receiver = if let Some(receiver) = existing { + receiver + } else { + let Ok(permit) = Arc::clone(&self.permits).acquire_owned().await else { + return resolver_error(); + }; + let (receiver, leader) = { + let mut in_flight = self.in_flight.lock().await; + if let Some(sender) = in_flight.get(&cache_key) { + (sender.subscribe(), None) + } else { + let (sender, receiver) = watch::channel(None); + in_flight.insert(cache_key.clone(), sender.clone()); + (receiver, Some((sender, permit))) + } + }; + + if let Some((sender, permit)) = leader { + let in_flight = Arc::clone(&self.in_flight); + let resolve = Arc::clone(&self.resolve); + let request_url = request_url.to_string(); + tokio::spawn(async move { + let decision = match tokio::task::spawn_blocking(move || { + let _permit = permit; + resolve(&request_url) + }) + .await + { + Ok(decision) => decision, + Err(_) => resolver_error(), + }; + let _ = sender.send_replace(Some(decision)); + let _ = in_flight.lock().await.remove(&cache_key); + }); + } + + receiver + }; + + loop { + if let Some(decision) = receiver.borrow().clone() { + return decision; + } + if receiver.changed().await.is_err() { + return resolver_error(); + } + } + } +} + +fn resolver_error() -> SystemProxyRouteDecision { + SystemProxyRouteDecision::Unavailable { + failure: RouteFailureClass::ResolverError, + } +} + +#[cfg(test)] +#[path = "async_resolver_tests.rs"] +mod tests; diff --git a/codex-rs/codex-client/src/outbound_proxy/async_resolver_tests.rs b/codex-rs/codex-client/src/outbound_proxy/async_resolver_tests.rs new file mode 100644 index 0000000000..510df125d3 --- /dev/null +++ b/codex-rs/codex-client/src/outbound_proxy/async_resolver_tests.rs @@ -0,0 +1,86 @@ +use std::sync::Arc; +use std::sync::Condvar; +use std::sync::Mutex as StdMutex; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use pretty_assertions::assert_eq; +use tokio::sync::Notify; +use tokio::time::timeout; + +use super::*; + +const SLOW_URL: &str = "https://example.test/slow"; +const FAST_URL: &str = "https://example.test/fast"; +const TEST_TIMEOUT: Duration = Duration::from_secs(/*secs*/ 1); + +#[tokio::test] +async fn slow_resolution_keeps_runtime_responsive_and_coalesces_duplicates() { + let slow_calls = Arc::new(AtomicUsize::new(0)); + let slow_started = Arc::new(Notify::new()); + let release_slow = Arc::new((StdMutex::new(false), Condvar::new())); + let resolver = Arc::new(AsyncSystemProxyResolver::new( + /*max_concurrent_resolutions*/ 2, + { + let slow_calls = Arc::clone(&slow_calls); + let slow_started = Arc::clone(&slow_started); + let release_slow = Arc::clone(&release_slow); + move |request_url| { + if request_url == SLOW_URL { + slow_calls.fetch_add(/*val*/ 1, Ordering::SeqCst); + slow_started.notify_one(); + let (released, condition) = &*release_slow; + let mut released = released.lock().unwrap(); + while !*released { + released = condition.wait(released).unwrap(); + } + } + SystemProxyRouteDecision::Direct + } + }, + )); + + let first = tokio::spawn({ + let resolver = Arc::clone(&resolver); + async move { resolver.resolve(SLOW_URL).await } + }); + timeout(TEST_TIMEOUT, slow_started.notified()) + .await + .unwrap(); + + let duplicate = tokio::spawn({ + let resolver = Arc::clone(&resolver); + async move { resolver.resolve(SLOW_URL).await } + }); + let cache_key = system_proxy_cache_key(SLOW_URL); + let duplicate_joined = timeout(TEST_TIMEOUT, async { + loop { + let receiver_count = resolver + .in_flight + .lock() + .await + .get(&cache_key) + .map(ResolutionSender::receiver_count) + .unwrap_or_default(); + if receiver_count >= 2 { + break; + } + tokio::task::yield_now().await; + } + }) + .await; + let fast = timeout(TEST_TIMEOUT, resolver.resolve(FAST_URL)).await; + + let (released, condition) = &*release_slow; + *released.lock().unwrap() = true; + condition.notify_all(); + let first = first.await.unwrap(); + let duplicate = duplicate.await.unwrap(); + + assert!(duplicate_joined.is_ok()); + assert_eq!(fast.unwrap(), SystemProxyRouteDecision::Direct); + assert_eq!(first, SystemProxyRouteDecision::Direct); + assert_eq!(duplicate, SystemProxyRouteDecision::Direct); + assert_eq!(slow_calls.load(Ordering::SeqCst), 1); +} diff --git a/codex-rs/network-proxy/src/http_proxy.rs b/codex-rs/network-proxy/src/http_proxy.rs index 24801ea68e..41acdfad7c 100644 --- a/codex-rs/network-proxy/src/http_proxy.rs +++ b/codex-rs/network-proxy/src/http_proxy.rs @@ -401,7 +401,7 @@ async fn http_connect_proxy(upgraded: Upgraded) -> Result<(), Infallible> { }; let target_host = target.host.to_string(); let target_url = https_url_for_host_port(&target_host, target.port); - proxy_for_connect(&target_url, respect_system_proxy) + proxy_for_connect(&target_url, respect_system_proxy).await } else { None }; diff --git a/codex-rs/network-proxy/src/upstream.rs b/codex-rs/network-proxy/src/upstream.rs index 669acbe224..989b80cc80 100644 --- a/codex-rs/network-proxy/src/upstream.rs +++ b/codex-rs/network-proxy/src/upstream.rs @@ -2,7 +2,7 @@ use crate::connect_policy::TargetCheckedTcpConnector; use crate::state::NetworkProxyState; use codex_client::RouteFailureClass; use codex_client::SystemProxyRouteDecision; -use codex_client::resolve_system_proxy_for_url; +use codex_client::resolve_system_proxy_for_url_async; use codex_utils_rustls_provider::ensure_rustls_crypto_provider; use rama_core::Layer; use rama_core::Service; @@ -102,9 +102,9 @@ impl ProxyConfig { } } - fn proxy_for_url(&self, request_url: &str, is_secure: bool) -> Option { + async fn proxy_for_url(&self, request_url: &str, is_secure: bool) -> Option { if self.respect_system_proxy { - match system_proxy_for_url(request_url) { + match system_proxy_for_url(request_url).await { Ok(proxy) => return proxy, Err(err) => { warn!("system proxy unavailable; falling back to env proxy ({err})"); @@ -115,8 +115,10 @@ impl ProxyConfig { } } -fn system_proxy_for_url(request_url: &str) -> Result, UpstreamProxyError> { - match resolve_system_proxy_for_url(request_url) { +async fn system_proxy_for_url( + request_url: &str, +) -> Result, UpstreamProxyError> { + match resolve_system_proxy_for_url_async(request_url).await { SystemProxyRouteDecision::Direct => Ok(None), SystemProxyRouteDecision::Proxy { url } => proxy_address_from_system_url(&url).map(Some), SystemProxyRouteDecision::Unavailable { failure } => { @@ -173,7 +175,7 @@ fn read_proxy_env(keys: &[&str]) -> Option { None } -pub(crate) fn proxy_for_connect( +pub(crate) async fn proxy_for_connect( request_url: &str, respect_system_proxy: bool, ) -> Option { @@ -182,7 +184,7 @@ pub(crate) fn proxy_for_connect( } else { ProxyConfig::from_env() }; - config.proxy_for_url(request_url, /*is_secure*/ true) + config.proxy_for_url(request_url, /*is_secure*/ true).await } #[derive(Clone)] @@ -285,11 +287,14 @@ impl Service> for UpstreamClient { .as_ref() .map(|ctx| ctx.host_with_port().to_string()) .unwrap_or_else(|| "".to_string()); - let proxy = request_context.as_ref().and_then(|request_context| { + let proxy = if let Some(request_context) = request_context.as_ref() { let request_url = request_target_url(&req, request_context); self.proxy_config .proxy_for_url(&request_url, request_context.protocol.is_secure()) - }); + .await + } else { + None + }; match proxy.as_ref() { Some(proxy) => info!( "HTTP upstream route selected (target={authority}, route=upstream_proxy, proxy={})",