From 361015dd5128a28b10599be76c706fcdd7311f98 Mon Sep 17 00:00:00 2001 From: Weilu Jia Date: Fri, 24 Jul 2026 14:34:01 -0700 Subject: [PATCH] fix(python-sdk): enable TCP keepalive on API and envd HTTP transports Add keepalive socket options (SO_KEEPALIVE plus TCP_KEEPIDLE / macOS TCP_KEEPALIVE with a 60s initial delay where supported) and keepalive httpx transport subclasses in e2b.api, and wire them into the sync and async API, envd, and proxy transports. On Windows and platforms without a keepalive-idle constant, only SO_KEEPALIVE is set. Custom network backends force the socket options at TCP connect time, since httpcore 1.0.x drops socket_options when constructing proxy connections. Because an explicit transport disables httpx trust_env proxy handling, a shared _build_env_proxy_mounts helper rebuilds the HTTP_PROXY/HTTPS_PROXY/NO_PROXY mounts for ApiClient. Also drop a dead verify parameter from _create_transport. --- .changeset/warm-sockets-keep-python.md | 5 + packages/python-sdk/e2b/api/__init__.py | 203 +++++++++++++- .../e2b/api/client_async/__init__.py | 28 +- .../e2b/api/client_sync/__init__.py | 34 ++- packages/python-sdk/pyproject.toml | 1 + .../tests/test_api_client_transport.py | 263 +++++++++++++++++- packages/python-sdk/uv.lock | 4 +- 7 files changed, 500 insertions(+), 38 deletions(-) create mode 100644 .changeset/warm-sockets-keep-python.md diff --git a/.changeset/warm-sockets-keep-python.md b/.changeset/warm-sockets-keep-python.md new file mode 100644 index 0000000000..d8626015fe --- /dev/null +++ b/.changeset/warm-sockets-keep-python.md @@ -0,0 +1,5 @@ +--- +'@e2b/python-sdk': patch +--- + +Enable TCP keepalive with a 60-second initial delay where supported across Python SDK HTTP transports, including API, sandbox, volume, proxy, and template upload connections. Environment proxies (`HTTP_PROXY`/`HTTPS_PROXY`/`NO_PROXY`) are now honored by these transports, including template uploads. diff --git a/packages/python-sdk/e2b/api/__init__.py b/packages/python-sdk/e2b/api/__init__.py index 15ef836bfb..51cf958a80 100644 --- a/packages/python-sdk/e2b/api/__init__.py +++ b/packages/python-sdk/e2b/api/__init__.py @@ -3,15 +3,22 @@ import logging import os import re +import socket +import sys import threading import weakref from dataclasses import dataclass from types import TracebackType -from typing import Callable, Optional, Protocol, Union +from typing import Any, Callable, Iterable, Optional, Protocol, TypeVar, Union, cast +import httpcore import httpx from httpx import AsyncBaseTransport, BaseTransport, Limits, Timeout +# NOTE: httpx private API. Kept in sync with the `httpx>=0.27.0,<1.0.0` pin in +# pyproject.toml — re-verify this import when bumping httpx. +from httpx._utils import get_environment_proxies + from e2b.api.client.client import AuthenticatedClient from e2b.api.client.types import Response from e2b.api.metadata import default_headers @@ -71,6 +78,183 @@ async def on_response(response: Response) -> None: connection_retries = int(os.getenv("E2B_CONNECTION_RETRIES") or "3") +def _get_socket_options( + platform: str, + tcp_keepidle: Optional[int], + tcp_keepalive: Optional[int], +) -> tuple[tuple[int, int, int], ...]: + """Build platform-specific TCP keepalive options for httpcore. + + The 60-second initial-delay tuning uses ``TCP_KEEPIDLE`` where the + constant is available (Linux and other platforms exposing it) or macOS's + ``TCP_KEEPALIVE``. Windows CPython also defines ``TCP_KEEPIDLE``, but + setting it raises ``OSError`` on Windows releases older than 10 1709, so + Windows only enables ``SO_KEEPALIVE`` and keeps the OS-default probe + timing. Platforms without either constant likewise fall back to + enabling keepalive only. + """ + options = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)] + if platform == "win32": + return tuple(options) + if tcp_keepidle is not None: + options.append((socket.IPPROTO_TCP, tcp_keepidle, 60)) + elif platform == "darwin" and tcp_keepalive is not None: + options.append((socket.IPPROTO_TCP, tcp_keepalive, 60)) + return tuple(options) + + +_TCP_KEEPALIVE_SOCKET_OPTIONS = _get_socket_options( + sys.platform, + getattr(socket, "TCP_KEEPIDLE", None), + getattr(socket, "TCP_KEEPALIVE", None), +) + + +class _TCPKeepaliveNetworkBackend(httpcore.NetworkBackend): + """Force TCP keepalive options through direct and proxy connections.""" + + def __init__(self, backend: httpcore.NetworkBackend): + self._backend = backend + + def connect_tcp( + self, + host: str, + port: int, + timeout: Optional[float] = None, + local_address: Optional[str] = None, + socket_options: Optional[Iterable[tuple]] = None, + ) -> httpcore.NetworkStream: + return self._backend.connect_tcp( + host, + port, + timeout=timeout, + local_address=local_address, + socket_options=_TCP_KEEPALIVE_SOCKET_OPTIONS, + ) + + def connect_unix_socket( + self, + path: str, + timeout: Optional[float] = None, + socket_options: Optional[Iterable[tuple]] = None, + ) -> httpcore.NetworkStream: + return self._backend.connect_unix_socket( + path, + timeout=timeout, + socket_options=socket_options, + ) + + +class _TCPKeepaliveAsyncNetworkBackend(httpcore.AsyncNetworkBackend): + """Async counterpart of :class:`_TCPKeepaliveNetworkBackend`.""" + + def __init__(self, backend: httpcore.AsyncNetworkBackend): + self._backend = backend + + async def connect_tcp( + self, + host: str, + port: int, + timeout: Optional[float] = None, + local_address: Optional[str] = None, + socket_options: Optional[Iterable[tuple]] = None, + ) -> httpcore.AsyncNetworkStream: + return await self._backend.connect_tcp( + host, + port, + timeout=timeout, + local_address=local_address, + socket_options=_TCP_KEEPALIVE_SOCKET_OPTIONS, + ) + + async def connect_unix_socket( + self, + path: str, + timeout: Optional[float] = None, + socket_options: Optional[Iterable[tuple]] = None, + ) -> httpcore.AsyncNetworkStream: + return await self._backend.connect_unix_socket( + path, + timeout=timeout, + socket_options=socket_options, + ) + + async def sleep(self, seconds: float) -> None: + await self._backend.sleep(seconds) + + +class _TCPKeepaliveHTTPTransport(httpx.HTTPTransport): + """HTTPX transport that also covers HTTPcore proxy sockets.""" + + def __init__(self, *args, **kwargs): + kwargs["socket_options"] = _TCP_KEEPALIVE_SOCKET_OPTIONS + super().__init__(*args, **kwargs) + # NOTE: `_pool` / `_network_backend` are httpx/httpcore private API. + # HTTPcore 1.0.x drops `socket_options` when it builds proxy + # connections (HTTPProxy/SOCKSProxy), so wrap the backend to force + # the keepalive options at actual TCP connect time. Verified against + # the `httpcore>=1.0.5,<2.0.0` pin in pyproject.toml — re-check on bumps. + pool = cast(Any, self._pool) + pool._network_backend = _TCPKeepaliveNetworkBackend(pool._network_backend) + + +class _TCPKeepaliveAsyncHTTPTransport(httpx.AsyncHTTPTransport): + """Async HTTPX transport that also covers HTTPcore proxy sockets.""" + + def __init__(self, *args, **kwargs): + kwargs["socket_options"] = _TCP_KEEPALIVE_SOCKET_OPTIONS + super().__init__(*args, **kwargs) + # NOTE: `_pool` / `_network_backend` are httpx/httpcore private API — + # see _TCPKeepaliveHTTPTransport for the rationale and version pins. + pool = cast(Any, self._pool) + pool._network_backend = _TCPKeepaliveAsyncNetworkBackend(pool._network_backend) + + +class _TCPKeepaliveTransport(httpx.BaseTransport, httpx.AsyncBaseTransport): + """Dual sync/async transport for directly constructed public clients.""" + + def __init__(self, **kwargs): + self._sync_transport = _TCPKeepaliveHTTPTransport(**kwargs) + self._async_transport = _TCPKeepaliveAsyncHTTPTransport(**kwargs) + + def handle_request(self, request: httpx.Request) -> httpx.Response: + return self._sync_transport.handle_request(request) + + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: + return await self._async_transport.handle_async_request(request) + + def close(self) -> None: + self._sync_transport.close() + + async def aclose(self) -> None: + await self._async_transport.aclose() + + +# Any transport type produced by the factory passed to +# _build_env_proxy_mounts; keeps the returned mount dict precisely typed for +# both sync and async httpx clients. +_TransportT = TypeVar("_TransportT", bound=Union[BaseTransport, AsyncBaseTransport]) + + +def _build_env_proxy_mounts( + transport_factory: Callable[[str], _TransportT], +) -> dict[str, Optional[_TransportT]]: + """Rebuild httpx's environment-proxy mounts for clients that pass an + explicit ``transport=``. + + Passing an explicit transport to an httpx client disables its + ``trust_env`` proxy handling entirely, so ``HTTP_PROXY``/``HTTPS_PROXY``/ + ``ALL_PROXY`` would silently be ignored. This reconstructs the same + mounts httpx would have built, creating a proxied transport per proxy URL + via ``transport_factory``. ``NO_PROXY`` entries map to ``None`` mounts, + which httpx routes to the default (direct) transport. + """ + return { + pattern: None if proxy_url is None else transport_factory(proxy_url) + for pattern, proxy_url in get_environment_proxies().items() + } + + @dataclass class SandboxCreateResponse: sandbox_id: str @@ -222,14 +406,25 @@ def __init__( httpx_args = { "event_hooks": self._logging_event_hooks(), } - if transport is not None: - httpx_args["transport"] = transport if ( transport is None and transport_factory is None and async_transport_factory is None ): - httpx_args["proxy"] = config.proxy + transport_options = {"verify": kwargs.get("verify_ssl", True)} + transport = _TCPKeepaliveTransport( + proxy=config.proxy, + **transport_options, + ) + if config.proxy is None: + httpx_args["mounts"] = _build_env_proxy_mounts( + lambda proxy_url: _TCPKeepaliveTransport( + proxy=proxy_url, + **transport_options, + ) + ) + if transport is not None: + httpx_args["transport"] = transport # config.request_timeout is None when the timeout is explicitly # disabled (request_timeout=0), which httpx.Timeout(None) preserves. diff --git a/packages/python-sdk/e2b/api/client_async/__init__.py b/packages/python-sdk/e2b/api/client_async/__init__.py index dd001a3a41..5708444428 100644 --- a/packages/python-sdk/e2b/api/client_async/__init__.py +++ b/packages/python-sdk/e2b/api/client_async/__init__.py @@ -2,11 +2,14 @@ import weakref from typing import Dict, Optional, Tuple -import httpx - from httpx._types import ProxyTypes -from e2b.api import AsyncApiClient, connection_retries, limits +from e2b.api import ( + _TCPKeepaliveAsyncHTTPTransport, + AsyncApiClient, + connection_retries, + limits, +) from e2b.connection_config import ConnectionConfig TransportKey = Tuple[bool, Optional[ProxyTypes]] @@ -20,7 +23,7 @@ def get_api_client(config: ConnectionConfig, **kwargs) -> AsyncApiClient: ) -class AsyncTransportWithLogger(httpx.AsyncHTTPTransport): +class AsyncTransportWithLogger(_TCPKeepaliveAsyncHTTPTransport): # Keyed weakly by the event loop object itself, not id(loop) — CPython # reuses object ids, so a new loop could otherwise inherit a transport # bound to a previous, closed loop. @@ -34,6 +37,16 @@ def pool(self): return self._pool +def _create_transport(cls, config: ConnectionConfig, http2: bool): + """Build a keepalive transport of the given class for this config.""" + return cls( + limits=limits, + proxy=config.proxy, + http2=http2, + retries=connection_retries, + ) + + def _get_cached_transport(cls, config: ConnectionConfig, http2: bool): loop = asyncio.get_running_loop() loop_instances = cls._instances.get(loop) @@ -44,12 +57,7 @@ def _get_cached_transport(cls, config: ConnectionConfig, http2: bool): key: TransportKey = (http2, config.proxy) transport = loop_instances.get(key) if transport is None: - transport = cls( - limits=limits, - proxy=config.proxy, - http2=http2, - retries=connection_retries, - ) + transport = _create_transport(cls, config, http2) loop_instances[key] = transport return transport diff --git a/packages/python-sdk/e2b/api/client_sync/__init__.py b/packages/python-sdk/e2b/api/client_sync/__init__.py index 41fb4f023d..94ac02282a 100644 --- a/packages/python-sdk/e2b/api/client_sync/__init__.py +++ b/packages/python-sdk/e2b/api/client_sync/__init__.py @@ -1,11 +1,15 @@ from typing import Dict, Optional, Tuple -import httpx import threading from httpx._types import ProxyTypes -from e2b.api import ApiClient, connection_retries, limits +from e2b.api import ( + _TCPKeepaliveHTTPTransport, + ApiClient, + connection_retries, + limits, +) from e2b.connection_config import ConnectionConfig TransportKey = Tuple[bool, Optional[ProxyTypes]] @@ -19,7 +23,7 @@ def get_api_client(config: ConnectionConfig, **kwargs) -> ApiClient: ) -class TransportWithLogger(httpx.HTTPTransport): +class TransportWithLogger(_TCPKeepaliveHTTPTransport): _thread_local = threading.local() @property @@ -27,6 +31,16 @@ def pool(self): return self._pool +def _create_transport(cls, config: ConnectionConfig, http2: bool): + """Build a keepalive transport of the given class for this config.""" + return cls( + limits=limits, + proxy=config.proxy, + http2=http2, + retries=connection_retries, + ) + + def get_transport(config: ConnectionConfig, http2: bool = True) -> TransportWithLogger: instances: Dict[TransportKey, TransportWithLogger] = getattr( TransportWithLogger._thread_local, "instances", {} @@ -36,12 +50,7 @@ def get_transport(config: ConnectionConfig, http2: bool = True) -> TransportWith if cached is not None: return cached - transport = TransportWithLogger( - limits=limits, - proxy=config.proxy, - http2=http2, - retries=connection_retries, - ) + transport = _create_transport(TransportWithLogger, config, http2) instances[key] = transport TransportWithLogger._thread_local.instances = instances return transport @@ -62,12 +71,7 @@ def get_envd_transport( if cached is not None: return cached - transport = EnvdTransportWithLogger( - limits=limits, - proxy=config.proxy, - http2=http2, - retries=connection_retries, - ) + transport = _create_transport(EnvdTransportWithLogger, config, http2) instances[key] = transport EnvdTransportWithLogger._thread_local.instances = instances return transport diff --git a/packages/python-sdk/pyproject.toml b/packages/python-sdk/pyproject.toml index d6f4696bdf..c3b3600bfe 100644 --- a/packages/python-sdk/pyproject.toml +++ b/packages/python-sdk/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "wcmatch>=10.1,<11", "protobuf-py>=0.1.1,<0.2", "httpx>=0.27.0,<1.0.0", + "httpcore>=1.0.5,<2.0.0", "h2>=4,<5", "attrs>=23.2.0", "packaging>=24.1", diff --git a/packages/python-sdk/tests/test_api_client_transport.py b/packages/python-sdk/tests/test_api_client_transport.py index e54848e98a..a7287d415d 100644 --- a/packages/python-sdk/tests/test_api_client_transport.py +++ b/packages/python-sdk/tests/test_api_client_transport.py @@ -1,10 +1,21 @@ import asyncio import gc +import socket +import ssl from concurrent.futures import ThreadPoolExecutor +from unittest.mock import AsyncMock, Mock import httpx import pytest +from e2b.api import ( + _TCP_KEEPALIVE_SOCKET_OPTIONS, + _TCPKeepaliveAsyncNetworkBackend, + _TCPKeepaliveNetworkBackend, + _TCPKeepaliveTransport, + ApiClient, + _get_socket_options, +) from e2b.api.client_async import AsyncEnvdTransportWithLogger, AsyncTransportWithLogger from e2b.api.client_async import get_api_client as get_async_api_client from e2b.api.client_async import get_envd_transport as get_async_envd_transport @@ -16,6 +27,67 @@ from e2b.connection_config import ConnectionConfig +def test_socket_options_use_tcp_keepidle_when_available(): + assert _get_socket_options("linux", tcp_keepidle=4, tcp_keepalive=None) == ( + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), + (socket.IPPROTO_TCP, 4, 60), + ) + + +def test_socket_options_use_macos_tcp_keepalive_fallback(): + assert _get_socket_options("darwin", tcp_keepidle=None, tcp_keepalive=16) == ( + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), + (socket.IPPROTO_TCP, 16, 60), + ) + + +def test_socket_options_fall_back_to_enabling_keepalive_only(): + assert _get_socket_options("win32", tcp_keepidle=None, tcp_keepalive=16) == ( + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), + ) + + +def test_socket_options_windows_ignores_tcp_keepidle(): + # Windows CPython defines TCP_KEEPIDLE, but setting it raises OSError on + # Windows releases older than 10 1709, so Windows stays enable-only with + # OS-default probe timing. + assert _get_socket_options("win32", tcp_keepidle=4, tcp_keepalive=None) == ( + (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1), + ) + + +def test_sync_network_backend_forces_keepalive_options(): + backend = Mock() + wrapped = _TCPKeepaliveNetworkBackend(backend) + + wrapped.connect_tcp("example.com", 443, socket_options=None) + + backend.connect_tcp.assert_called_once_with( + "example.com", + 443, + timeout=None, + local_address=None, + socket_options=_TCP_KEEPALIVE_SOCKET_OPTIONS, + ) + + +@pytest.mark.asyncio +async def test_async_network_backend_forces_keepalive_options(): + backend = Mock() + backend.connect_tcp = AsyncMock() + wrapped = _TCPKeepaliveAsyncNetworkBackend(backend) + + await wrapped.connect_tcp("example.com", 443, socket_options=None) + + backend.connect_tcp.assert_awaited_once_with( + "example.com", + 443, + timeout=None, + local_address=None, + socket_options=_TCP_KEEPALIVE_SOCKET_OPTIONS, + ) + + def reset_sync_api_transports(): TransportWithLogger._thread_local.instances = {} @@ -29,6 +101,25 @@ def run_in_worker_thread(fn): return executor.submit(fn).result() +def get_transport_option(transport, name): + return getattr(getattr(transport, "_pool"), name) + + +def set_only_https_proxy(monkeypatch, proxy): + for name in ( + "ALL_PROXY", + "HTTP_PROXY", + "HTTPS_PROXY", + "NO_PROXY", + "all_proxy", + "http_proxy", + "https_proxy", + "no_proxy", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.setenv("HTTPS_PROXY", proxy) + + def test_sync_api_client_proxy_uses_explicit_transport(test_api_key): reset_sync_api_transports() config = ConnectionConfig( @@ -43,11 +134,82 @@ def test_sync_api_client_proxy_uses_explicit_transport(test_api_key): assert "proxy" not in api_client._httpx_args assert httpx_client._transport is get_sync_transport(config) assert httpx_client._mounts == {} + assert ( + get_transport_option(httpx_client._transport, "_socket_options") + == _TCP_KEEPALIVE_SOCKET_OPTIONS + ) + assert isinstance( + get_transport_option(httpx_client._transport, "_network_backend"), + _TCPKeepaliveNetworkBackend, + ) finally: httpx_client.close() reset_sync_api_transports() +def test_direct_api_client_uses_owned_sync_keepalive_transport(test_api_key): + reset_sync_api_transports() + config = ConnectionConfig( + api_key=test_api_key, + proxy="http://127.0.0.1:9999", + ) + api_client = ApiClient(config, verify_ssl=False) + httpx_client = api_client.get_httpx_client() + + try: + assert isinstance(httpx_client._transport, _TCPKeepaliveTransport) + transport = httpx_client._transport._sync_transport + assert get_transport_option(transport, "_http2") is False + assert ( + get_transport_option(transport, "_socket_options") + == _TCP_KEEPALIVE_SOCKET_OPTIONS + ) + assert ( + get_transport_option(transport, "_ssl_context").verify_mode == ssl.CERT_NONE + ) + finally: + httpx_client.close() + reset_sync_api_transports() + + +def test_direct_api_client_preserves_explicit_sync_transport(test_api_key): + transport = httpx.MockTransport(lambda request: httpx.Response(200)) + api_client = ApiClient(ConnectionConfig(api_key=test_api_key), transport=transport) + + try: + assert api_client.get_httpx_client()._transport is transport + finally: + api_client.get_httpx_client().close() + + +def test_direct_api_client_preserves_set_httpx_client(test_api_key): + user_client = httpx.Client(transport=httpx.MockTransport(lambda request: None)) + api_client = ApiClient(ConnectionConfig(api_key=test_api_key)) + api_client.set_httpx_client(user_client) + + try: + assert api_client.get_httpx_client() is user_client + finally: + user_client.close() + + +def test_direct_api_client_preserves_environment_proxy(test_api_key, monkeypatch): + set_only_https_proxy(monkeypatch, "http://127.0.0.1:9999") + api_client = ApiClient(ConnectionConfig(api_key=test_api_key), verify_ssl=False) + httpx_client = api_client.get_httpx_client() + + try: + transport = httpx_client._transport_for_url(httpx.URL("https://example.com")) + assert isinstance(transport, _TCPKeepaliveTransport) + assert transport is not httpx_client._transport + assert ( + get_transport_option(transport._sync_transport, "_proxy_url").origin.host + == b"127.0.0.1" + ) + finally: + httpx_client.close() + + def test_sync_get_transport_http2_opt_out_returns_distinct_instance(test_api_key): reset_sync_api_transports() config = ConnectionConfig(api_key=test_api_key) @@ -57,8 +219,8 @@ def test_sync_get_transport_http2_opt_out_returns_distinct_instance(test_api_key http1_transport = get_sync_transport(config, http2=False) assert http2_transport is not http1_transport - assert http2_transport._pool._http2 is True - assert http1_transport._pool._http2 is False + assert get_transport_option(http2_transport, "_http2") is True + assert get_transport_option(http1_transport, "_http2") is False # Subsequent calls with the same http2 flag return the cached # instance. assert get_sync_transport(config) is http2_transport @@ -134,7 +296,11 @@ def test_sync_envd_transport_uses_separate_cache(test_api_key): assert api_transport is not envd_transport assert get_sync_transport(config) is api_transport assert get_sync_envd_transport(config) is envd_transport - assert envd_transport._pool._http2 is True + assert get_transport_option(envd_transport, "_http2") is True + assert ( + get_transport_option(envd_transport, "_socket_options") + == _TCP_KEEPALIVE_SOCKET_OPTIONS + ) finally: reset_sync_api_transports() reset_sync_envd_transports() @@ -207,8 +373,8 @@ def test_sync_envd_transport_cache_is_thread_local(test_api_key): assert main_transport is get_sync_envd_transport(config) assert thread_transport is not main_transport - assert main_transport._pool._http2 is True - assert thread_transport._pool._http2 is True + assert get_transport_option(main_transport, "_http2") is True + assert get_transport_option(thread_transport, "_http2") is True finally: reset_sync_envd_transports() @@ -231,11 +397,88 @@ async def test_async_api_client_proxy_uses_explicit_transport(test_api_key): assert "proxy" not in api_client._httpx_args assert httpx_client._transport is transport assert httpx_client._mounts == {} + assert ( + get_transport_option(httpx_client._transport, "_socket_options") + == _TCP_KEEPALIVE_SOCKET_OPTIONS + ) + assert isinstance( + get_transport_option(httpx_client._transport, "_network_backend"), + _TCPKeepaliveAsyncNetworkBackend, + ) finally: await httpx_client.aclose() AsyncTransportWithLogger._instances.clear() +@pytest.mark.asyncio +async def test_direct_api_client_uses_owned_async_keepalive_transport(test_api_key): + AsyncTransportWithLogger._instances.clear() + config = ConnectionConfig( + api_key=test_api_key, + proxy="http://127.0.0.1:9999", + ) + api_client = ApiClient(config, verify_ssl=False) + httpx_client = api_client.get_async_httpx_client() + + try: + assert isinstance(httpx_client._transport, _TCPKeepaliveTransport) + transport = httpx_client._transport._async_transport + assert get_transport_option(transport, "_http2") is False + assert ( + get_transport_option(transport, "_socket_options") + == _TCP_KEEPALIVE_SOCKET_OPTIONS + ) + assert ( + get_transport_option(transport, "_ssl_context").verify_mode == ssl.CERT_NONE + ) + finally: + await httpx_client.aclose() + AsyncTransportWithLogger._instances.clear() + + +@pytest.mark.asyncio +async def test_direct_api_client_preserves_explicit_async_transport(test_api_key): + transport = httpx.MockTransport(lambda request: httpx.Response(200)) + api_client = ApiClient(ConnectionConfig(api_key=test_api_key), transport=transport) + + try: + assert api_client.get_async_httpx_client()._transport is transport + finally: + await api_client.get_async_httpx_client().aclose() + + +@pytest.mark.asyncio +async def test_direct_api_client_preserves_set_async_httpx_client(test_api_key): + user_client = httpx.AsyncClient(transport=httpx.MockTransport(lambda request: None)) + api_client = ApiClient(ConnectionConfig(api_key=test_api_key)) + api_client.set_async_httpx_client(user_client) + + try: + assert api_client.get_async_httpx_client() is user_client + finally: + await user_client.aclose() + + +@pytest.mark.asyncio +async def test_direct_api_client_preserves_async_environment_proxy( + test_api_key, monkeypatch +): + set_only_https_proxy(monkeypatch, "http://127.0.0.1:9999") + api_client = ApiClient(ConnectionConfig(api_key=test_api_key), verify_ssl=False) + httpx_client = api_client.get_async_httpx_client() + + try: + transport = httpx_client._transport_for_url(httpx.URL("https://example.com")) + assert isinstance(transport, _TCPKeepaliveTransport) + assert transport is not httpx_client._transport + assert ( + get_transport_option(transport._async_transport, "_proxy_url").origin.host + == b"127.0.0.1" + ) + finally: + await httpx_client.aclose() + + @pytest.mark.asyncio async def test_async_get_transport_http2_opt_out_returns_distinct_instance( test_api_key, @@ -248,8 +491,8 @@ async def test_async_get_transport_http2_opt_out_returns_distinct_instance( http1_transport = get_async_transport(config, http2=False) assert http2_transport is not http1_transport - assert http2_transport._pool._http2 is True - assert http1_transport._pool._http2 is False + assert get_transport_option(http2_transport, "_http2") is True + assert get_transport_option(http1_transport, "_http2") is False # Subsequent calls with the same http2 flag return the cached # instance. assert get_async_transport(config) is http2_transport @@ -415,7 +658,11 @@ async def test_async_envd_transport_uses_separate_cache(test_api_key): assert api_transport is not envd_transport assert get_async_transport(config) is api_transport assert get_async_envd_transport(config) is envd_transport - assert envd_transport._pool._http2 is True + assert get_transport_option(envd_transport, "_http2") is True + assert ( + get_transport_option(envd_transport, "_socket_options") + == _TCP_KEEPALIVE_SOCKET_OPTIONS + ) finally: AsyncTransportWithLogger._instances.clear() AsyncEnvdTransportWithLogger._instances.clear() diff --git a/packages/python-sdk/uv.lock b/packages/python-sdk/uv.lock index efdf1a50ff..574b62c9f6 100644 --- a/packages/python-sdk/uv.lock +++ b/packages/python-sdk/uv.lock @@ -187,6 +187,7 @@ dependencies = [ { name = "connectrpc" }, { name = "dockerfile-parse" }, { name = "h2" }, + { name = "httpcore" }, { name = "httpx" }, { name = "packaging" }, { name = "protobuf-py" }, @@ -222,6 +223,7 @@ requires-dist = [ { name = "connectrpc", specifier = ">=0.11.1,<0.12" }, { name = "dockerfile-parse", specifier = ">=2.0.1,<3" }, { name = "h2", specifier = ">=4,<5" }, + { name = "httpcore", specifier = ">=1.0.5,<2.0.0" }, { name = "httpx", specifier = ">=0.27.0,<1.0.0" }, { name = "packaging", specifier = ">=24.1" }, { name = "protobuf-py", specifier = ">=0.1.1,<0.2" }, @@ -278,7 +280,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [