feat(runtime): make connection pooling configurable

engine/baml-lib/baml/tests/validation_files/client/http_config_aws_pooling.baml

@@ -0,0 +1,18 @@
+client<llm> AwsPooling {
+  provider aws-bedrock
+  options {
+    model "anthropic.claude-3-sonnet-20240229-v1:0"
+    http {
+      enable_connection_pooling true
+    }
+  }
+}
+
+// error: enable_connection_pooling is not supported for aws clients because their HTTP transport is managed externally
+//   -->  client/http_config_aws_pooling.baml:5
+//    | 
+//  4 |     model "anthropic.claude-3-sonnet-20240229-v1:0"
+//  5 |     http {
+//  6 |       enable_connection_pooling true
+//  7 |     }
+//    | 

engine/baml-lib/baml/tests/validation_files/client/http_config_invalid_pooling.baml

@@ -0,0 +1,16 @@
+client<llm> InvalidPoolingType {
+  provider openai
+  options {
+    model "gpt-4"
+    http {
+      enable_connection_pooling "yes"
+    }
+  }
+}
+
+// error: enable_connection_pooling must be a bool. Got: string
+//   -->  client/http_config_invalid_pooling.baml:6
+//    | 
+//  5 |     http {
+//  6 |       enable_connection_pooling "yes"
+//    | 

engine/baml-lib/baml/tests/validation_files/client/http_config_regular_with_total.baml

@@ -8,7 +8,7 @@ client<llm> RegularWithTotal {
   }
 }
 
-// error: Unrecognized field 'total_timeout_ms' in http configuration block. 'total_timeout_ms' is only available for composite clients (fallback/round-robin). For regular clients, use: connect_timeout_ms, request_timeout_ms, time_to_first_token_timeout_ms, idle_timeout_ms
+// error: Unrecognized field 'total_timeout_ms' in http configuration block. 'total_timeout_ms' is only available for composite clients (fallback/round-robin). For regular clients, use: enable_connection_pooling, connect_timeout_ms, request_timeout_ms, time_to_first_token_timeout_ms, idle_timeout_ms
 //   -->  client/http_config_regular_with_total.baml:5
 //    | 
 //  4 |     model "gpt-4"

engine/baml-lib/baml/tests/validation_files/client/http_config_valid.baml

