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()); + } }