From e0bc0c05469010a6728683d1342964e09e5adf50 Mon Sep 17 00:00:00 2001 From: bianbiandashen <282758717@qq.com> Date: Sat, 14 Feb 2026 20:38:16 +0800 Subject: [PATCH] feat: add configurable timeouts and connection pooling for upstream requests Add HTTP client timeout and connection pool configuration to prevent resource exhaustion when upstream LLM APIs are slow or unresponsive. Configuration options (in [upstream] section): - connect_timeout_secs: TCP connection timeout (default: 30s) - request_timeout_secs: Full request/response timeout (default: 300s) - pool_idle_timeout_secs: Idle connection TTL (default: 90s) - pool_max_idle_per_host: Max idle connections per host (default: 32) The default request timeout of 5 minutes accommodates long-running LLM completions while still providing protection against hung connections. Changes: - config.rs: Add timeout fields to UpstreamConfig with sensible defaults - proxy.rs: Add ClientConfig struct and with_config() constructor - lib.rs: Wire up configuration in AppState::from_config() - clawshell.example.toml: Document new configuration options --- clawshell.example.toml | 16 +++++++ src/config.rs | 41 ++++++++++++++++++ src/lib.rs | 15 ++++++- src/proxy.rs | 94 ++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 164 insertions(+), 2 deletions(-) diff --git a/clawshell.example.toml b/clawshell.example.toml index fa0eb11..961cf0a 100644 --- a/clawshell.example.toml +++ b/clawshell.example.toml @@ -11,6 +11,22 @@ port = 18790 base_url = "https://api.openai.com" anthropic_base_url = "https://api.anthropic.com" +# HTTP client timeout and connection pool settings. +# These control how ClawShell connects to upstream LLM APIs. + +# Maximum time (seconds) to wait for TCP connection establishment. +# connect_timeout_secs = 30 + +# Maximum time (seconds) for the entire request/response cycle. +# LLM completions can take several minutes, so this is set high by default. +# request_timeout_secs = 300 + +# How long (seconds) idle connections stay in the pool before being closed. +# pool_idle_timeout_secs = 90 + +# Maximum number of idle connections to keep per host. +# pool_max_idle_per_host = 32 + # Virtual-to-real API key mappings # Multiple virtual keys can map to the same real key. # The "provider" field determines which upstream to use ("openai" or "anthropic"). diff --git a/src/config.rs b/src/config.rs index 9c10138..ec8d3f4 100644 --- a/src/config.rs +++ b/src/config.rs @@ -60,6 +60,31 @@ pub struct UpstreamConfig { pub anthropic_base_url: Option, #[serde(default = "default_anthropic_version")] pub anthropic_version: String, + + /// Connection timeout in seconds. This is the maximum time to wait for a + /// TCP connection to be established with the upstream server. + /// Default: 30 seconds. + #[serde(default = "default_connect_timeout_secs")] + pub connect_timeout_secs: u64, + + /// Request timeout in seconds. This is the maximum time to wait for the + /// entire request/response cycle, including connection, sending the request, + /// and receiving the response. For LLM APIs, this should be set high enough + /// to accommodate long-running completions. + /// Default: 300 seconds (5 minutes). + #[serde(default = "default_request_timeout_secs")] + pub request_timeout_secs: u64, + + /// Connection pool idle timeout in seconds. Connections that remain idle + /// longer than this duration will be closed. + /// Default: 90 seconds. + #[serde(default = "default_pool_idle_timeout_secs")] + pub pool_idle_timeout_secs: u64, + + /// Maximum number of idle connections per host in the connection pool. + /// Default: 32. + #[serde(default = "default_pool_max_idle_per_host")] + pub pool_max_idle_per_host: usize, } fn default_anthropic_version() -> String { @@ -70,6 +95,22 @@ fn default_base_url() -> String { "https://api.openai.com".to_string() } +fn default_connect_timeout_secs() -> u64 { + 30 +} + +fn default_request_timeout_secs() -> u64 { + 300 // 5 minutes - LLM completions can take a while +} + +fn default_pool_idle_timeout_secs() -> u64 { + 90 +} + +fn default_pool_max_idle_per_host() -> usize { + 32 +} + #[derive(Debug, Deserialize, Clone)] pub struct KeyMapping { pub virtual_key: String, diff --git a/src/lib.rs b/src/lib.rs index c67987b..a13cb63 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,7 +14,7 @@ pub mod tui; use crate::config::{Config, Provider}; use crate::dlp::DlpScanner; use crate::keys::{KeyManager, ResolvedKey}; -use crate::proxy::ProxyClient; +use crate::proxy::{ClientConfig, ProxyClient}; use axum::Router; use axum::body::Body; @@ -38,6 +38,8 @@ pub struct AppState { impl AppState { pub fn from_config(config: &Config) -> Self { + use std::time::Duration; + let mut upstream_urls = BTreeMap::new(); upstream_urls.insert(Provider::Openai, config.upstream_url(Provider::Openai)); upstream_urls.insert( @@ -59,15 +61,24 @@ impl AppState { }) .collect(); + // Build client configuration from upstream settings + let client_config = ClientConfig { + connect_timeout: Duration::from_secs(config.upstream.connect_timeout_secs), + request_timeout: Duration::from_secs(config.upstream.request_timeout_secs), + pool_idle_timeout: Duration::from_secs(config.upstream.pool_idle_timeout_secs), + pool_max_idle_per_host: config.upstream.pool_max_idle_per_host, + }; + Self { key_manager: Arc::new(KeyManager::new(key_mappings)), dlp_scanner: Arc::new( DlpScanner::new(&config.dlp.patterns, config.dlp.scan_responses) .expect("Failed to compile DLP patterns"), ), - proxy_client: Arc::new(ProxyClient::with_upstream_urls( + proxy_client: Arc::new(ProxyClient::with_config( upstream_urls, config.upstream.anthropic_version.clone(), + client_config, )), } } diff --git a/src/proxy.rs b/src/proxy.rs index c620aaa..b10734c 100644 --- a/src/proxy.rs +++ b/src/proxy.rs @@ -9,8 +9,37 @@ use futures_util::TryStreamExt; use reqwest::Client; use std::collections::BTreeMap; use std::io::Error as IoError; +use std::time::Duration; use tracing::{debug, trace}; +/// Configuration for HTTP client timeouts and connection pooling. +/// +/// These settings control how the proxy client behaves when connecting to +/// upstream LLM APIs. The defaults are tuned for typical LLM workloads where +/// responses may take several minutes for long completions. +#[derive(Debug, Clone)] +pub struct ClientConfig { + /// Maximum time to wait for TCP connection establishment. + pub connect_timeout: Duration, + /// Maximum time for the entire request/response cycle. + pub request_timeout: Duration, + /// How long idle connections stay in the pool before being closed. + pub pool_idle_timeout: Duration, + /// Maximum idle connections to keep per host. + pub pool_max_idle_per_host: usize, +} + +impl Default for ClientConfig { + fn default() -> Self { + Self { + connect_timeout: Duration::from_secs(30), + request_timeout: Duration::from_secs(300), // 5 minutes for LLM completions + pool_idle_timeout: Duration::from_secs(90), + pool_max_idle_per_host: 32, + } + } +} + #[derive(Debug)] pub struct ProxyClient { client: Client, @@ -19,14 +48,45 @@ pub struct ProxyClient { } impl ProxyClient { + /// Creates a new ProxyClient with default timeout configuration. pub fn with_upstream_urls( upstream_urls: BTreeMap, anthropic_version: String, + ) -> Self { + Self::with_config(upstream_urls, anthropic_version, ClientConfig::default()) + } + + /// Creates a new ProxyClient with custom timeout and pool configuration. + /// + /// # Arguments + /// * `upstream_urls` - Map of provider to base URL + /// * `anthropic_version` - Anthropic API version header value + /// * `config` - Timeout and connection pool settings + /// + /// # Panics + /// Panics if the HTTP client fails to build (e.g., TLS initialization failure). + pub fn with_config( + upstream_urls: BTreeMap, + anthropic_version: String, + config: ClientConfig, ) -> Self { let client = Client::builder() .redirect(reqwest::redirect::Policy::none()) + .connect_timeout(config.connect_timeout) + .timeout(config.request_timeout) + .pool_idle_timeout(config.pool_idle_timeout) + .pool_max_idle_per_host(config.pool_max_idle_per_host) .build() .expect("Failed to build reqwest client"); + + debug!( + connect_timeout_ms = config.connect_timeout.as_millis(), + request_timeout_ms = config.request_timeout.as_millis(), + pool_idle_timeout_ms = config.pool_idle_timeout.as_millis(), + pool_max_idle_per_host = config.pool_max_idle_per_host, + "ProxyClient initialized with timeout configuration" + ); + Self { client, upstream_urls, @@ -308,4 +368,38 @@ mod tests { let json: serde_json::Value = serde_json::from_slice(&body).unwrap(); assert!(json["error"].as_str().unwrap().contains("TRACE")); } + + #[test] + fn test_client_config_default() { + let config = ClientConfig::default(); + assert_eq!(config.connect_timeout, Duration::from_secs(30)); + assert_eq!(config.request_timeout, Duration::from_secs(300)); + assert_eq!(config.pool_idle_timeout, Duration::from_secs(90)); + assert_eq!(config.pool_max_idle_per_host, 32); + } + + #[test] + fn test_proxy_client_with_custom_config() { + let mut urls = BTreeMap::new(); + urls.insert(Provider::Openai, "https://api.openai.com".to_string()); + + let custom_config = ClientConfig { + connect_timeout: Duration::from_secs(10), + request_timeout: Duration::from_secs(60), + pool_idle_timeout: Duration::from_secs(30), + pool_max_idle_per_host: 16, + }; + + // This should not panic - verifies the client builds successfully + let _client = ProxyClient::with_config(urls, "2023-06-01".to_string(), custom_config); + } + + #[test] + fn test_proxy_client_with_default_config() { + let mut urls = BTreeMap::new(); + urls.insert(Provider::Openai, "https://api.openai.com".to_string()); + + // with_upstream_urls should use default config internally + let _client = ProxyClient::with_upstream_urls(urls, "2023-06-01".to_string()); + } }