@@ -4,6 +4,7 @@ client<llm> ValidHttpClient {
   options {
     model "gpt-4"
     http {
+      enable_connection_pooling true
       connect_timeout_ms 5000
       request_timeout_ms 30000
     }
@@ -16,6 +17,7 @@ client<llm> ValidStreamingClient {
   options {
     model "claude-3"
     http {
+      enable_connection_pooling false
       time_to_first_token_timeout_ms 15000
       idle_timeout_ms 10000
     }
@@ -42,4 +44,4 @@ client<llm> ValidFallbackClient {
       total_timeout_ms 60000
     }
   }
-}
+}

engine/baml-lib/llm-client/src/clients/helpers.rs

@@ -10,9 +10,11 @@ use crate::{
     UnresolvedRolesSelection,
 };
 
-/// Configuration for HTTP timeouts
+/// Configuration for HTTP clients
 #[derive(Debug, Clone, Default, BamlHash)]
 pub struct HttpConfig {
+    /// Disabled by default because a pooled client must not be inherited across `fork()`.
+    pub enable_connection_pooling: bool,
     pub connect_timeout_ms: Option<u64>,
     pub request_timeout_ms: Option<u64>,
     pub time_to_first_token_timeout_ms: Option<u64>,
@@ -515,23 +517,41 @@ impl<Meta: Clone> PropertyHandler<Meta> {
                     // Define allowed fields based on provider type
                     let is_composite =
                         provider_type == "fallback" || provider_type == "round-robin";
+                    let supports_connection_pooling = !is_composite && provider_type != "aws";
+                    let mut regular_http_fields = vec![
+                        "connect_timeout_ms",
+                        "request_timeout_ms",
+                        "time_to_first_token_timeout_ms",
+                        "idle_timeout_ms",
+                    ];
+                    if supports_connection_pooling {
+                        regular_http_fields.insert(0, "enable_connection_pooling");
+                    }
                     let allowed_fields: HashSet<&str> = if is_composite {
                         // Composite clients only support total_timeout_ms
                         vec!["total_timeout_ms"].into_iter().collect()
                     } else {
-                        // Regular clients support all timeout types except total_timeout_ms
-                        vec![
-                            "connect_timeout_ms",
-                            "request_timeout_ms",
-                            "time_to_first_token_timeout_ms",
-                            "idle_timeout_ms",
-                        ]
-                        .into_iter()
-                        .collect()
+                        regular_http_fields.iter().copied().collect()
                     };
 
                     for (key, (_, value)) in config_map {
                         match key.as_str() {
+                            "enable_connection_pooling" if supports_connection_pooling => {
+                                match value.into_bool() {
+                                    Ok((enabled, _)) => {
+                                        http_config.enable_connection_pooling = enabled;
+                                    }
+                                    Err(other) => {
+                                        self.push_error(
+                                            format!(
+                                                "enable_connection_pooling must be a bool. Got: {}",
+                                                other.r#type()
+                                            ),
+                                            other.meta().clone(),
+                                        );
+                                    }
+                                }
+                            }
                             "connect_timeout_ms" if !is_composite => {
                                 let value_meta = value.meta().clone();
                                 match value.into_numeric() {
@@ -684,22 +704,23 @@ impl<Meta: Clone> PropertyHandler<Meta> {
                                 }
                             } else {
                                 // For regular clients
-                                let all_timeout_fields = vec![
-                                    "connect_timeout_ms",
-                                    "request_timeout_ms",
-                                    "time_to_first_token_timeout_ms",
-                                    "idle_timeout_ms",
-                                    "total_timeout_ms", // Include for suggestions
-                                ];
+                                let mut all_http_fields = regular_http_fields.clone();
+                                all_http_fields.push("total_timeout_ms");
+                                let supported_fields = regular_http_fields.join(", ");
 
-                                if unrecognized_field == "total_timeout_ms" {
+                                if unrecognized_field == "enable_connection_pooling"
+                                    && !supports_connection_pooling
+                                {
+                                    format!(
+                                        "enable_connection_pooling is not supported for {provider_type} clients because their HTTP transport is managed externally"
+                                    )
+                                } else if unrecognized_field == "total_timeout_ms" {
                                     // Special case for total_timeout_ms in regular clients
-                                    "Unrecognized field 'total_timeout_ms' in http configuration block. \
+                                    format!("Unrecognized field 'total_timeout_ms' in http configuration block. \
                                         'total_timeout_ms' is only available for composite clients (fallback/round-robin). \
-                                        For regular clients, use: connect_timeout_ms, request_timeout_ms, \
-                                        time_to_first_token_timeout_ms, idle_timeout_ms".to_string()
+                                        For regular clients, use: {supported_fields}")
                                 } else if let Some(suggestion) =
-                                    find_best_match(unrecognized_field, &all_timeout_fields)
+                                    find_best_match(unrecognized_field, &all_http_fields)
                                 {
                                     if suggestion == "total_timeout_ms" {
                                         format!(
@@ -715,8 +736,7 @@ impl<Meta: Clone> PropertyHandler<Meta> {
                                 } else {
                                     format!(
                                         "Unrecognized field '{unrecognized_field}' in http configuration block. \
-                                        Supported timeout fields are: connect_timeout_ms, request_timeout_ms, \
-                                        time_to_first_token_timeout_ms, idle_timeout_ms"
+                                        Supported fields are: {supported_fields}"
                                     )
                                 }
                             };

engine/baml-runtime/src/request/mod.rs

@@ -43,7 +43,10 @@ pub fn create_http_client(
             let danger_accept_invalid_certs = matches!(std::env::var("DANGER_ACCEPT_INVALID_CERTS").as_deref(), Ok("1"));
             let mut builder = reqwest::Client::builder()
                 .danger_accept_invalid_certs(danger_accept_invalid_certs)
-                .http2_keep_alive_interval(Some(Duration::from_secs(10)))
+                .http2_keep_alive_interval(Some(Duration::from_secs(10)));
+
+            if !http_config.enable_connection_pooling {
+                builder = builder
                 // To prevent stalling in python, we set the pool to 0 and idle timeout to 0.
                 // See:
                 // https://github.com/seanmonstar/reqwest/issues/600
@@ -52,6 +55,7 @@ pub fn create_http_client(
                 // https://github.com/Azure/azure-sdk-for-rust/pull/1550
                 .pool_max_idle_per_host(0)
                 .pool_idle_timeout(std::time::Duration::from_nanos(1));
+            }
 
             // Apply connect timeout if specified
             // Note: 0 means infinite timeout (no timeout)
@@ -85,3 +89,78 @@ pub(crate) fn create_tracing_client() -> Result<reqwest::Client> {
 
     cb.build().context("Failed to create reqwest client")
 }
+
+#[cfg(all(test, not(target_arch = "wasm32")))]
+mod tests {
+    use std::sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    };
+
+    use tokio::{
+        io::{AsyncReadExt, AsyncWriteExt},
+        net::{TcpListener, TcpStream},
+    };
+
+    use super::create_http_client;
+
+    const RESPONSE: &[u8] =
+        b"HTTP/1.1 200 OK\r\ncontent-length: 2\r\nconnection: keep-alive\r\n\r\nok";
+
+    async fn serve_connection(mut socket: TcpStream) {
+        let mut pending = Vec::new();
+        let mut buffer = [0_u8; 4096];
+
+        loop {
+            let bytes_read = socket.read(&mut buffer).await.unwrap();
+            if bytes_read == 0 {
+                return;
+            }
+            pending.extend_from_slice(&buffer[..bytes_read]);
+
+            while let Some(header_end) = pending.windows(4).position(|bytes| bytes == b"\r\n\r\n") {
+                pending.drain(..header_end + 4);
+                socket.write_all(RESPONSE).await.unwrap();
+            }
+        }
+    }
+
+    async fn accepted_connections(enable_connection_pooling: bool) -> usize {
+        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let connection_count = Arc::new(AtomicUsize::new(0));
+        let server_count = Arc::clone(&connection_count);
+        let server = tokio::spawn(async move {
+            loop {
+                let (socket, _) = listener.accept().await.unwrap();
+                server_count.fetch_add(1, Ordering::Relaxed);
+                tokio::spawn(serve_connection(socket));
+            }
+        });
+
+        let client = create_http_client(&internal_llm_client::HttpConfig {
+            enable_connection_pooling,
+            ..Default::default()
+        })
+        .unwrap();
+
+        for _ in 0..3 {
+            let response = client
+                .get(format!("http://{address}"))
+                .send()
+                .await
+                .unwrap();
+            assert_eq!(response.text().await.unwrap(), "ok");
+        }
+
+        let accepted = connection_count.load(Ordering::Relaxed);
+        server.abort();
+        accepted
+    }
+
+    #[tokio::test]
+    async fn connection_pooling_is_opt_in() {
+        assert_eq!(accepted_connections(false).await, 3);
+        assert_eq!(accepted_connections(true).await, 1);
+    }
+}

fern/03-reference/baml/clients/connection-pooling.mdx

@@ -0,0 +1,40 @@
+---
+title: Connection Pooling
+subtitle: Reuse HTTP connections across requests
+---
+
+BAML disables HTTP connection pooling by default. This preserves compatibility
+with applications that create a BAML runtime before using Python's `fork()`
+start method.
+
+Enable connection reuse for a leaf client with the `http` option:
+
+```baml
+client<llm> MyClient {
+  provider openai
+  options {
+    model "gpt-4o"
+    api_key env.OPENAI_API_KEY
+    http {
+      enable_connection_pooling true
+    }
+  }
+}
+```
+
+Pooling avoids a new TCP and TLS handshake for every request. It is useful for
+asyncio, threaded, and multiprocessing applications that use the `spawn` start
+method.
+
+<Warning>
+Do not enable pooling when child processes inherit an existing BAML runtime or
+client through `fork()`. Create the BAML runtime inside each child process, or
+use the `spawn` start method instead.
+</Warning>
+
+Connection pooling is not available for
+[fallback](/reference/baml/clients/strategy/fallback) or
+[round-robin](/reference/baml/clients/strategy/round-robin) clients. Set it on
+their leaf clients instead. It is also not available for
+[AWS Bedrock](/reference/baml/clients/providers/aws-bedrock) clients, whose HTTP
+transport is managed by the AWS SDK.

fern/docs.yml

@@ -853,6 +853,9 @@ navigation:
           - page: "Timeout Configuration"
             slug: timeouts
             path: 03-reference/baml/clients/timeouts.mdx
+          - page: "Connection Pooling"
+            slug: connection-pooling
+            path: 03-reference/baml/clients/connection-pooling.mdx
           - page: "Retry Policy"
             path: 03-reference/baml/clients/strategy/retry.mdx
           - page: "Fallback"