fix: enable TCP keepalive on default httpx transports to prevent NAT idle-timeout drops

pyproject.toml

@@ -9,7 +9,7 @@ authors = [
 ]
 
 dependencies = [
-  "httpx>=0.23.0, <1",
+  "httpx>=0.25.0, <1",
   "pydantic>=1.9.0, <3",
     "typing-extensions>=4.11, <5", "typing-extensions>=4.14, <5",
   "anyio>=3.5.0, <5",

src/openai/_base_client.py

@@ -5,6 +5,7 @@
 import time
 import uuid
 import email
+import socket
 import asyncio
 import inspect
 import logging
@@ -831,11 +832,36 @@ def _idempotency_key(self) -> str:
         return f"stainless-python-retry-{uuid.uuid4()}"
 
 
+def _build_keepalive_socket_options() -> list[tuple[int, int, int]]:
+    # Enable TCP keepalive to prevent silent connection drops by NAT gateways and
+    # load balancers during long-running inference calls (generation can hold a
+    # TCP connection idle for 300–600 s, exceeding the ~350 s AWS NAT Gateway
+    # idle timeout). Mirrors the pattern used by the Anthropic Python SDK.
+    options: list[tuple[int, int, int]] = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
+    # TCP_KEEPIDLE (Linux) / TCP_KEEPALIVE (macOS): seconds idle before first probe
+    if hasattr(socket, "TCP_KEEPIDLE"):
+        options.append((socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, 60))
+    elif hasattr(socket, "TCP_KEEPALIVE"):
+        options.append((socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, 60))
+    # TCP_KEEPINTVL: seconds between subsequent probes
+    if hasattr(socket, "TCP_KEEPINTVL"):
+        options.append((socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60))
+    # TCP_KEEPCNT: drop the connection after this many unanswered probes
+    if hasattr(socket, "TCP_KEEPCNT"):
+        options.append((socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5))
+    return options
+
+
 class _DefaultHttpxClient(httpx.Client):
     def __init__(self, **kwargs: Any) -> None:
         kwargs.setdefault("timeout", DEFAULT_TIMEOUT)
         kwargs.setdefault("limits", DEFAULT_CONNECTION_LIMITS)
         kwargs.setdefault("follow_redirects", True)
+        if "transport" not in kwargs:
+            kwargs["transport"] = httpx.HTTPTransport(
+                limits=kwargs.get("limits", DEFAULT_CONNECTION_LIMITS),
+                socket_options=_build_keepalive_socket_options(),
+            )
         super().__init__(**kwargs)
 
 
@@ -1423,6 +1449,11 @@ def __init__(self, **kwargs: Any) -> None:
         kwargs.setdefault("timeout", DEFAULT_TIMEOUT)
         kwargs.setdefault("limits", DEFAULT_CONNECTION_LIMITS)
         kwargs.setdefault("follow_redirects", True)
+        if "transport" not in kwargs:
+            kwargs["transport"] = httpx.AsyncHTTPTransport(
+                limits=kwargs.get("limits", DEFAULT_CONNECTION_LIMITS),
+                socket_options=_build_keepalive_socket_options(),
+            )
         super().__init__(**kwargs)
 
 

tests/test_client.py

@@ -1304,6 +1304,22 @@ def test_default_client_creation(self) -> None:
             limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
         )
 
+    def test_default_transport_has_tcp_keepalive(self) -> None:
+        import socket as socket_module
+
+        client = OpenAI(base_url=base_url, api_key=api_key)
+        transport = client._client._transport
+        assert isinstance(transport, httpx.HTTPTransport)
+        socket_options = transport._pool._socket_options
+        assert any(
+            opt == (socket_module.SOL_SOCKET, socket_module.SO_KEEPALIVE, 1) for opt in socket_options
+        ), "Default sync transport should have SO_KEEPALIVE enabled to survive NAT idle timeouts"
+
+    def test_custom_http_client_transport_is_not_overridden(self) -> None:
+        with httpx.Client() as http_client:
+            client = OpenAI(base_url=base_url, api_key=api_key, http_client=http_client)
+            assert client._client is http_client, "A caller-supplied http_client must not be replaced"
+
     @pytest.mark.respx(base_url=base_url)
     def test_follow_redirects(self, respx_mock: MockRouter, client: OpenAI) -> None:
         # Test that the default follow_redirects=True allows following redirects
@@ -2564,6 +2580,22 @@ async def test_default_client_creation(self) -> None:
             limits=httpx.Limits(max_connections=100, max_keepalive_connections=20),
         )
 
+    async def test_default_transport_has_tcp_keepalive(self) -> None:
+        import socket as socket_module
+
+        client = AsyncOpenAI(base_url=base_url, api_key=api_key)
+        transport = client._client._transport
+        assert isinstance(transport, httpx.AsyncHTTPTransport)
+        socket_options = transport._pool._socket_options
+        assert any(
+            opt == (socket_module.SOL_SOCKET, socket_module.SO_KEEPALIVE, 1) for opt in socket_options
+        ), "Default async transport should have SO_KEEPALIVE enabled to survive NAT idle timeouts"
+
+    async def test_custom_async_http_client_transport_is_not_overridden(self) -> None:
+        async with httpx.AsyncClient() as http_client:
+            client = AsyncOpenAI(base_url=base_url, api_key=api_key, http_client=http_client)
+            assert client._client is http_client, "A caller-supplied http_client must not be replaced"
+
     @pytest.mark.respx(base_url=base_url)
     async def test_follow_redirects(self, respx_mock: MockRouter, async_client: AsyncOpenAI) -> None:
         # Test that the default follow_redirects=True allows following redirects