From f6b30e8e1e1e6e021eaa14ea7a625c19912a7efd Mon Sep 17 00:00:00 2001
From: KlyneChrysler <klynechrislu@gmail.com>
Date: Sun, 12 Jul 2026 10:07:18 +0800
Subject: [PATCH] fix: enable TCP keepalive on default httpx transports

Long running generate_content calls can stay silent on the wire for
longer than the idle timeout of NAT gateways and stateful firewalls,
which silently drop the connection and leave the response read hanging
forever. Enable SO_KEEPALIVE with a 60 second idle time on the default
sync and async transports, with platform guards for the option names.
User supplied transports and mounts are left untouched, and proxy
configuration keeps working since httpx mounts proxies over the base
transport.

Fixes #2705
---
 google/genai/_api_client.py                   | 32 +++++++++
 .../genai/tests/client/test_socket_options.py | 67 +++++++++++++++++++
 2 files changed, 99 insertions(+)
 create mode 100644 google/genai/tests/client/test_socket_options.py

diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py
index 4d90ba86e..07b821f03 100644
--- a/google/genai/_api_client.py
+++ b/google/genai/_api_client.py
@@ -30,6 +30,7 @@
 import math
 import os
 import random
+import socket
 import ssl
 import sys
 import threading
@@ -557,12 +558,39 @@ def retry_args(options: Optional[HttpRetryOptions]) -> _common.StringDict:
   }
 
 
+def _default_socket_options() -> list[tuple[int, int, int]]:
+  """Returns socket options that enable TCP keepalive on connections.
+
+  Long-running calls can stay silent on the wire for longer than the idle
+  timeout of NAT gateways and stateful firewalls on the egress path. Without
+  keepalive probes such connections are dropped silently and the response
+  read hangs forever (https://github.com/googleapis/python-genai/issues/2705).
+  """
+  options = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
+  # Start sending keepalive probes after 60 seconds of idleness, well below
+  # common NAT idle timeouts. The option name is platform dependent: Linux
+  # and Windows expose TCP_KEEPIDLE; macOS uses TCP_KEEPALIVE.
+  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))
+  if hasattr(socket, 'TCP_KEEPINTVL'):
+    options.append((socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, 60))
+  if hasattr(socket, 'TCP_KEEPCNT'):
+    options.append((socket.IPPROTO_TCP, socket.TCP_KEEPCNT, 5))
+  return options
+
+
 class SyncHttpxClient(httpx.Client):
   """Sync httpx client."""
 
   def __init__(self, **kwargs: Any) -> None:
     """Initializes the httpx client."""
     kwargs.setdefault('follow_redirects', True)
+    if 'transport' not in kwargs and 'mounts' not in kwargs:
+      kwargs['transport'] = httpx.HTTPTransport(
+          socket_options=_default_socket_options()
+      )
     super().__init__(**kwargs)
 
   def __del__(self) -> None:
@@ -584,6 +612,10 @@ class AsyncHttpxClient(httpx.AsyncClient):
   def __init__(self, **kwargs: Any) -> None:
     """Initializes the httpx client."""
     kwargs.setdefault('follow_redirects', True)
+    if 'transport' not in kwargs and 'mounts' not in kwargs:
+      kwargs['transport'] = httpx.AsyncHTTPTransport(
+          socket_options=_default_socket_options()
+      )
     super().__init__(**kwargs)
 
   def __del__(self) -> None:
diff --git a/google/genai/tests/client/test_socket_options.py b/google/genai/tests/client/test_socket_options.py
new file mode 100644
index 000000000..d1ff66eed
--- /dev/null
+++ b/google/genai/tests/client/test_socket_options.py
@@ -0,0 +1,67 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Tests for TCP keepalive on the default httpx transports.
+
+Regression tests for https://github.com/googleapis/python-genai/issues/2705:
+without TCP keepalive, long-running calls that stay silent on the wire are
+evicted by NAT gateways and stateful firewalls, hanging the read forever.
+"""
+
+import socket
+
+import httpx
+
+from ... import _api_client as api_client
+
+
+def test_default_socket_options_enable_keepalive():
+  options = api_client._default_socket_options()
+  assert (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1) in options
+  # At least one platform-specific keepalive tuning option must be present so
+  # probes start before typical NAT idle timeouts.
+  assert len(options) > 1
+
+
+def test_sync_client_gets_keepalive_transport_by_default():
+  client = api_client.SyncHttpxClient()
+  assert isinstance(client._transport, httpx.HTTPTransport)
+  assert (
+      socket.SOL_SOCKET,
+      socket.SO_KEEPALIVE,
+      1,
+  ) in client._transport._pool._socket_options
+
+
+def test_async_client_gets_keepalive_transport_by_default():
+  client = api_client.AsyncHttpxClient()
+  assert isinstance(client._transport, httpx.AsyncHTTPTransport)
+  assert (
+      socket.SOL_SOCKET,
+      socket.SO_KEEPALIVE,
+      1,
+  ) in client._transport._pool._socket_options
+
+
+def test_sync_client_preserves_user_transport():
+  transport = httpx.MockTransport(lambda request: httpx.Response(200))
+  client = api_client.SyncHttpxClient(transport=transport)
+  assert client._transport is transport
+
+
+def test_async_client_preserves_user_transport():
+  transport = httpx.MockTransport(lambda request: httpx.Response(200))
+  client = api_client.AsyncHttpxClient(transport=transport)
+  assert client._transport is transport