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