|
3 | 3 | import logging |
4 | 4 | import os |
5 | 5 | import re |
| 6 | +import socket |
| 7 | +import sys |
6 | 8 | import threading |
7 | 9 | import weakref |
8 | 10 | from dataclasses import dataclass |
9 | 11 | from types import TracebackType |
10 | | -from typing import Callable, Optional, Protocol, Union |
| 12 | +from typing import Any, Callable, Iterable, Optional, Protocol, TypeVar, Union, cast |
11 | 13 |
|
| 14 | +import httpcore |
12 | 15 | import httpx |
13 | 16 | from httpx import AsyncBaseTransport, BaseTransport, Limits, Timeout |
14 | 17 |
|
| 18 | +# NOTE: httpx private API. Kept in sync with the `httpx>=0.27.0,<1.0.0` pin in |
| 19 | +# pyproject.toml — re-verify this import when bumping httpx. |
| 20 | +from httpx._utils import get_environment_proxies |
| 21 | + |
15 | 22 | from e2b.api.client.client import AuthenticatedClient |
16 | 23 | from e2b.api.client.types import Response |
17 | 24 | from e2b.api.metadata import default_headers |
@@ -71,6 +78,183 @@ async def on_response(response: Response) -> None: |
71 | 78 | connection_retries = int(os.getenv("E2B_CONNECTION_RETRIES") or "3") |
72 | 79 |
|
73 | 80 |
|
| 81 | +def _get_socket_options( |
| 82 | + platform: str, |
| 83 | + tcp_keepidle: Optional[int], |
| 84 | + tcp_keepalive: Optional[int], |
| 85 | +) -> tuple[tuple[int, int, int], ...]: |
| 86 | + """Build platform-specific TCP keepalive options for httpcore. |
| 87 | +
|
| 88 | + The 60-second initial-delay tuning uses ``TCP_KEEPIDLE`` where the |
| 89 | + constant is available (Linux and other platforms exposing it) or macOS's |
| 90 | + ``TCP_KEEPALIVE``. Windows CPython also defines ``TCP_KEEPIDLE``, but |
| 91 | + setting it raises ``OSError`` on Windows releases older than 10 1709, so |
| 92 | + Windows only enables ``SO_KEEPALIVE`` and keeps the OS-default probe |
| 93 | + timing. Platforms without either constant likewise fall back to |
| 94 | + enabling keepalive only. |
| 95 | + """ |
| 96 | + options = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)] |
| 97 | + if platform == "win32": |
| 98 | + return tuple(options) |
| 99 | + if tcp_keepidle is not None: |
| 100 | + options.append((socket.IPPROTO_TCP, tcp_keepidle, 60)) |
| 101 | + elif platform == "darwin" and tcp_keepalive is not None: |
| 102 | + options.append((socket.IPPROTO_TCP, tcp_keepalive, 60)) |
| 103 | + return tuple(options) |
| 104 | + |
| 105 | + |
| 106 | +_TCP_KEEPALIVE_SOCKET_OPTIONS = _get_socket_options( |
| 107 | + sys.platform, |
| 108 | + getattr(socket, "TCP_KEEPIDLE", None), |
| 109 | + getattr(socket, "TCP_KEEPALIVE", None), |
| 110 | +) |
| 111 | + |
| 112 | + |
| 113 | +class _TCPKeepaliveNetworkBackend(httpcore.NetworkBackend): |
| 114 | + """Force TCP keepalive options through direct and proxy connections.""" |
| 115 | + |
| 116 | + def __init__(self, backend: httpcore.NetworkBackend): |
| 117 | + self._backend = backend |
| 118 | + |
| 119 | + def connect_tcp( |
| 120 | + self, |
| 121 | + host: str, |
| 122 | + port: int, |
| 123 | + timeout: Optional[float] = None, |
| 124 | + local_address: Optional[str] = None, |
| 125 | + socket_options: Optional[Iterable[tuple]] = None, |
| 126 | + ) -> httpcore.NetworkStream: |
| 127 | + return self._backend.connect_tcp( |
| 128 | + host, |
| 129 | + port, |
| 130 | + timeout=timeout, |
| 131 | + local_address=local_address, |
| 132 | + socket_options=_TCP_KEEPALIVE_SOCKET_OPTIONS, |
| 133 | + ) |
| 134 | + |
| 135 | + def connect_unix_socket( |
| 136 | + self, |
| 137 | + path: str, |
| 138 | + timeout: Optional[float] = None, |
| 139 | + socket_options: Optional[Iterable[tuple]] = None, |
| 140 | + ) -> httpcore.NetworkStream: |
| 141 | + return self._backend.connect_unix_socket( |
| 142 | + path, |
| 143 | + timeout=timeout, |
| 144 | + socket_options=socket_options, |
| 145 | + ) |
| 146 | + |
| 147 | + |
| 148 | +class _TCPKeepaliveAsyncNetworkBackend(httpcore.AsyncNetworkBackend): |
| 149 | + """Async counterpart of :class:`_TCPKeepaliveNetworkBackend`.""" |
| 150 | + |
| 151 | + def __init__(self, backend: httpcore.AsyncNetworkBackend): |
| 152 | + self._backend = backend |
| 153 | + |
| 154 | + async def connect_tcp( |
| 155 | + self, |
| 156 | + host: str, |
| 157 | + port: int, |
| 158 | + timeout: Optional[float] = None, |
| 159 | + local_address: Optional[str] = None, |
| 160 | + socket_options: Optional[Iterable[tuple]] = None, |
| 161 | + ) -> httpcore.AsyncNetworkStream: |
| 162 | + return await self._backend.connect_tcp( |
| 163 | + host, |
| 164 | + port, |
| 165 | + timeout=timeout, |
| 166 | + local_address=local_address, |
| 167 | + socket_options=_TCP_KEEPALIVE_SOCKET_OPTIONS, |
| 168 | + ) |
| 169 | + |
| 170 | + async def connect_unix_socket( |
| 171 | + self, |
| 172 | + path: str, |
| 173 | + timeout: Optional[float] = None, |
| 174 | + socket_options: Optional[Iterable[tuple]] = None, |
| 175 | + ) -> httpcore.AsyncNetworkStream: |
| 176 | + return await self._backend.connect_unix_socket( |
| 177 | + path, |
| 178 | + timeout=timeout, |
| 179 | + socket_options=socket_options, |
| 180 | + ) |
| 181 | + |
| 182 | + async def sleep(self, seconds: float) -> None: |
| 183 | + await self._backend.sleep(seconds) |
| 184 | + |
| 185 | + |
| 186 | +class _TCPKeepaliveHTTPTransport(httpx.HTTPTransport): |
| 187 | + """HTTPX transport that also covers HTTPcore proxy sockets.""" |
| 188 | + |
| 189 | + def __init__(self, *args, **kwargs): |
| 190 | + kwargs["socket_options"] = _TCP_KEEPALIVE_SOCKET_OPTIONS |
| 191 | + super().__init__(*args, **kwargs) |
| 192 | + # NOTE: `_pool` / `_network_backend` are httpx/httpcore private API. |
| 193 | + # HTTPcore 1.0.x drops `socket_options` when it builds proxy |
| 194 | + # connections (HTTPProxy/SOCKSProxy), so wrap the backend to force |
| 195 | + # the keepalive options at actual TCP connect time. Verified against |
| 196 | + # the `httpcore>=1.0.5,<2.0.0` pin in pyproject.toml — re-check on bumps. |
| 197 | + pool = cast(Any, self._pool) |
| 198 | + pool._network_backend = _TCPKeepaliveNetworkBackend(pool._network_backend) |
| 199 | + |
| 200 | + |
| 201 | +class _TCPKeepaliveAsyncHTTPTransport(httpx.AsyncHTTPTransport): |
| 202 | + """Async HTTPX transport that also covers HTTPcore proxy sockets.""" |
| 203 | + |
| 204 | + def __init__(self, *args, **kwargs): |
| 205 | + kwargs["socket_options"] = _TCP_KEEPALIVE_SOCKET_OPTIONS |
| 206 | + super().__init__(*args, **kwargs) |
| 207 | + # NOTE: `_pool` / `_network_backend` are httpx/httpcore private API — |
| 208 | + # see _TCPKeepaliveHTTPTransport for the rationale and version pins. |
| 209 | + pool = cast(Any, self._pool) |
| 210 | + pool._network_backend = _TCPKeepaliveAsyncNetworkBackend(pool._network_backend) |
| 211 | + |
| 212 | + |
| 213 | +class _TCPKeepaliveTransport(httpx.BaseTransport, httpx.AsyncBaseTransport): |
| 214 | + """Dual sync/async transport for directly constructed public clients.""" |
| 215 | + |
| 216 | + def __init__(self, **kwargs): |
| 217 | + self._sync_transport = _TCPKeepaliveHTTPTransport(**kwargs) |
| 218 | + self._async_transport = _TCPKeepaliveAsyncHTTPTransport(**kwargs) |
| 219 | + |
| 220 | + def handle_request(self, request: httpx.Request) -> httpx.Response: |
| 221 | + return self._sync_transport.handle_request(request) |
| 222 | + |
| 223 | + async def handle_async_request(self, request: httpx.Request) -> httpx.Response: |
| 224 | + return await self._async_transport.handle_async_request(request) |
| 225 | + |
| 226 | + def close(self) -> None: |
| 227 | + self._sync_transport.close() |
| 228 | + |
| 229 | + async def aclose(self) -> None: |
| 230 | + await self._async_transport.aclose() |
| 231 | + |
| 232 | + |
| 233 | +# Any transport type produced by the factory passed to |
| 234 | +# _build_env_proxy_mounts; keeps the returned mount dict precisely typed for |
| 235 | +# both sync and async httpx clients. |
| 236 | +_TransportT = TypeVar("_TransportT", bound=Union[BaseTransport, AsyncBaseTransport]) |
| 237 | + |
| 238 | + |
| 239 | +def _build_env_proxy_mounts( |
| 240 | + transport_factory: Callable[[str], _TransportT], |
| 241 | +) -> dict[str, Optional[_TransportT]]: |
| 242 | + """Rebuild httpx's environment-proxy mounts for clients that pass an |
| 243 | + explicit ``transport=``. |
| 244 | +
|
| 245 | + Passing an explicit transport to an httpx client disables its |
| 246 | + ``trust_env`` proxy handling entirely, so ``HTTP_PROXY``/``HTTPS_PROXY``/ |
| 247 | + ``ALL_PROXY`` would silently be ignored. This reconstructs the same |
| 248 | + mounts httpx would have built, creating a proxied transport per proxy URL |
| 249 | + via ``transport_factory``. ``NO_PROXY`` entries map to ``None`` mounts, |
| 250 | + which httpx routes to the default (direct) transport. |
| 251 | + """ |
| 252 | + return { |
| 253 | + pattern: None if proxy_url is None else transport_factory(proxy_url) |
| 254 | + for pattern, proxy_url in get_environment_proxies().items() |
| 255 | + } |
| 256 | + |
| 257 | + |
74 | 258 | @dataclass |
75 | 259 | class SandboxCreateResponse: |
76 | 260 | sandbox_id: str |
@@ -222,14 +406,25 @@ def __init__( |
222 | 406 | httpx_args = { |
223 | 407 | "event_hooks": self._logging_event_hooks(), |
224 | 408 | } |
225 | | - if transport is not None: |
226 | | - httpx_args["transport"] = transport |
227 | 409 | if ( |
228 | 410 | transport is None |
229 | 411 | and transport_factory is None |
230 | 412 | and async_transport_factory is None |
231 | 413 | ): |
232 | | - httpx_args["proxy"] = config.proxy |
| 414 | + transport_options = {"verify": kwargs.get("verify_ssl", True)} |
| 415 | + transport = _TCPKeepaliveTransport( |
| 416 | + proxy=config.proxy, |
| 417 | + **transport_options, |
| 418 | + ) |
| 419 | + if config.proxy is None: |
| 420 | + httpx_args["mounts"] = _build_env_proxy_mounts( |
| 421 | + lambda proxy_url: _TCPKeepaliveTransport( |
| 422 | + proxy=proxy_url, |
| 423 | + **transport_options, |
| 424 | + ) |
| 425 | + ) |
| 426 | + if transport is not None: |
| 427 | + httpx_args["transport"] = transport |
233 | 428 |
|
234 | 429 | # config.request_timeout is None when the timeout is explicitly |
235 | 430 | # disabled (request_timeout=0), which httpx.Timeout(None) preserves. |
|
0 commit comments