From ad9eb51f1fbf6a4aeea066e0b92d1f11f0843a7c Mon Sep 17 00:00:00 2001 From: Mish Ushakov <10400064+mishushakov@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:56:24 +0200 Subject: [PATCH 1/2] refactor(python-sdk): unify the pyqwest connection pools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every persistent HTTP stack in the SDK now draws its connection pool from `e2b.api.client_sync`/`client_async`, keyed on (proxy, idle read bound), instead of caching four of its own: the control-plane REST API, the envd HTTP API, the envd RPC clients, and the volume content API. reqwest pools per host internally, so one pool serves the API host and every per-sandbox host without interference — and because envd RPC and the envd HTTP API hit the same host, an active sandbox needs a single HTTP/2 connection instead of one per stack. Two accessors expose it: `get_pyqwest_transport` hands connectrpc the pool behind the connect-only retries, and `get_httpx_transport` hands the generated httpx clients the `PyqwestTransport` adapter over that same pool. Layers above stay per-consumer, as the design calls for: `PlainHTTPErrorTransport` is now a stateless per-client wrapper rather than a cached transport, so Connect-error normalization stays RPC-only. Streamed downloads keep a pool of their own — the only one carrying the idle `read_timeout`, since reqwest's read timer runs during body send and TTFB and would otherwise cut off long uploads. Sharing puts the sandbox health probe on the connection the failed RPC was using, so `tests/test_shared_transport_pool.py` pins that at the frame level with a new multi-connection HTTP/2 server serving both routes on one pool: an RST_STREAM kills only the stream and the probe reuses the same connection, while a dropped TCP connection makes reqwest redial. Both paths still answer, so `handle_rpc_exception_with_health` keeps telling a wedged connection apart from a dead sandbox. Co-Authored-By: Claude Opus 5 (1M context) --- .../python-unify-pyqwest-connection-pools.md | 16 ++ .../e2b/api/client_async/__init__.py | 181 ++++++++++-------- .../e2b/api/client_sync/__init__.py | 180 +++++++++-------- .../e2b/envd/client_async/__init__.py | 36 ++-- packages/python-sdk/e2b/envd/client_shared.py | 17 +- .../e2b/envd/client_sync/__init__.py | 40 ++-- .../sandbox_async/filesystem/filesystem.py | 4 +- .../e2b/sandbox_sync/filesystem/filesystem.py | 4 +- .../e2b/volume/client_async/__init__.py | 62 ++---- .../e2b/volume/client_sync/__init__.py | 59 ++---- .../python-sdk/tests/envd_frame_server.py | 136 ++++++++++++- .../tests/test_api_client_transport.py | 169 +++++++--------- .../tests/test_envd_client_transport.py | 109 +++++++---- .../tests/test_shared_transport_pool.py | 158 +++++++++++++++ .../python-sdk/tests/test_volume_client.py | 32 +++- packages/python-sdk/tests/transport_caches.py | 19 ++ 16 files changed, 764 insertions(+), 458 deletions(-) create mode 100644 .changeset/python-unify-pyqwest-connection-pools.md create mode 100644 packages/python-sdk/tests/test_shared_transport_pool.py create mode 100644 packages/python-sdk/tests/transport_caches.py diff --git a/.changeset/python-unify-pyqwest-connection-pools.md b/.changeset/python-unify-pyqwest-connection-pools.md new file mode 100644 index 0000000000..ec3655cd00 --- /dev/null +++ b/.changeset/python-unify-pyqwest-connection-pools.md @@ -0,0 +1,16 @@ +--- +"@e2b/python-sdk": patch +--- + +Run every persistent HTTP stack in the SDK on one shared pyqwest connection pool +per proxy instead of four: the control-plane REST API, the envd HTTP API, the +envd RPC clients, and the volume content API now all draw from +`e2b.api.client_sync`/`client_async`. reqwest pools per host internally, so one +pool serves the API host and every per-sandbox host without interference — and +since envd RPC and the envd HTTP API hit the same host, an active sandbox now +needs a single HTTP/2 connection instead of one per stack. Streamed downloads +keep a pool of their own, the only one carrying the idle `read_timeout`: +reqwest's read timer runs during body send and TTFB, so on a shared pool it +would cut off long uploads. Layers above the pool are unchanged — the RPC +stack's plain-HTTP error normalization still wraps it, and connect-only retries +are now part of it. diff --git a/packages/python-sdk/e2b/api/client_async/__init__.py b/packages/python-sdk/e2b/api/client_async/__init__.py index d40ca257e6..45f7bc1cc9 100644 --- a/packages/python-sdk/e2b/api/client_async/__init__.py +++ b/packages/python-sdk/e2b/api/client_async/__init__.py @@ -24,14 +24,15 @@ def get_api_client(config: ConnectionConfig, **kwargs) -> AsyncApiClient: class ConnectionRetryTransport(RetryTransport): - """Retry only failures establishing the connection — shared by the REST - API and envd RPC stacks: pyqwest raises the builtin ``ConnectionError`` - only before the request was written, so these retries can never replay a - request the server may have received (a delivered REST call or unary RPC - like ``SendInput``). This matches the connect-only ``retries`` of the - httpx transports this replaced; the retry middleware's default policy - would otherwise also retry I/O errors and 429/5xx responses for - idempotent methods.""" + """Retry only failures establishing the connection — part of the shared + transport stack, so it covers the REST API, the envd HTTP API, the envd RPC + clients and the volume content API alike: pyqwest raises the builtin + ``ConnectionError`` only before the request was written, so these retries + can never replay a request the server may have received (a delivered REST + call or unary RPC like ``SendInput``). This matches the connect-only + ``retries`` of the httpx transports this replaced; the retry middleware's + default policy would otherwise also retry I/O errors and 429/5xx responses + for idempotent methods.""" def should_retry_response( self, request: Request, response: Union[Response, Exception] @@ -39,104 +40,122 @@ def should_retry_response( return isinstance(response, ConnectionError) -def retrying_http_transport( +TransportKey = Tuple[Optional[ProxyConfig], Optional[float]] +"""Cache key of the shared transports: the proxy and the idle read bound.""" + +_transport_lock = threading.Lock() +# One pyqwest transport — one reqwest connection pool — per proxy and idle read +# bound; `None` is the direct pool. Every HTTP stack in the SDK draws from +# these: the control-plane REST API, the envd HTTP API, the envd RPC clients +# (`e2b.envd.client_async`) and the volume content API +# (`e2b.volume.client_async`). reqwest pools per host internally, so a single +# pool serves the API host and every per-sandbox host without interference — +# and because envd RPC and the envd HTTP API share it, a sandbox needs one +# HTTP/2 connection instead of one per stack. +# +# pyqwest's I/O runs on its own Rust runtime, so unlike the httpx transports +# they replaced, the transports are not bound to an event loop and the caches +# are process-global rather than per-loop. +_transports: Dict[TransportKey, ConnectionRetryTransport] = {} +# The httpx adapter over each pool, shared by every httpx client on it. +_httpx_transports: Dict[TransportKey, AsyncPyqwestTransport] = {} + + +def get_pyqwest_transport( proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None ) -> ConnectionRetryTransport: - """A fresh pyqwest transport (= its own connection pool) with the SDK's - shared tuning — system CA certs (without which TLS through an - intercepting proxy fails), the httpx-equivalent pool limits, and - connect-only retries. The REST API, envd RPC, and envd HTTP API stacks - each cache their own instances (pool unification is a follow-up). - - ``read_timeout`` bounds every read on the transport's connections; see - :func:`get_envd_transport` for when that is (and isn't) appropriate. + """The shared pyqwest transport (= one connection pool) with the SDK's + tuning — system CA certs (without which TLS through an intercepting proxy + fails) and the httpx-equivalent pool limits — behind connect-only retries. + For TLS connections ALPN negotiates the HTTP version (HTTP/2 against the + E2B API and envd), like the http2-enabled httpx transports this replaced. + + Consumers speaking pyqwest natively (the envd RPC clients) take this; + consumers speaking httpx take :func:`get_httpx_transport`, the adapter over + the very same pool. Layer concerns above it rather than into it — the RPC + stack's plain-HTTP-error normalization wraps it, headers and codecs are + per-request — so that the pool stays shareable. + + ``read_timeout`` bounds every read on the pool's connections, which is why + it is part of the cache key: reqwest's read timer keeps running while a + request body is sent and while waiting for the response head, so a pool + carrying one would cut off long uploads and slow responses. Only streamed + downloads ask for it, as an idle bound (see :func:`get_transport`). Requests are logged by pyqwest itself on the ``pyqwest.access`` and ``pyqwest`` loggers at ``DEBUG`` (off unless enabled) — the transport-level diagnostics httpcore used to provide. The SDK's own ``logger`` option is separate and sits above this, on the httpx client.""" - return ConnectionRetryTransport( - HTTPTransport( - tls_include_system_certs=True, - proxy=proxy.to_pyqwest() if proxy is not None else None, - pool_idle_timeout=pool_idle_timeout, - pool_max_idle_per_host=pool_max_idle_per_host, - read_timeout=read_timeout, - # Redirects belong to the httpx client above (which the generated - # clients leave off), not to reqwest. - follow_redirects=False, - ), - max_retries=connection_retries, - ) - - -_transport_lock = threading.Lock() -# One transport (= one connection pool) per proxy; None is the direct pool. -# pyqwest's I/O runs on its own Rust runtime, so unlike the httpx transports -# they replaced, the transports are not bound to an event loop and the -# caches are process-global rather than per-loop. -_transports: Dict[Optional[ProxyConfig], AsyncPyqwestTransport] = {} - - -def get_transport(config: ConnectionConfig) -> AsyncPyqwestTransport: - """The shared pyqwest-backed httpx transport for REST API calls. For TLS - connections ALPN negotiates the HTTP version (HTTP/2 against the E2B - API), like the http2-enabled httpx transport this replaced.""" - proxy = proxy_to_config(config.proxy) + key = (proxy, read_timeout) with _transport_lock: - transport = _transports.get(proxy) + transport = _transports.get(key) if transport is None: - transport = AsyncPyqwestTransport(retrying_http_transport(proxy)) - _transports[proxy] = transport + transport = ConnectionRetryTransport( + HTTPTransport( + tls_include_system_certs=True, + proxy=proxy.to_pyqwest() if proxy is not None else None, + pool_idle_timeout=pool_idle_timeout, + pool_max_idle_per_host=pool_max_idle_per_host, + read_timeout=read_timeout, + # Redirects belong to the httpx client above (which the + # generated clients leave off), not to reqwest. + follow_redirects=False, + ), + max_retries=connection_retries, + ) + _transports[key] = transport return transport -# One transport per (proxy, streaming) pair, separate from the REST API -# pools — envd traffic goes to per-sandbox hosts. -_envd_transports: Dict[Tuple[Optional[ProxyConfig], bool], AsyncPyqwestTransport] = {} +def get_httpx_transport( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> AsyncPyqwestTransport: + """The httpx adapter over the shared pool of + :func:`get_pyqwest_transport`, for the generated httpx clients (control + plane, envd HTTP API, volume content). The adapter holds no state of its + own and does not close the pool, so closing an httpx client leaves the + pool intact for the other clients on it.""" + key = (proxy, read_timeout) + # Resolve the pool before taking the lock: it takes the same one. + pool = get_pyqwest_transport(proxy, read_timeout) + with _transport_lock: + transport = _httpx_transports.get(key) + if transport is None: + transport = AsyncPyqwestTransport(pool) + _httpx_transports[key] = transport + return transport -def get_envd_transport( +def get_transport( config: ConnectionConfig, *, for_streaming: bool = False ) -> AsyncPyqwestTransport: - """The shared pyqwest-backed httpx transports for the envd HTTP API - (file transfers, health checks). - - The streaming transport carries ``read_timeout``, the idle bound on - every read: it resets after each successful read, so it caps how long a - streamed download may stall without limiting total transfer time. It is - fixed per transport — the adapter's per-request timeouts are - whole-request deadlines. Only streamed downloads use it: reqwest's read - timer keeps running while a request body is sent and while waiting for - the response head, so on the regular transport it would cut off uploads - and slow unary responses longer than the idle bound (those stay bounded - by their whole-request deadlines instead). + """The shared httpx transport for the control-plane REST API and the envd + HTTP API (file transfers, health checks) — one pool serves both, keyed by + the connection's proxy. + + ``for_streaming`` selects the pool carrying ``READ_TIMEOUT``, the idle + bound on every read: it resets after each successful read, so it caps how + long a streamed download may stall without limiting total transfer time. + It is fixed per pool — the adapter's per-request timeouts are + whole-request deadlines rather than idle bounds — so only streamed + downloads take it, and they get their own pool + (see :func:`get_pyqwest_transport`). """ - proxy = proxy_to_config(config.proxy) - key = (proxy, for_streaming) - with _transport_lock: - transport = _envd_transports.get(key) - if transport is None: - transport = AsyncPyqwestTransport( - retrying_http_transport( - proxy, - read_timeout=READ_TIMEOUT if for_streaming else None, - ) - ) - _envd_transports[key] = transport - return transport + return get_httpx_transport( + proxy_to_config(config.proxy), READ_TIMEOUT if for_streaming else None + ) def get_envd_api( config: ConnectionConfig, base_url: str, *, for_streaming: bool = False ) -> httpx.AsyncClient: """An httpx client for a sandbox's envd HTTP API (file transfers, health - checks) on the shared pyqwest transports. The client itself is a cheap - stateless wrapper — one per consumer is fine — while the pooled transport - underneath is shared and loop-independent.""" + checks) on the shared transports. The client itself is a cheap stateless + wrapper — one per consumer is fine — while the pooled transport underneath + is shared and loop-independent.""" return httpx.AsyncClient( base_url=base_url, - transport=get_envd_transport(config, for_streaming=for_streaming), + transport=get_transport(config, for_streaming=for_streaming), headers=config.sandbox_headers, event_hooks=make_async_logging_event_hooks(config.logger), ) diff --git a/packages/python-sdk/e2b/api/client_sync/__init__.py b/packages/python-sdk/e2b/api/client_sync/__init__.py index d7331d8cc4..eb30c4cd26 100644 --- a/packages/python-sdk/e2b/api/client_sync/__init__.py +++ b/packages/python-sdk/e2b/api/client_sync/__init__.py @@ -24,14 +24,15 @@ def get_api_client(config: ConnectionConfig, **kwargs) -> ApiClient: class ConnectionRetryTransport(SyncRetryTransport): - """Retry only failures establishing the connection — shared by the REST - API and envd RPC stacks: pyqwest raises the builtin ``ConnectionError`` - only before the request was written, so these retries can never replay a - request the server may have received (a delivered REST call or unary RPC - like ``SendInput``). This matches the connect-only ``retries`` of the - httpx transports this replaced; the retry middleware's default policy - would otherwise also retry I/O errors and 429/5xx responses for - idempotent methods.""" + """Retry only failures establishing the connection — part of the shared + transport stack, so it covers the REST API, the envd HTTP API, the envd RPC + clients and the volume content API alike: pyqwest raises the builtin + ``ConnectionError`` only before the request was written, so these retries + can never replay a request the server may have received (a delivered REST + call or unary RPC like ``SendInput``). This matches the connect-only + ``retries`` of the httpx transports this replaced; the retry middleware's + default policy would otherwise also retry I/O errors and 429/5xx responses + for idempotent methods.""" def should_retry_response( self, request: SyncRequest, response: Union[SyncResponse, Exception] @@ -39,104 +40,121 @@ def should_retry_response( return isinstance(response, ConnectionError) -def retrying_http_transport( +TransportKey = Tuple[Optional[ProxyConfig], Optional[float]] +"""Cache key of the shared transports: the proxy and the idle read bound.""" + +_transport_lock = threading.Lock() +# One pyqwest transport — one reqwest connection pool — per proxy and idle read +# bound; `None` is the direct pool. Every HTTP stack in the SDK draws from +# these: the control-plane REST API, the envd HTTP API, the envd RPC clients +# (`e2b.envd.client_sync`) and the volume content API (`e2b.volume.client_sync`). +# reqwest pools per host internally, so a single pool serves the API host and +# every per-sandbox host without interference — and because envd RPC and the +# envd HTTP API share it, a sandbox needs one HTTP/2 connection instead of one +# per stack. +# +# pyqwest transports are thread-safe, so unlike the httpx transports they +# replaced, the caches are process-global rather than per-thread. +_transports: Dict[TransportKey, ConnectionRetryTransport] = {} +# The httpx adapter over each pool, shared by every httpx client on it. +_httpx_transports: Dict[TransportKey, PyqwestTransport] = {} + + +def get_pyqwest_transport( proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None ) -> ConnectionRetryTransport: - """A fresh pyqwest transport (= its own connection pool) with the SDK's - shared tuning — system CA certs (without which TLS through an - intercepting proxy fails), the httpx-equivalent pool limits, and - connect-only retries. The REST API, envd RPC, and envd HTTP API stacks - each cache their own instances (pool unification is a follow-up). - - ``read_timeout`` bounds every read on the transport's connections; see - :func:`get_envd_transport` for when that is (and isn't) appropriate. + """The shared pyqwest transport (= one connection pool) with the SDK's + tuning — system CA certs (without which TLS through an intercepting proxy + fails) and the httpx-equivalent pool limits — behind connect-only retries. + For TLS connections ALPN negotiates the HTTP version (HTTP/2 against the + E2B API and envd), like the http2-enabled httpx transports this replaced. + + Consumers speaking pyqwest natively (the envd RPC clients) take this; + consumers speaking httpx take :func:`get_httpx_transport`, the adapter over + the very same pool. Layer concerns above it rather than into it — the RPC + stack's plain-HTTP-error normalization wraps it, headers and codecs are + per-request — so that the pool stays shareable. + + ``read_timeout`` bounds every read on the pool's connections, which is why + it is part of the cache key: reqwest's read timer keeps running while a + request body is sent and while waiting for the response head, so a pool + carrying one would cut off long uploads and slow responses. Only streamed + downloads ask for it, as an idle bound (see :func:`get_transport`). Requests are logged by pyqwest itself on the ``pyqwest.access`` and ``pyqwest`` loggers at ``DEBUG`` (off unless enabled) — the transport-level diagnostics httpcore used to provide. The SDK's own ``logger`` option is separate and sits above this, on the httpx client.""" - return ConnectionRetryTransport( - SyncHTTPTransport( - tls_include_system_certs=True, - proxy=proxy.to_pyqwest() if proxy is not None else None, - pool_idle_timeout=pool_idle_timeout, - pool_max_idle_per_host=pool_max_idle_per_host, - read_timeout=read_timeout, - # Redirects belong to the httpx client above (which the generated - # clients leave off), not to reqwest. - follow_redirects=False, - ), - max_retries=connection_retries, - ) - - -_transport_lock = threading.Lock() -# One transport (= one connection pool) per proxy; None is the direct pool. -# pyqwest transports are thread-safe, so unlike the httpx transports they -# replaced, the caches are process-global rather than per-thread. -_transports: Dict[Optional[ProxyConfig], PyqwestTransport] = {} - - -def get_transport(config: ConnectionConfig) -> PyqwestTransport: - """The shared pyqwest-backed httpx transport for REST API calls. For TLS - connections ALPN negotiates the HTTP version (HTTP/2 against the E2B - API), like the http2-enabled httpx transport this replaced.""" - proxy = proxy_to_config(config.proxy) + key = (proxy, read_timeout) with _transport_lock: - transport = _transports.get(proxy) + transport = _transports.get(key) if transport is None: - transport = PyqwestTransport(retrying_http_transport(proxy)) - _transports[proxy] = transport + transport = ConnectionRetryTransport( + SyncHTTPTransport( + tls_include_system_certs=True, + proxy=proxy.to_pyqwest() if proxy is not None else None, + pool_idle_timeout=pool_idle_timeout, + pool_max_idle_per_host=pool_max_idle_per_host, + read_timeout=read_timeout, + # Redirects belong to the httpx client above (which the + # generated clients leave off), not to reqwest. + follow_redirects=False, + ), + max_retries=connection_retries, + ) + _transports[key] = transport return transport -# One transport per (proxy, streaming) pair, separate from the REST API -# pools — envd traffic goes to per-sandbox hosts. -_envd_transports: Dict[Tuple[Optional[ProxyConfig], bool], PyqwestTransport] = {} +def get_httpx_transport( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> PyqwestTransport: + """The httpx adapter over the shared pool of + :func:`get_pyqwest_transport`, for the generated httpx clients (control + plane, envd HTTP API, volume content). The adapter holds no state of its + own and does not close the pool, so closing an httpx client leaves the + pool intact for the other clients on it.""" + key = (proxy, read_timeout) + # Resolve the pool before taking the lock: it takes the same one. + pool = get_pyqwest_transport(proxy, read_timeout) + with _transport_lock: + transport = _httpx_transports.get(key) + if transport is None: + transport = PyqwestTransport(pool) + _httpx_transports[key] = transport + return transport -def get_envd_transport( +def get_transport( config: ConnectionConfig, *, for_streaming: bool = False ) -> PyqwestTransport: - """The shared pyqwest-backed httpx transports for the envd HTTP API - (file transfers, health checks). - - The streaming transport carries ``read_timeout``, the idle bound on - every read: it resets after each successful read, so it caps how long a - streamed download may stall without limiting total transfer time. It is - fixed per transport — the adapter's per-request timeouts are - whole-request deadlines rather than idle bounds. Only streamed downloads - use it: reqwest's read timer keeps - running while a request body is sent and while waiting for the response - head, so on the regular transport it would cut off uploads and slow - unary responses longer than the idle bound (those stay bounded by their - whole-request deadlines instead). + """The shared httpx transport for the control-plane REST API and the envd + HTTP API (file transfers, health checks) — one pool serves both, keyed by + the connection's proxy. + + ``for_streaming`` selects the pool carrying ``READ_TIMEOUT``, the idle + bound on every read: it resets after each successful read, so it caps how + long a streamed download may stall without limiting total transfer time. + It is fixed per pool — the adapter's per-request timeouts are + whole-request deadlines rather than idle bounds — so only streamed + downloads take it, and they get their own pool + (see :func:`get_pyqwest_transport`). """ - proxy = proxy_to_config(config.proxy) - key = (proxy, for_streaming) - with _transport_lock: - transport = _envd_transports.get(key) - if transport is None: - transport = PyqwestTransport( - retrying_http_transport( - proxy, - read_timeout=READ_TIMEOUT if for_streaming else None, - ) - ) - _envd_transports[key] = transport - return transport + return get_httpx_transport( + proxy_to_config(config.proxy), READ_TIMEOUT if for_streaming else None + ) def get_envd_api( config: ConnectionConfig, base_url: str, *, for_streaming: bool = False ) -> httpx.Client: """An httpx client for a sandbox's envd HTTP API (file transfers, health - checks) on the shared pyqwest transports. The client itself is a cheap - stateless wrapper — one per consumer is fine — while the pooled transport - underneath is shared and thread-safe.""" + checks) on the shared transports. The client itself is a cheap stateless + wrapper — one per consumer is fine — while the pooled transport underneath + is shared and thread-safe.""" return httpx.Client( base_url=base_url, - transport=get_envd_transport(config, for_streaming=for_streaming), + transport=get_transport(config, for_streaming=for_streaming), headers=config.sandbox_headers, event_hooks=make_logging_event_hooks(config.logger), ) diff --git a/packages/python-sdk/e2b/envd/client_async/__init__.py b/packages/python-sdk/e2b/envd/client_async/__init__.py index 36a19bc1ac..dcc979d730 100644 --- a/packages/python-sdk/e2b/envd/client_async/__init__.py +++ b/packages/python-sdk/e2b/envd/client_async/__init__.py @@ -1,7 +1,6 @@ -"""Async envd RPC clients: shared pyqwest transports and client factory.""" +"""Async envd RPC clients: the plain-error transport layer and client factory.""" import asyncio -import threading from typing import ( Any, AsyncGenerator, @@ -16,8 +15,8 @@ from connectrpc.errors import ConnectError from pyqwest import Client, Request, Response, Transport -from e2b.api import ProxyConfig, proxy_to_config -from e2b.api.client_async import retrying_http_transport +from e2b.api import proxy_to_config +from e2b.api.client_async import get_pyqwest_transport from e2b.connection_config import ConnectionConfig from e2b.envd.client_shared import ( ENVD_JSON_CODEC, @@ -30,10 +29,6 @@ RES = TypeVar("RES") TClient = TypeVar("TClient") -_transport_lock = threading.Lock() -# One transport (= one connection pool) per proxy; None is the direct pool. -_transports: dict[Optional[ProxyConfig], "PlainHTTPErrorTransport"] = {} - class PlainHTTPErrorTransport: """Raise plain (non-Connect-encoded) HTTP error responses — e.g. an edge @@ -65,19 +60,6 @@ async def execute(self, request: Request) -> Response: raise error -def get_transport(proxy: Optional[ProxyConfig]) -> "PlainHTTPErrorTransport": - with _transport_lock: - transport = _transports.get(proxy) - if transport is None: - # connectrpc arms the per-call deadline around the transport, so - # retry backoff counts against the request timeout. The plain- - # error normalization sits outside the retries so it converts - # the settled response once. - transport = PlainHTTPErrorTransport(retrying_http_transport(proxy)) - _transports[proxy] = transport - return transport - - def create_rpc_client( client_cls: Callable[..., TClient], base_url: str, @@ -88,8 +70,18 @@ def create_rpc_client( see :class:`e2b.api.client_async.ConnectionRetryTransport`), the envd JSON codec, and the SDK's default-header and logging interceptors. Compression is disabled (see ``ENVD_RPC_COMPRESSION``). + + The plain-error normalization is the one RPC-only transport concern, so it + wraps the shared pool per client instead of being cached with it — a + stateless wrapper over the pool the envd HTTP API uses for the same + sandbox, which is what lets both share a single HTTP/2 connection. + connectrpc arms the per-call deadline around the transport, so retry + backoff counts against the request timeout, and the normalization sits + outside the retries so it converts the settled response once. """ - http_client = Client(get_transport(proxy_to_config(config.proxy))) + http_client = Client( + PlainHTTPErrorTransport(get_pyqwest_transport(proxy_to_config(config.proxy))) + ) return client_cls( base_url, codec=ENVD_JSON_CODEC, diff --git a/packages/python-sdk/e2b/envd/client_shared.py b/packages/python-sdk/e2b/envd/client_shared.py index aa152cd1d6..0b7f1f9dfa 100644 --- a/packages/python-sdk/e2b/envd/client_shared.py +++ b/packages/python-sdk/e2b/envd/client_shared.py @@ -1,14 +1,15 @@ """envd RPC client plumbing shared by the sync and async flavors. The envd RPC clients (process, filesystem) run on `connectrpc`, whose HTTP -layer is `pyqwest` (Rust reqwest/hyper) — built on the same -`retrying_http_transport` pieces as the REST API client in `e2b.api`, in a -separately cached pool. Only the multipart file transfer endpoints stay on -the `httpx` envd transports. Unlike the previous httpcore-based transport, -hyper sends RST_STREAM when a server stream is closed early, so abandoned -command/watch streams don't leak on the shared HTTP/2 connection. - -The flavor-specific transports and client factories live in +layer is `pyqwest` (Rust reqwest/hyper) — on the very connection pool the REST +clients in `e2b.api` use (`get_pyqwest_transport`), so RPCs and the envd HTTP +API share one HTTP/2 connection per sandbox. Only the multipart file transfer +endpoints stay on the `httpx` side of that pool. Unlike the previous +httpcore-based transport, hyper sends RST_STREAM when a server stream is +closed early, so abandoned command/watch streams don't leak on the shared +HTTP/2 connection. + +The flavor-specific transport layer and client factories live in :mod:`e2b.envd.client_sync` and :mod:`e2b.envd.client_async`, mirroring the `e2b.api.client_sync`/`client_async` layout. This module exists because `e2b/envd/__init__.py`, the natural home by that analogy, is owned by the diff --git a/packages/python-sdk/e2b/envd/client_sync/__init__.py b/packages/python-sdk/e2b/envd/client_sync/__init__.py index fc10cf9881..fadf66323f 100644 --- a/packages/python-sdk/e2b/envd/client_sync/__init__.py +++ b/packages/python-sdk/e2b/envd/client_sync/__init__.py @@ -1,7 +1,6 @@ -"""Sync envd RPC clients: shared pyqwest transports and client factory.""" +"""Sync envd RPC clients: the plain-error transport layer and client factory.""" -import threading -from typing import Any, Callable, Generator, Iterator, Optional, TypeVar, cast +from typing import Any, Callable, Generator, Iterator, TypeVar, cast from pyqwest import ( SyncClient, @@ -10,8 +9,8 @@ SyncTransport, ) -from e2b.api import ProxyConfig, proxy_to_config -from e2b.api.client_sync import retrying_http_transport +from e2b.api import proxy_to_config +from e2b.api.client_sync import get_pyqwest_transport from e2b.connection_config import ConnectionConfig from e2b.envd.client_shared import ( ENVD_JSON_CODEC, @@ -23,10 +22,6 @@ RES = TypeVar("RES") TClient = TypeVar("TClient") -_transport_lock = threading.Lock() -# One transport (= one connection pool) per proxy; None is the direct pool. -_transports: dict[Optional[ProxyConfig], "PlainHTTPErrorTransport"] = {} - class PlainHTTPErrorTransport: """Raise plain (non-Connect-encoded) HTTP error responses — e.g. an edge @@ -58,19 +53,6 @@ def execute_sync(self, request: SyncRequest) -> SyncResponse: raise error -def get_transport(proxy: Optional[ProxyConfig]) -> "PlainHTTPErrorTransport": - with _transport_lock: - transport = _transports.get(proxy) - if transport is None: - # connectrpc arms the per-call deadline around the transport, so - # retry backoff counts against the request timeout. The plain- - # error normalization sits outside the retries so it converts - # the settled response once. - transport = PlainHTTPErrorTransport(retrying_http_transport(proxy)) - _transports[proxy] = transport - return transport - - def create_rpc_client( client_cls: Callable[..., TClient], base_url: str, @@ -81,10 +63,20 @@ def create_rpc_client( see :class:`e2b.api.client_sync.ConnectionRetryTransport`), the envd JSON codec, and the SDK's default-header and logging interceptors. Compression is disabled (see ``ENVD_RPC_COMPRESSION``). The client is stateless per - call and its transport is process-global, so one instance serves all + call and its connection pool is process-global, so one instance serves all threads. + + The plain-error normalization is the one RPC-only transport concern, so it + wraps the shared pool per client instead of being cached with it — a + stateless wrapper over the pool the envd HTTP API uses for the same + sandbox, which is what lets both share a single HTTP/2 connection. + connectrpc arms the per-call deadline around the transport, so retry + backoff counts against the request timeout, and the normalization sits + outside the retries so it converts the settled response once. """ - http_client = SyncClient(get_transport(proxy_to_config(config.proxy))) + http_client = SyncClient( + PlainHTTPErrorTransport(get_pyqwest_transport(proxy_to_config(config.proxy))) + ) return client_cls( base_url, codec=ENVD_JSON_CODEC, diff --git a/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py b/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py index c36a21c86f..725662fad0 100644 --- a/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py +++ b/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py @@ -101,7 +101,7 @@ def __init__( ) self._envd_api = envd_api # Streamed downloads default to a sibling client whose transport - # carries the idle read timeout (see `get_envd_transport`). + # carries the idle read timeout (see `get_transport`). self._envd_api_streaming = get_envd_api( connection_config, envd_api_url, for_streaming=True ) @@ -218,7 +218,7 @@ async def read( # sent only when the caller set `request_timeout` explicitly # (making it the total-transfer deadline). By default a stalled # stream is bounded by the streaming transport's idle read - # timeout (see `get_envd_transport`); an explicit + # timeout (see `get_transport`); an explicit # `stream_idle_timeout` is applied per read with `wait_for` on # the regular transport instead — so values above the transport # bound aren't capped by it and `0` disables idle bounding diff --git a/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py b/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py index f249c98957..d3a2ba6463 100644 --- a/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py +++ b/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py @@ -95,7 +95,7 @@ def __init__( ) self._envd_api = envd_api # Streamed downloads default to a sibling client whose transport - # carries the idle read timeout (see `get_envd_transport`). Like the + # carries the idle read timeout (see `get_transport`). Like the # RPC client, the pyqwest transports underneath are thread-safe, so # one client serves all threads. self._envd_api_streaming = get_envd_api( @@ -211,7 +211,7 @@ def read( # sent only when the caller set `request_timeout` explicitly # (making it the total-transfer deadline). A stalled stream is # instead bounded by the streaming transport's idle read timeout - # (see `get_envd_transport`), which resets on every chunk. + # (see `get_transport`), which resets on every chunk. stream_timeout = ConnectionConfig._get_request_timeout( None, request_timeout ) diff --git a/packages/python-sdk/e2b/volume/client_async/__init__.py b/packages/python-sdk/e2b/volume/client_async/__init__.py index 0416c32a4b..3aed1272a5 100644 --- a/packages/python-sdk/e2b/volume/client_async/__init__.py +++ b/packages/python-sdk/e2b/volume/client_async/__init__.py @@ -1,19 +1,11 @@ -import threading -from typing import Dict, Optional, Tuple - import httpx -from pyqwest import HTTPTransport from pyqwest.httpx import AsyncPyqwestTransport from e2b.api import ( - ProxyConfig, - connection_retries, make_async_logging_event_hooks, - pool_idle_timeout, - pool_max_idle_per_host, proxy_to_config, ) -from e2b.api.client_async import ConnectionRetryTransport +from e2b.api.client_async import get_httpx_transport from e2b.api.metadata import default_headers from e2b.exceptions import AuthenticationException from e2b.volume.client.client import AuthenticatedClient as AsyncVolumeApiClient @@ -69,18 +61,11 @@ def _api_client( ) -_transport_lock = threading.Lock() -# One transport (= one connection pool) per proxy and read timeout; None is -# the direct pool. pyqwest's I/O runs on its own Rust runtime, so unlike the -# httpx transport this replaced, the transport is not bound to an event loop -# and the cache is process-global rather than per-loop. -_transports: Dict[ - Tuple[Optional[ProxyConfig], Optional[float]], AsyncPyqwestTransport -] = {} - - def get_transport(config: VolumeConnectionConfig) -> AsyncPyqwestTransport: - """The shared pyqwest-backed httpx transport for volume content API calls. + """The shared pyqwest-backed httpx transport for volume content API calls — + the same pool the control-plane REST API and the envd HTTP API draw from + (see :func:`e2b.api.client_async.get_pyqwest_transport`); reqwest pools per + host, so the volume host gets its own connections within it. It carries no idle read bound: reqwest's read timer keeps running while a request body is sent and while waiting for the response head, so one here @@ -88,7 +73,7 @@ def get_transport(config: VolumeConnectionConfig) -> AsyncPyqwestTransport: their whole-request deadlines instead). Streamed downloads, which do need an idle bound, use :func:`get_streaming_transport`. """ - return _transport(config, read_timeout=None) + return get_httpx_transport(proxy_to_config(config.proxy)) def get_streaming_transport( @@ -97,34 +82,9 @@ def get_streaming_transport( """The transport for streamed downloads, carrying ``READ_TIMEOUT`` as the idle bound on every read: it resets after each successful read, so it caps how long a streamed download may stall without limiting total transfer - time. It is fixed per transport — the adapter's per-request timeouts are - whole-request deadlines rather than idle bounds. + time. It is fixed per pool — the adapter's per-request timeouts are + whole-request deadlines rather than idle bounds — so streamed downloads get + their own, shared with the sandbox filesystem's streaming transport + whenever the two bounds agree. """ - return _transport(config, read_timeout=READ_TIMEOUT) - - -def _transport( - config: VolumeConnectionConfig, *, read_timeout: Optional[float] -) -> AsyncPyqwestTransport: - proxy = proxy_to_config(config.proxy) - key = (proxy, read_timeout) - with _transport_lock: - transport = _transports.get(key) - if transport is None: - transport = AsyncPyqwestTransport( - ConnectionRetryTransport( - HTTPTransport( - tls_include_system_certs=True, - proxy=proxy.to_pyqwest() if proxy is not None else None, - pool_idle_timeout=pool_idle_timeout, - pool_max_idle_per_host=pool_max_idle_per_host, - read_timeout=read_timeout, - # Redirects belong to the httpx client above (which the - # generated clients leave off), not to reqwest. - follow_redirects=False, - ), - max_retries=connection_retries, - ) - ) - _transports[key] = transport - return transport + return get_httpx_transport(proxy_to_config(config.proxy), READ_TIMEOUT) diff --git a/packages/python-sdk/e2b/volume/client_sync/__init__.py b/packages/python-sdk/e2b/volume/client_sync/__init__.py index adbe9c3e72..57344abbd2 100644 --- a/packages/python-sdk/e2b/volume/client_sync/__init__.py +++ b/packages/python-sdk/e2b/volume/client_sync/__init__.py @@ -1,19 +1,11 @@ -import threading -from typing import Dict, Optional, Tuple - import httpx -from pyqwest import SyncHTTPTransport from pyqwest.httpx import PyqwestTransport from e2b.api import ( - ProxyConfig, - connection_retries, make_logging_event_hooks, - pool_idle_timeout, - pool_max_idle_per_host, proxy_to_config, ) -from e2b.api.client_sync import ConnectionRetryTransport +from e2b.api.client_sync import get_httpx_transport from e2b.api.metadata import default_headers from e2b.exceptions import AuthenticationException from e2b.volume.client.client import AuthenticatedClient as VolumeApiClient @@ -69,15 +61,11 @@ def _api_client( ) -_transport_lock = threading.Lock() -# One transport (= one connection pool) per proxy and read timeout; None is -# the direct pool. pyqwest transports are thread-safe, so unlike the httpx -# transport this replaced, the cache is process-global rather than per-thread. -_transports: Dict[Tuple[Optional[ProxyConfig], Optional[float]], PyqwestTransport] = {} - - def get_transport(config: VolumeConnectionConfig) -> PyqwestTransport: - """The shared pyqwest-backed httpx transport for volume content API calls. + """The shared pyqwest-backed httpx transport for volume content API calls — + the same pool the control-plane REST API and the envd HTTP API draw from + (see :func:`e2b.api.client_sync.get_pyqwest_transport`); reqwest pools per + host, so the volume host gets its own connections within it. It carries no idle read bound: reqwest's read timer keeps running while a request body is sent and while waiting for the response head, so one here @@ -85,41 +73,16 @@ def get_transport(config: VolumeConnectionConfig) -> PyqwestTransport: their whole-request deadlines instead). Streamed downloads, which do need an idle bound, use :func:`get_streaming_transport`. """ - return _transport(config, read_timeout=None) + return get_httpx_transport(proxy_to_config(config.proxy)) def get_streaming_transport(config: VolumeConnectionConfig) -> PyqwestTransport: """The transport for streamed downloads, carrying ``READ_TIMEOUT`` as the idle bound on every read: it resets after each successful read, so it caps how long a streamed download may stall without limiting total transfer - time. It is fixed per transport — the adapter's per-request timeouts are - whole-request deadlines rather than idle bounds. + time. It is fixed per pool — the adapter's per-request timeouts are + whole-request deadlines rather than idle bounds — so streamed downloads get + their own, shared with the sandbox filesystem's streaming transport + whenever the two bounds agree. """ - return _transport(config, read_timeout=READ_TIMEOUT) - - -def _transport( - config: VolumeConnectionConfig, *, read_timeout: Optional[float] -) -> PyqwestTransport: - proxy = proxy_to_config(config.proxy) - key = (proxy, read_timeout) - with _transport_lock: - transport = _transports.get(key) - if transport is None: - transport = PyqwestTransport( - ConnectionRetryTransport( - SyncHTTPTransport( - tls_include_system_certs=True, - proxy=proxy.to_pyqwest() if proxy is not None else None, - pool_idle_timeout=pool_idle_timeout, - pool_max_idle_per_host=pool_max_idle_per_host, - read_timeout=read_timeout, - # Redirects belong to the httpx client above (which the - # generated clients leave off), not to reqwest. - follow_redirects=False, - ), - max_retries=connection_retries, - ) - ) - _transports[key] = transport - return transport + return get_httpx_transport(proxy_to_config(config.proxy), READ_TIMEOUT) diff --git a/packages/python-sdk/tests/envd_frame_server.py b/packages/python-sdk/tests/envd_frame_server.py index df7f33447b..dfa3b6c514 100644 --- a/packages/python-sdk/tests/envd_frame_server.py +++ b/packages/python-sdk/tests/envd_frame_server.py @@ -16,10 +16,11 @@ import struct import threading import time -from typing import Iterator, Optional +from typing import Iterator, List, Optional import h2.config import h2.connection +import h2.errors import h2.events from protobuf import Oneof from pyqwest import ( @@ -33,6 +34,7 @@ ) from e2b.connection_config import ConnectionConfig +from e2b.envd.api import ENVD_API_HEALTH_ROUTE from e2b.envd.client_shared import ENVD_JSON_CODEC, ENVD_RPC_COMPRESSION from e2b.envd.interceptors import build_interceptors from e2b.envd.process.process_connect import ProcessClient, ProcessClientSync @@ -177,6 +179,138 @@ def frame_recording_server( server.listener.close() +class SharedPoolServer(threading.Thread): + """Multi-connection plaintext HTTP/2 server for the shared-pool tests. + + Serves both sides of a sandbox's traffic — the process ``connect`` server + stream and the envd HTTP ``/health`` route — so a single connection pool + can be pointed at it the way the SDK points one at a real sandbox. After + the first stream event it breaks the RPC the way a sandbox going away + does: + + * ``fault="reset"`` sends ``RST_STREAM``, which kills the RPC stream and + leaves the HTTP/2 connection healthy; + * ``fault="drop"`` tears the whole TCP connection down with a RST. + + ``/health`` is always answered 200, so a probe that comes back anything + but ``True`` failed at the transport layer. ``connections`` counts the + accepted TCP connections, which is what tells reuse from a redial. + """ + + def __init__(self, fault: str): + super().__init__(daemon=True) + self.fault = fault + self.listener = socket.create_server(("127.0.0.1", 0)) + self.listener.settimeout(10) + self.port = self.listener.getsockname()[1] + self.connections: List[socket.socket] = [] + self.paths: List[str] = [] + self.errors: List[str] = [] + self._lock = threading.Lock() + + def run(self): + while True: + try: + sock, _ = self.listener.accept() + except (OSError, socket.timeout): + # The listener was closed by the context manager, or nothing + # else connected — either way there is nothing left to serve. + return + with self._lock: + self.connections.append(sock) + threading.Thread(target=self._serve, args=(sock,), daemon=True).start() + + def _serve(self, sock: socket.socket): + sock.settimeout(0.1) + conn = h2.connection.H2Connection( + config=h2.config.H2Configuration(client_side=False) + ) + conn.initiate_connection() + sock.sendall(conn.data_to_send()) + paths: dict[int, str] = {} + drop_connection = False + deadline = time.monotonic() + 10 + try: + while time.monotonic() < deadline: + try: + data = sock.recv(65535) + except socket.timeout: + continue + except OSError: + break + if not data: + break + for event in conn.receive_data(data): + if isinstance(event, h2.events.RequestReceived): + path = dict(event.headers).get(b":path", b"").decode() + paths[event.stream_id] = path + with self._lock: + self.paths.append(path) + elif isinstance(event, h2.events.StreamEnded): + # The client finished sending the request: respond. + if paths.get(event.stream_id, "") == ENVD_API_HEALTH_ROUTE: + body = b'{"version":"0.0.0"}' + conn.send_headers( + event.stream_id, + [ + (":status", "200"), + ("content-type", "application/json"), + ("content-length", str(len(body))), + ], + ) + conn.send_data(event.stream_id, body, end_stream=True) + continue + conn.send_headers( + event.stream_id, + [ + (":status", "200"), + ("content-type", "application/connect+json"), + ], + ) + conn.send_data(event.stream_id, event_envelope()) + if self.fault == "reset": + conn.reset_stream( + event.stream_id, + error_code=h2.errors.ErrorCodes.INTERNAL_ERROR, + ) + else: + drop_connection = True + elif isinstance(event, h2.events.DataReceived): + conn.acknowledge_received_data( + event.flow_controlled_length, event.stream_id + ) + elif isinstance(event, h2.events.ConnectionTerminated): + return + out = conn.data_to_send() + if out: + sock.sendall(out) + if drop_connection: + # RST the connection rather than closing it cleanly: a + # GOAWAY would tell the client to retire the connection, + # which is not what a sandbox disappearing looks like. + sock.setsockopt( + socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0) + ) + return + except Exception as e: # noqa: BLE001 — surfaced via assert_no_errors + self.errors.append(repr(e)) + finally: + sock.close() + + def assert_no_errors(self): + assert not self.errors, self.errors + + +@contextlib.contextmanager +def shared_pool_server(fault: str) -> Iterator[SharedPoolServer]: + server = SharedPoolServer(fault) + server.start() + try: + yield server + finally: + server.listener.close() + + def assert_stdout_event(event: ConnectResponse): assert event.event is not None match event.event.event: diff --git a/packages/python-sdk/tests/test_api_client_transport.py b/packages/python-sdk/tests/test_api_client_transport.py index c37d034642..bca84d2bb1 100644 --- a/packages/python-sdk/tests/test_api_client_transport.py +++ b/packages/python-sdk/tests/test_api_client_transport.py @@ -10,37 +10,24 @@ import httpx import pytest from pyqwest.httpx import AsyncPyqwestTransport, PyqwestTransport +from transport_caches import reset_transport_caches -import e2b.api.client_async as client_async -import e2b.api.client_sync as client_sync from e2b.api.client_async import get_api_client as get_async_api_client from e2b.api.client_async import get_envd_api as get_async_envd_api -from e2b.api.client_async import get_envd_transport as get_async_envd_transport from e2b.api.client_async import get_transport as get_async_transport from e2b.api.client_sync import get_api_client as get_sync_api_client from e2b.api.client_sync import get_envd_api as get_sync_envd_api -from e2b.api.client_sync import get_envd_transport as get_sync_envd_transport from e2b.api.client_sync import get_transport as get_sync_transport from e2b.connection_config import ConnectionConfig -def reset_sync_api_transports(): - client_sync._transports.clear() - client_sync._envd_transports.clear() - - -def reset_async_api_transports(): - client_async._transports.clear() - client_async._envd_transports.clear() - - def run_in_worker_thread(fn): with ThreadPoolExecutor(max_workers=1) as executor: return executor.submit(fn).result() def test_sync_api_client_proxy_uses_explicit_transport(test_api_key): - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig( api_key=test_api_key, proxy="http://127.0.0.1:9999", @@ -56,11 +43,11 @@ def test_sync_api_client_proxy_uses_explicit_transport(test_api_key): assert httpx_client._mounts == {} finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() def test_sync_get_transport_keyed_by_proxy(test_api_key): - reset_sync_api_transports() + reset_transport_caches() proxied_config = ConnectionConfig( api_key=test_api_key, proxy="http://127.0.0.1:9999", @@ -83,11 +70,11 @@ def test_sync_get_transport_keyed_by_proxy(test_api_key): assert get_sync_transport(proxied_config) is proxied_transport assert get_sync_transport(direct_config) is direct_transport finally: - reset_sync_api_transports() + reset_transport_caches() def test_sync_api_client_applies_request_timeout(test_api_key): - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, request_timeout=1.5) api_client = get_sync_api_client(config) @@ -97,11 +84,11 @@ def test_sync_api_client_applies_request_timeout(test_api_key): assert httpx_client.timeout == httpx.Timeout(1.5) finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() def test_sync_api_client_request_timeout_zero_disables_timeout(test_api_key): - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, request_timeout=0) api_client = get_sync_api_client(config) @@ -111,34 +98,30 @@ def test_sync_api_client_request_timeout_zero_disables_timeout(test_api_key): assert httpx_client.timeout == httpx.Timeout(None) finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() -def test_sync_envd_transports_keyed_by_streaming(test_api_key): - # The envd HTTP API pools are separate from the REST API pools, and the - # streaming variant (which carries the idle read timeout) is its own - # pool per proxy. - reset_sync_api_transports() +def test_sync_envd_and_api_share_one_transport(test_api_key): + # The envd HTTP API draws from the same pool as the control-plane REST API + # — reqwest pools per host, so one pool serves both. Only the streaming + # variant, which carries the idle read timeout, is a pool of its own. + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) try: api_transport = get_sync_transport(config) - envd_transport = get_sync_envd_transport(config) - streaming_transport = get_sync_envd_transport(config, for_streaming=True) - - assert isinstance(envd_transport, PyqwestTransport) - assert envd_transport is not api_transport - assert streaming_transport is not envd_transport - assert get_sync_envd_transport(config) is envd_transport - assert ( - get_sync_envd_transport(config, for_streaming=True) is streaming_transport - ) + streaming_transport = get_sync_transport(config, for_streaming=True) + + assert isinstance(api_transport, PyqwestTransport) + assert api_transport is get_sync_transport(config, for_streaming=False) + assert streaming_transport is not api_transport + assert get_sync_transport(config, for_streaming=True) is streaming_transport finally: - reset_sync_api_transports() + reset_transport_caches() def test_sync_envd_api_client_wiring(test_api_key): - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, access_token="tok") client = get_sync_envd_api(config, "https://sandbox.e2b.app") @@ -146,23 +129,21 @@ def test_sync_envd_api_client_wiring(test_api_key): try: assert client.base_url == "https://sandbox.e2b.app" - assert client._transport is get_sync_envd_transport(config) - assert streaming._transport is get_sync_envd_transport( - config, for_streaming=True - ) + assert client._transport is get_sync_transport(config) + assert streaming._transport is get_sync_transport(config, for_streaming=True) for header, value in config.sandbox_headers.items(): assert client.headers[header] == value finally: client.close() streaming.close() - reset_sync_api_transports() + reset_transport_caches() def test_sync_api_client_is_shared_across_threads(test_api_key): # httpx.Client is thread-safe and the pyqwest transport underneath is # too, so a single client (and its pool) serves all threads — the # per-thread client caching this replaced is gone. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) api_client = get_sync_api_client(config) @@ -174,12 +155,12 @@ def test_sync_api_client_is_shared_across_threads(test_api_key): assert worker_client is main_client finally: main_client.close() - reset_sync_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_api_client_proxy_uses_explicit_transport(test_api_key): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig( api_key=test_api_key, proxy="http://127.0.0.1:9999", @@ -195,12 +176,12 @@ async def test_async_api_client_proxy_uses_explicit_transport(test_api_key): assert httpx_client._mounts == {} finally: await httpx_client.aclose() - reset_async_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_get_transport_keyed_by_proxy(test_api_key): - reset_async_api_transports() + reset_transport_caches() proxied_config = ConnectionConfig( api_key=test_api_key, proxy="http://127.0.0.1:9999", @@ -216,7 +197,7 @@ async def test_async_get_transport_keyed_by_proxy(test_api_key): assert get_async_transport(proxied_config) is proxied_transport assert get_async_transport(direct_config) is direct_transport finally: - reset_async_api_transports() + reset_transport_caches() @pytest.mark.asyncio @@ -225,7 +206,7 @@ async def test_async_api_client_is_shared_across_loops(test_api_key): # nor the httpx client wrapper is bound to an event loop — a single # client serves all loops (the per-loop client caching this replaced is # gone). - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) api_client = get_async_api_client(config) @@ -243,45 +224,41 @@ async def get_client(): assert other_loop_client is main_client finally: await main_client.aclose() - reset_async_api_transports() + reset_transport_caches() @pytest.mark.asyncio -async def test_async_envd_transports_keyed_by_streaming(test_api_key): - reset_async_api_transports() +async def test_async_envd_and_api_share_one_transport(test_api_key): + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) try: api_transport = get_async_transport(config) - envd_transport = get_async_envd_transport(config) - streaming_transport = get_async_envd_transport(config, for_streaming=True) - - assert isinstance(envd_transport, AsyncPyqwestTransport) - assert envd_transport is not api_transport - assert streaming_transport is not envd_transport - assert get_async_envd_transport(config) is envd_transport - assert ( - get_async_envd_transport(config, for_streaming=True) is streaming_transport - ) + streaming_transport = get_async_transport(config, for_streaming=True) + + assert isinstance(api_transport, AsyncPyqwestTransport) + assert api_transport is get_async_transport(config, for_streaming=False) + assert streaming_transport is not api_transport + assert get_async_transport(config, for_streaming=True) is streaming_transport finally: - reset_async_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_envd_api_client_wiring(test_api_key): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, access_token="tok") client = get_async_envd_api(config, "https://sandbox.e2b.app") try: assert client.base_url == "https://sandbox.e2b.app" - assert client._transport is get_async_envd_transport(config) + assert client._transport is get_async_transport(config) for header, value in config.sandbox_headers.items(): assert client.headers[header] == value finally: await client.aclose() - reset_async_api_transports() + reset_transport_caches() class _EchoHandler(BaseHTTPRequestHandler): @@ -349,7 +326,7 @@ def test_sync_transport_sends_proxy_credentials_and_headers(test_api_key, echo_s # Everything an httpx.Proxy can express reaches the proxy: the echo server # stands in for one, so the request arrives in absolute form with the # credentials and the extra headers configured for it. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig( api_key=test_api_key, proxy=httpx.Proxy( @@ -367,14 +344,14 @@ def test_sync_transport_sends_proxy_credentials_and_headers(test_api_key, echo_s assert echoed["headers"]["x-proxy-token"] == "t" finally: client.close() - reset_sync_api_transports() + reset_transport_caches() def test_transport_emits_pyqwest_access_log(test_api_key, echo_server, caplog): # pyqwest logs every request on `pyqwest.access` at DEBUG — the # transport-level diagnostics httpcore used to provide, and separate from # the SDK's own `logger` option. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_sync_api_client(config) httpx_client = api_client.get_httpx_client() @@ -392,11 +369,11 @@ def test_transport_emits_pyqwest_access_log(test_api_key, echo_server, caplog): ] finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() def test_sync_api_client_round_trips_through_pyqwest(test_api_key, echo_server): - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_sync_api_client(config) httpx_client = api_client.get_httpx_client() @@ -411,13 +388,13 @@ def test_sync_api_client_round_trips_through_pyqwest(test_api_key, echo_server): assert echoed["headers"]["package_version"] finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() def test_sync_api_client_serves_concurrent_threads(test_api_key, echo_server): # The scenario the removed per-thread client caching used to guard: one # client, one shared pyqwest pool, many threads at once. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_sync_api_client(config) httpx_client = api_client.get_httpx_client() @@ -433,12 +410,12 @@ def request(i: int) -> tuple[int, str]: assert results == [(200, f"/sandboxes/{i}") for i in range(32)] finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_api_client_serves_concurrent_requests(test_api_key, echo_server): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_async_api_client(config) httpx_client = api_client.get_async_httpx_client() @@ -452,12 +429,12 @@ async def request(i: int) -> tuple[int, str]: assert list(results) == [(200, f"/sandboxes/{i}") for i in range(32)] finally: await httpx_client.aclose() - reset_async_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_api_client_round_trips_through_pyqwest(test_api_key, echo_server): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_async_api_client(config) httpx_client = api_client.get_async_httpx_client() @@ -472,14 +449,14 @@ async def test_async_api_client_round_trips_through_pyqwest(test_api_key, echo_s assert echoed["headers"]["package_version"] finally: await httpx_client.aclose() - reset_async_api_transports() + reset_transport_caches() def test_sync_api_client_leaves_redirects_to_httpx(test_api_key, echo_server): # reqwest would otherwise follow redirects inside the transport, hiding them # from httpx: the generated client asks for no redirect following, so a 302 # must surface as-is, and opting in must record the hop in `history`. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_sync_api_client(config) httpx_client = api_client.get_httpx_client() @@ -497,12 +474,12 @@ def test_sync_api_client_leaves_redirects_to_httpx(test_api_key, echo_server): assert [r.status_code for r in followed.history] == [302] finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_api_client_leaves_redirects_to_httpx(test_api_key, echo_server): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_async_api_client(config) httpx_client = api_client.get_async_httpx_client() @@ -520,13 +497,13 @@ async def test_async_api_client_leaves_redirects_to_httpx(test_api_key, echo_ser assert [r.status_code for r in followed.history] == [302] finally: await httpx_client.aclose() - reset_async_api_transports() + reset_transport_caches() def test_sync_api_client_timeout_raises_httpx_read_timeout(test_api_key, echo_server): # pyqwest raises the builtin TimeoutError; the transport re-raises it as # httpx.ReadTimeout to keep the httpx.TimeoutException contract. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_sync_api_client(config) httpx_client = api_client.get_httpx_client() @@ -536,14 +513,14 @@ def test_sync_api_client_timeout_raises_httpx_read_timeout(test_api_key, echo_se httpx_client.request("GET", "/slow", timeout=0.2) finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_api_client_timeout_raises_httpx_read_timeout( test_api_key, echo_server ): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_async_api_client(config) httpx_client = api_client.get_async_httpx_client() @@ -553,7 +530,7 @@ async def test_async_api_client_timeout_raises_httpx_read_timeout( await httpx_client.request("GET", "/slow", timeout=0.2) finally: await httpx_client.aclose() - reset_async_api_transports() + reset_transport_caches() def test_sync_api_client_body_timeout_raises_httpx_read_timeout( @@ -561,7 +538,7 @@ def test_sync_api_client_body_timeout_raises_httpx_read_timeout( ): # The head arrives in time and the body never does: httpx reads the body # after the transport returned, so that timeout is mapped on the stream. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_sync_api_client(config) httpx_client = api_client.get_httpx_client() @@ -571,14 +548,14 @@ def test_sync_api_client_body_timeout_raises_httpx_read_timeout( httpx_client.request("GET", "/stall", timeout=0.2) finally: httpx_client.close() - reset_sync_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_api_client_body_timeout_raises_httpx_read_timeout( test_api_key, echo_server ): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) api_client = get_async_api_client(config) httpx_client = api_client.get_async_httpx_client() @@ -588,7 +565,7 @@ async def test_async_api_client_body_timeout_raises_httpx_read_timeout( await httpx_client.request("GET", "/stall", timeout=0.2) finally: await httpx_client.aclose() - reset_async_api_transports() + reset_transport_caches() def test_sync_transport_sends_multipart_bodies(test_api_key, echo_server): @@ -597,11 +574,9 @@ def test_sync_transport_sends_multipart_bodies(test_api_key, echo_server): # sync path used to match AsyncByteStream first and raise from inside the # body iterator, surfacing as a WriteError mid-request; pyqwest 0.8 matches # the sync case first, so the SDK no longer rewraps the stream. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) - client = httpx.Client( - base_url=echo_server, transport=client_sync.get_envd_transport(config) - ) + client = httpx.Client(base_url=echo_server, transport=get_sync_transport(config)) try: response = client.post("/files", files=[("file", ("a.txt", b"x" * 4096))]) @@ -609,4 +584,4 @@ def test_sync_transport_sends_multipart_bodies(test_api_key, echo_server): assert response.json()["received"] > 4096 finally: client.close() - reset_sync_api_transports() + reset_transport_caches() diff --git a/packages/python-sdk/tests/test_envd_client_transport.py b/packages/python-sdk/tests/test_envd_client_transport.py index 683d65faea..1b8c6203e5 100644 --- a/packages/python-sdk/tests/test_envd_client_transport.py +++ b/packages/python-sdk/tests/test_envd_client_transport.py @@ -3,22 +3,22 @@ import httpx import pytest from pyqwest import Proxy +from transport_caches import reset_transport_caches import e2b.api.client_async as api_client_async import e2b.api.client_sync as api_client_sync from e2b.api import ProxyConfig, proxy_to_config -from e2b.connection_config import ProxyTypes +from e2b.connection_config import ConnectionConfig, ProxyTypes from e2b.envd import client_async, client_sync +from e2b.envd.process.process_connect import ProcessClient, ProcessClientSync from e2b.exceptions import InvalidArgumentException @pytest.fixture(autouse=True) -def reset_transport_caches(): - client_sync._transports.clear() - client_async._transports.clear() +def clear_transport_caches(): + reset_transport_caches() yield - client_sync._transports.clear() - client_async._transports.clear() + reset_transport_caches() def test_proxy_to_config_none(): @@ -97,26 +97,28 @@ def test_proxy_to_config_rejects_unknown_types(): proxy_to_config(cast(ProxyTypes, object())) -def test_sync_transport_is_cached_per_proxy(): +def test_sync_pool_is_cached_per_proxy(): proxy = ProxyConfig("http://127.0.0.1:8080") - transport_a = client_sync.get_transport(None) - transport_b = client_sync.get_transport(None) - transport_c = client_sync.get_transport(proxy) + pool_a = api_client_sync.get_pyqwest_transport(None) + pool_b = api_client_sync.get_pyqwest_transport(None) + pool_c = api_client_sync.get_pyqwest_transport(proxy) # A second, equal config keys the same pool. - transport_d = client_sync.get_transport(ProxyConfig("http://127.0.0.1:8080")) + pool_d = api_client_sync.get_pyqwest_transport(ProxyConfig("http://127.0.0.1:8080")) - assert transport_a is transport_b - assert transport_c is transport_d - assert transport_a is not transport_c + assert pool_a is pool_b + assert pool_c is pool_d + assert pool_a is not pool_c -def test_sync_transport_is_not_shared_across_proxy_credentials(): +def test_sync_pool_is_not_shared_across_proxy_credentials(): # Same proxy URL, different credentials or headers: separate pools, since # the proxy configuration is fixed per transport. url = "http://127.0.0.1:8080" - plain = client_sync.get_transport(ProxyConfig(url)) - with_auth = client_sync.get_transport(ProxyConfig(url, auth=("user", "pass"))) - with_headers = client_sync.get_transport( + plain = api_client_sync.get_pyqwest_transport(ProxyConfig(url)) + with_auth = api_client_sync.get_pyqwest_transport( + ProxyConfig(url, auth=("user", "pass")) + ) + with_headers = api_client_sync.get_pyqwest_transport( ProxyConfig(url, headers=(("x-custom", "1"),)) ) @@ -125,27 +127,56 @@ def test_sync_transport_is_not_shared_across_proxy_credentials(): assert with_auth is not with_headers -def test_async_transport_is_cached_per_proxy(): - transport_a = client_async.get_transport(None) - transport_b = client_async.get_transport(None) - transport_c = client_async.get_transport(ProxyConfig("http://127.0.0.1:8080")) - - assert transport_a is transport_b - assert transport_a is not transport_c - assert client_sync.get_transport(None) is not transport_a - +def test_async_pool_is_cached_per_proxy(): + pool_a = api_client_async.get_pyqwest_transport(None) + pool_b = api_client_async.get_pyqwest_transport(None) + pool_c = api_client_async.get_pyqwest_transport( + ProxyConfig("http://127.0.0.1:8080") + ) -def test_transport_stack_normalizes_plain_errors_and_retries_connects(): - # The shared transports are the plain-HTTP-error normalization wrapping - # the connection retries; `E2B_CONNECTION_RETRIES` must flow into the - # retry layer the way it does into the httpx REST transports. + assert pool_a is pool_b + assert pool_a is not pool_c + # Sync and async are separate stacks all the way down. + assert api_client_sync.get_pyqwest_transport(None) is not pool_a + + +def test_rpc_clients_run_on_the_shared_pool(test_api_key, monkeypatch): + # The RPC stack is the plain-HTTP-error normalization wrapping the very + # pool the httpx clients use, so an envd RPC and an envd HTTP call to the + # same sandbox share one HTTP/2 connection. `pyqwest.SyncClient` doesn't + # hand its transport back, so record what the normalization is given. + config = ConnectionConfig(api_key=test_api_key) + pool = api_client_sync.get_pyqwest_transport(None) + async_pool = api_client_async.get_pyqwest_transport(None) + # The httpx adapters every REST client uses sit on those same pools. + assert api_client_sync.get_httpx_transport(None)._transport is pool + assert api_client_async.get_httpx_transport(None)._transport is async_pool + + wrapped = [] + for module in (client_sync, client_async): + normalization = module.PlainHTTPErrorTransport + monkeypatch.setattr( + module, + "PlainHTTPErrorTransport", + lambda inner, normalization=normalization: ( + wrapped.append(inner) or normalization(inner) + ), + ) + + client_sync.create_rpc_client(ProcessClientSync, "https://sandbox.e2b.app", config) + client_async.create_rpc_client(ProcessClient, "https://sandbox.e2b.app", config) + + assert wrapped == [pool, async_pool] + + +def test_shared_pool_retries_connects(): + # `E2B_CONNECTION_RETRIES` must flow into the retry layer of the shared + # pool, which every stack now inherits. from e2b.api import connection_retries - sync_transport = client_sync.get_transport(None) - async_transport = client_async.get_transport(None) - assert isinstance(sync_transport, client_sync.PlainHTTPErrorTransport) - assert isinstance(async_transport, client_async.PlainHTTPErrorTransport) - assert isinstance(sync_transport._inner, api_client_sync.ConnectionRetryTransport) - assert isinstance(async_transport._inner, api_client_async.ConnectionRetryTransport) - assert sync_transport._inner._max_retries == connection_retries - assert async_transport._inner._max_retries == connection_retries + pool = api_client_sync.get_pyqwest_transport(None) + apool = api_client_async.get_pyqwest_transport(None) + assert isinstance(pool, api_client_sync.ConnectionRetryTransport) + assert isinstance(apool, api_client_async.ConnectionRetryTransport) + assert pool._max_retries == connection_retries + assert apool._max_retries == connection_retries diff --git a/packages/python-sdk/tests/test_shared_transport_pool.py b/packages/python-sdk/tests/test_shared_transport_pool.py new file mode 100644 index 0000000000..d5a708a4cd --- /dev/null +++ b/packages/python-sdk/tests/test_shared_transport_pool.py @@ -0,0 +1,158 @@ +"""The envd RPC stack and the envd HTTP API share one connection pool, so a +sandbox costs one HTTP/2 connection instead of one per stack (SDK-291). + +That sharing puts the sandbox health probe on the connection the failed RPC +was using: ``handle_rpc_exception_with_health`` calls ``/health`` precisely +when an RPC died at the transport layer, and it must still get an answer, or +every dropped connection would be reported as an indeterminate state instead +of "the sandbox is gone". These tests pin that at the frame level rather than +trusting reqwest to discard broken connections — a real plaintext HTTP/2 +server serves both routes on one pool and counts the TCP connections the +client opens (see ``envd_frame_server``): + +* an ``RST_STREAM`` kills the RPC stream only, so the probe reuses the same + connection (which is what proves the pool is genuinely shared); +* a dropped TCP connection makes reqwest redial for the probe. + +The pool here mirrors the SDK's — the same connect-retry middleware under the +same plain-error normalization — but with HTTP/2 prior knowledge, since the +transports the SDK builds negotiate the version over TLS via ALPN and this +server is plaintext. +""" + +import httpx +import pytest +from connectrpc.errors import ConnectError +from envd_frame_server import ( + assert_stdout_event, + make_async_client, + make_sync_client, + shared_pool_server, +) +from pyqwest import HTTPTransport, HTTPVersion, SyncHTTPTransport +from pyqwest.httpx import AsyncPyqwestTransport, PyqwestTransport + +from e2b.api import connection_retries +from e2b.api.client_async import ConnectionRetryTransport +from e2b.api.client_sync import ( + ConnectionRetryTransport as SyncConnectionRetryTransport, +) +from e2b.envd.api import acheck_sandbox_health, check_sandbox_health +from e2b.envd.client_async import PlainHTTPErrorTransport +from e2b.envd.client_sync import ( + PlainHTTPErrorTransport as SyncPlainHTTPErrorTransport, +) +from e2b.envd.process.process_pb import ConnectRequest +from e2b.envd.rpc import is_transport_failure + + +def _sync_pool() -> SyncConnectionRetryTransport: + return SyncConnectionRetryTransport( + SyncHTTPTransport(http_version=HTTPVersion.HTTP2), + max_retries=connection_retries, + ) + + +def _async_pool() -> ConnectionRetryTransport: + return ConnectionRetryTransport( + HTTPTransport(http_version=HTTPVersion.HTTP2), + max_retries=connection_retries, + ) + + +def _break_sync_stream(events) -> ConnectError: + """Read the first event, then the failure the server injects after it.""" + assert_stdout_event(next(events)) + with pytest.raises(ConnectError) as excinfo: + next(events) + assert is_transport_failure(excinfo.value), excinfo.value + return excinfo.value + + +async def _break_async_stream(events) -> ConnectError: + assert_stdout_event(await events.__anext__()) + with pytest.raises(ConnectError) as excinfo: + await events.__anext__() + assert is_transport_failure(excinfo.value), excinfo.value + return excinfo.value + + +def test_sync_stream_reset_leaves_the_shared_connection_usable(): + with shared_pool_server("reset") as server: + pool = _sync_pool() + envd_api = httpx.Client( + base_url=f"http://127.0.0.1:{server.port}", + transport=PyqwestTransport(pool), + ) + try: + _break_sync_stream( + make_sync_client( + server.port, transport=SyncPlainHTTPErrorTransport(pool) + ).connect(ConnectRequest()) + ) + assert check_sandbox_health(envd_api) is True + # One TCP connection served the RPC and the probe: the reset took + # down the stream, not the connection. + assert len(server.connections) == 1 + server.assert_no_errors() + finally: + envd_api.close() + + +def test_sync_dropped_connection_redials_for_the_health_probe(): + with shared_pool_server("drop") as server: + pool = _sync_pool() + envd_api = httpx.Client( + base_url=f"http://127.0.0.1:{server.port}", + transport=PyqwestTransport(pool), + ) + try: + _break_sync_stream( + make_sync_client( + server.port, transport=SyncPlainHTTPErrorTransport(pool) + ).connect(ConnectRequest()) + ) + # The probe must not be answered from the dead pooled connection. + assert check_sandbox_health(envd_api) is True + assert len(server.connections) == 2 + finally: + envd_api.close() + + +async def test_async_stream_reset_leaves_the_shared_connection_usable(): + with shared_pool_server("reset") as server: + pool = _async_pool() + envd_api = httpx.AsyncClient( + base_url=f"http://127.0.0.1:{server.port}", + transport=AsyncPyqwestTransport(pool), + ) + try: + await _break_async_stream( + make_async_client( + server.port, transport=PlainHTTPErrorTransport(pool) + ).connect(ConnectRequest()) + ) + assert await acheck_sandbox_health(envd_api) is True + assert len(server.connections) == 1 + server.assert_no_errors() + finally: + await envd_api.aclose() + + +async def test_async_dropped_connection_redials_for_the_health_probe(): + with shared_pool_server("drop") as server: + pool = _async_pool() + envd_api = httpx.AsyncClient( + base_url=f"http://127.0.0.1:{server.port}", + transport=AsyncPyqwestTransport(pool), + ) + try: + await _break_async_stream( + make_async_client( + server.port, transport=PlainHTTPErrorTransport(pool) + ).connect(ConnectRequest()) + ) + assert await acheck_sandbox_health(envd_api) is True + assert len(server.connections) == 2 + finally: + await envd_api.aclose() diff --git a/packages/python-sdk/tests/test_volume_client.py b/packages/python-sdk/tests/test_volume_client.py index 1c157a6900..6429daca2e 100644 --- a/packages/python-sdk/tests/test_volume_client.py +++ b/packages/python-sdk/tests/test_volume_client.py @@ -8,8 +8,13 @@ import pytest from pyqwest.httpx import AsyncPyqwestTransport, PyqwestTransport +from transport_caches import reset_transport_caches + +import e2b.api.client_async as api_client_async +import e2b.api.client_sync as api_client_sync import e2b.volume.client_async as client_async import e2b.volume.client_sync as client_sync +from e2b.connection_config import ConnectionConfig from e2b.exceptions import AuthenticationException from e2b.volume.client_async import get_api_client as get_async_api_client from e2b.volume.client_async import ( @@ -27,8 +32,8 @@ def reset_volume_transports(): - client_sync._transports.clear() - client_async._transports.clear() + # The volume clients draw from the SDK-wide pools in `e2b.api.client_*`. + reset_transport_caches() def test_sync_client_requires_volume_token(monkeypatch): @@ -90,6 +95,29 @@ def test_sync_transport_is_cached_per_proxy(): reset_volume_transports() +def test_volume_transports_are_the_shared_sdk_pools(test_api_key): + # The volume content API draws from the same pools as the control-plane + # REST API and the envd HTTP API — reqwest pools per host, so the volume + # host doesn't cost the process a pool of its own. Streamed downloads land + # in the streaming pool, whose 60s idle read bound the sandbox + # filesystem's streamed downloads ask for too. + reset_volume_transports() + config = VolumeConnectionConfig(token="vol-token") + api_config = ConnectionConfig(api_key=test_api_key) + + try: + assert get_sync_transport(config) is api_client_sync.get_transport(api_config) + assert get_sync_streaming_transport(config) is api_client_sync.get_transport( + api_config, for_streaming=True + ) + assert get_async_transport(config) is api_client_async.get_transport(api_config) + assert get_async_streaming_transport(config) is api_client_async.get_transport( + api_config, for_streaming=True + ) + finally: + reset_volume_transports() + + def test_sync_transport_is_shared_across_threads(): # pyqwest transports are thread-safe, so one transport (and its pool) # serves all threads — the per-thread caching this replaced is gone. diff --git a/packages/python-sdk/tests/transport_caches.py b/packages/python-sdk/tests/transport_caches.py new file mode 100644 index 0000000000..611ca6ab5e --- /dev/null +++ b/packages/python-sdk/tests/transport_caches.py @@ -0,0 +1,19 @@ +"""Reset the SDK's process-global pyqwest transport caches. + +Every HTTP stack in the SDK — control-plane REST, envd HTTP API, envd RPC, +volume content, template uploads — draws its connection pool from +``e2b.api.client_sync``/``client_async``, so one helper clears them all. Tests +that assert on pool identity or rebuild a pool with different tuning call this +before and after. Not a test module itself — imported by the transport test +modules (``pythonpath = tests`` in pytest.ini makes it importable under +``--import-mode=importlib``). +""" + +import e2b.api.client_async as api_client_async +import e2b.api.client_sync as api_client_sync + + +def reset_transport_caches() -> None: + for module in (api_client_sync, api_client_async): + module._transports.clear() + module._httpx_transports.clear() From 0d6033d3b6723b9792322d3e226cc3874f7cc4ee Mon Sep 17 00:00:00 2001 From: Mish Ushakov <10400064+mishushakov@users.noreply.github.com> Date: Tue, 11 Aug 2026 13:42:06 +0200 Subject: [PATCH 2/2] fix(python-sdk): keep streamed uploads off the retrying transport MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit To be able to replay a request, pyqwest's retry middleware copies a streamed request body in full as it is sent, so uploads were streamed to the wire *and* held whole in RAM (64 MiB file -> 70 MB peak). The copy is what makes the retries work — reqwest reads ahead into the body while connecting, so a connect error leaves the iterator already started and unrewindable — so uploads skip the layer rather than the layer skipping the copy. `get_upload_transport` is the httpx adapter over the same pool without the retries, and envd `files.write`, `volume.write_file` and template context uploads take it whenever their body is streamed. Sharing the pool keeps the connection reuse (and lets template uploads stop building a pool per build); what a streamed upload gives up is the connect retry, which fires before any of the body was written. Writes of in-memory data stay on the retrying transport, where a `bytes` body is replayed without a copy. Co-Authored-By: Claude Opus 5 (1M context) --- .../python-uploads-skip-connect-retries.md | 16 + .../e2b/api/client_async/__init__.py | 138 ++++--- .../e2b/api/client_sync/__init__.py | 138 ++++--- .../sandbox_async/filesystem/filesystem.py | 30 +- .../e2b/sandbox_sync/filesystem/filesystem.py | 30 +- .../e2b/template_async/build_api.py | 22 +- .../python-sdk/e2b/template_sync/build_api.py | 22 +- .../e2b/volume/client_async/__init__.py | 19 + .../e2b/volume/client_sync/__init__.py | 17 + .../python-sdk/e2b/volume/volume_async.py | 6 +- packages/python-sdk/e2b/volume/volume_sync.py | 6 +- .../async/template_async/test_upload_file.py | 27 +- .../sync/template_sync/test_upload_file.py | 27 +- .../python-sdk/tests/test_upload_transport.py | 339 ++++++++++++++++++ .../python-sdk/tests/test_volume_client.py | 35 ++ packages/python-sdk/tests/transport_caches.py | 2 + 16 files changed, 728 insertions(+), 146 deletions(-) create mode 100644 .changeset/python-uploads-skip-connect-retries.md create mode 100644 packages/python-sdk/tests/test_upload_transport.py diff --git a/.changeset/python-uploads-skip-connect-retries.md b/.changeset/python-uploads-skip-connect-retries.md new file mode 100644 index 0000000000..9b570c7215 --- /dev/null +++ b/.changeset/python-uploads-skip-connect-retries.md @@ -0,0 +1,16 @@ +--- +"@e2b/python-sdk": patch +--- + +Stop mirroring streamed file uploads in memory. To be able to replay a request, +pyqwest's retry middleware copies a streamed request body in full as it is sent, +so an upload was streamed to the wire *and* held whole in RAM: writing a 64 MiB +file with `volume.write_file`, `sandbox.files.write` or a template context upload +peaked at ~67 MB of it. A streamed upload now goes out on an httpx transport over +the same connection pool with the retry layer left out, so it keeps the pooled +connections while the body is no longer copied — a 32 MiB upload peaks at 1.5 MB +instead of 38 MB. What those uploads give up is the connect retry, which fires +before any of the body was written, so a failure to connect surfaces to the +caller as it would have anyway; writes of in-memory data are replayable for free +and keep their retries. Template context uploads no longer build a connection +pool of their own per build either. diff --git a/packages/python-sdk/e2b/api/client_async/__init__.py b/packages/python-sdk/e2b/api/client_async/__init__.py index 45f7bc1cc9..8580acc824 100644 --- a/packages/python-sdk/e2b/api/client_async/__init__.py +++ b/packages/python-sdk/e2b/api/client_async/__init__.py @@ -32,7 +32,13 @@ class ConnectionRetryTransport(RetryTransport): call or unary RPC like ``SendInput``). This matches the connect-only ``retries`` of the httpx transports this replaced; the retry middleware's default policy would otherwise also retry I/O errors and 429/5xx responses - for idempotent methods.""" + for idempotent methods. + + Retries cost memory on a streamed request body: to be able to replay it, + the middleware copies the body in full as it is sent (reqwest reads ahead + into it while connecting, so a connect error leaves the iterator already + started and unrewindable). File uploads therefore skip this layer and go + straight to the pool underneath — see :func:`get_upload_transport`.""" def should_retry_response( self, request: Request, response: Union[Response, Exception] @@ -56,25 +62,27 @@ def should_retry_response( # pyqwest's I/O runs on its own Rust runtime, so unlike the httpx transports # they replaced, the transports are not bound to an event loop and the caches # are process-global rather than per-loop. +_pools: Dict[TransportKey, HTTPTransport] = {} +# The connect-retrying stack over each pool. _transports: Dict[TransportKey, ConnectionRetryTransport] = {} -# The httpx adapter over each pool, shared by every httpx client on it. +# The httpx adapter over each pool, shared by every httpx client on it — one per +# pool with the retries, one without for uploads. _httpx_transports: Dict[TransportKey, AsyncPyqwestTransport] = {} +_upload_transports: Dict[TransportKey, AsyncPyqwestTransport] = {} -def get_pyqwest_transport( +def get_pool( proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None -) -> ConnectionRetryTransport: - """The shared pyqwest transport (= one connection pool) with the SDK's - tuning — system CA certs (without which TLS through an intercepting proxy - fails) and the httpx-equivalent pool limits — behind connect-only retries. - For TLS connections ALPN negotiates the HTTP version (HTTP/2 against the - E2B API and envd), like the http2-enabled httpx transports this replaced. - - Consumers speaking pyqwest natively (the envd RPC clients) take this; - consumers speaking httpx take :func:`get_httpx_transport`, the adapter over - the very same pool. Layer concerns above it rather than into it — the RPC - stack's plain-HTTP-error normalization wraps it, headers and codecs are - per-request — so that the pool stays shareable. +) -> HTTPTransport: + """The shared connection pool with the SDK's tuning — system CA certs + (without which TLS through an intercepting proxy fails) and the + httpx-equivalent pool limits. For TLS connections ALPN negotiates the HTTP + version (HTTP/2 against the E2B API and envd), like the http2-enabled httpx + transports this replaced. + + Layer concerns above it rather than into it — connect retries wrap it, so + does the RPC stack's plain-HTTP-error normalization, and headers and codecs + are per-request — so that the pool stays shareable. ``read_timeout`` bounds every read on the pool's connections, which is why it is part of the cache key: reqwest's read timer keeps running while a @@ -87,22 +95,37 @@ def get_pyqwest_transport( diagnostics httpcore used to provide. The SDK's own ``logger`` option is separate and sits above this, on the httpx client.""" key = (proxy, read_timeout) + with _transport_lock: + pool = _pools.get(key) + if pool is None: + pool = HTTPTransport( + tls_include_system_certs=True, + proxy=proxy.to_pyqwest() if proxy is not None else None, + pool_idle_timeout=pool_idle_timeout, + pool_max_idle_per_host=pool_max_idle_per_host, + read_timeout=read_timeout, + # Redirects belong to the httpx client above (which the + # generated clients leave off), not to reqwest. + follow_redirects=False, + ) + _pools[key] = pool + return pool + + +def get_pyqwest_transport( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> ConnectionRetryTransport: + """The shared pool of :func:`get_pool` behind connect-only retries, for + consumers speaking pyqwest natively (the envd RPC clients). Consumers + speaking httpx take :func:`get_httpx_transport`, the adapter over the very + same stack.""" + key = (proxy, read_timeout) + # Resolve the pool before taking the lock: it takes the same one. + pool = get_pool(proxy, read_timeout) with _transport_lock: transport = _transports.get(key) if transport is None: - transport = ConnectionRetryTransport( - HTTPTransport( - tls_include_system_certs=True, - proxy=proxy.to_pyqwest() if proxy is not None else None, - pool_idle_timeout=pool_idle_timeout, - pool_max_idle_per_host=pool_max_idle_per_host, - read_timeout=read_timeout, - # Redirects belong to the httpx client above (which the - # generated clients leave off), not to reqwest. - follow_redirects=False, - ), - max_retries=connection_retries, - ) + transport = ConnectionRetryTransport(pool, max_retries=connection_retries) _transports[key] = transport return transport @@ -110,24 +133,48 @@ def get_pyqwest_transport( def get_httpx_transport( proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None ) -> AsyncPyqwestTransport: - """The httpx adapter over the shared pool of + """The httpx adapter over the retrying stack of :func:`get_pyqwest_transport`, for the generated httpx clients (control plane, envd HTTP API, volume content). The adapter holds no state of its own and does not close the pool, so closing an httpx client leaves the pool intact for the other clients on it.""" key = (proxy, read_timeout) - # Resolve the pool before taking the lock: it takes the same one. - pool = get_pyqwest_transport(proxy, read_timeout) + stack = get_pyqwest_transport(proxy, read_timeout) with _transport_lock: transport = _httpx_transports.get(key) if transport is None: - transport = AsyncPyqwestTransport(pool) + transport = AsyncPyqwestTransport(stack) _httpx_transports[key] = transport return transport +def get_upload_transport( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> AsyncPyqwestTransport: + """The httpx adapter over the *same pool* as :func:`get_httpx_transport` + but without the retry layer, for file uploads: envd ``files.write``, volume + ``write_file`` and template context uploads. + + Those are the requests whose body is streamed rather than held in memory, + and the retry middleware makes a body replayable by copying it in full as + it is sent — so uploading a file on the retrying stack would mirror the + whole file in RAM (SDK-332). Skipping the layer, not the pool, keeps the + connection reuse: an upload still travels the sandbox's or the volume + host's pooled connections. What it gives up is the connect retry, which + fires before any of the body was written, so the caller sees the connection + error and can retry the upload itself.""" + key = (proxy, read_timeout) + pool = get_pool(proxy, read_timeout) + with _transport_lock: + transport = _upload_transports.get(key) + if transport is None: + transport = AsyncPyqwestTransport(pool) + _upload_transports[key] = transport + return transport + + def get_transport( - config: ConnectionConfig, *, for_streaming: bool = False + config: ConnectionConfig, *, for_streaming: bool = False, for_upload: bool = False ) -> AsyncPyqwestTransport: """The shared httpx transport for the control-plane REST API and the envd HTTP API (file transfers, health checks) — one pool serves both, keyed by @@ -138,16 +185,25 @@ def get_transport( long a streamed download may stall without limiting total transfer time. It is fixed per pool — the adapter's per-request timeouts are whole-request deadlines rather than idle bounds — so only streamed - downloads take it, and they get their own pool - (see :func:`get_pyqwest_transport`). + downloads take it, and they get their own pool (see :func:`get_pool`). + + ``for_upload`` selects the same pool without the retry layer, so that a + streamed upload is not copied into memory to keep it replayable + (see :func:`get_upload_transport`). """ - return get_httpx_transport( - proxy_to_config(config.proxy), READ_TIMEOUT if for_streaming else None - ) + proxy = proxy_to_config(config.proxy) + read_timeout = READ_TIMEOUT if for_streaming else None + if for_upload: + return get_upload_transport(proxy, read_timeout) + return get_httpx_transport(proxy, read_timeout) def get_envd_api( - config: ConnectionConfig, base_url: str, *, for_streaming: bool = False + config: ConnectionConfig, + base_url: str, + *, + for_streaming: bool = False, + for_upload: bool = False, ) -> httpx.AsyncClient: """An httpx client for a sandbox's envd HTTP API (file transfers, health checks) on the shared transports. The client itself is a cheap stateless @@ -155,7 +211,9 @@ def get_envd_api( is shared and loop-independent.""" return httpx.AsyncClient( base_url=base_url, - transport=get_transport(config, for_streaming=for_streaming), + transport=get_transport( + config, for_streaming=for_streaming, for_upload=for_upload + ), headers=config.sandbox_headers, event_hooks=make_async_logging_event_hooks(config.logger), ) diff --git a/packages/python-sdk/e2b/api/client_sync/__init__.py b/packages/python-sdk/e2b/api/client_sync/__init__.py index eb30c4cd26..a39968163d 100644 --- a/packages/python-sdk/e2b/api/client_sync/__init__.py +++ b/packages/python-sdk/e2b/api/client_sync/__init__.py @@ -32,7 +32,13 @@ class ConnectionRetryTransport(SyncRetryTransport): call or unary RPC like ``SendInput``). This matches the connect-only ``retries`` of the httpx transports this replaced; the retry middleware's default policy would otherwise also retry I/O errors and 429/5xx responses - for idempotent methods.""" + for idempotent methods. + + Retries cost memory on a streamed request body: to be able to replay it, + the middleware copies the body in full as it is sent (reqwest reads ahead + into it while connecting, so a connect error leaves the iterator already + started and unrewindable). File uploads therefore skip this layer and go + straight to the pool underneath — see :func:`get_upload_transport`.""" def should_retry_response( self, request: SyncRequest, response: Union[SyncResponse, Exception] @@ -55,25 +61,27 @@ def should_retry_response( # # pyqwest transports are thread-safe, so unlike the httpx transports they # replaced, the caches are process-global rather than per-thread. +_pools: Dict[TransportKey, SyncHTTPTransport] = {} +# The connect-retrying stack over each pool. _transports: Dict[TransportKey, ConnectionRetryTransport] = {} -# The httpx adapter over each pool, shared by every httpx client on it. +# The httpx adapter over each pool, shared by every httpx client on it — one per +# pool with the retries, one without for uploads. _httpx_transports: Dict[TransportKey, PyqwestTransport] = {} +_upload_transports: Dict[TransportKey, PyqwestTransport] = {} -def get_pyqwest_transport( +def get_pool( proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None -) -> ConnectionRetryTransport: - """The shared pyqwest transport (= one connection pool) with the SDK's - tuning — system CA certs (without which TLS through an intercepting proxy - fails) and the httpx-equivalent pool limits — behind connect-only retries. - For TLS connections ALPN negotiates the HTTP version (HTTP/2 against the - E2B API and envd), like the http2-enabled httpx transports this replaced. - - Consumers speaking pyqwest natively (the envd RPC clients) take this; - consumers speaking httpx take :func:`get_httpx_transport`, the adapter over - the very same pool. Layer concerns above it rather than into it — the RPC - stack's plain-HTTP-error normalization wraps it, headers and codecs are - per-request — so that the pool stays shareable. +) -> SyncHTTPTransport: + """The shared connection pool with the SDK's tuning — system CA certs + (without which TLS through an intercepting proxy fails) and the + httpx-equivalent pool limits. For TLS connections ALPN negotiates the HTTP + version (HTTP/2 against the E2B API and envd), like the http2-enabled httpx + transports this replaced. + + Layer concerns above it rather than into it — connect retries wrap it, so + does the RPC stack's plain-HTTP-error normalization, and headers and codecs + are per-request — so that the pool stays shareable. ``read_timeout`` bounds every read on the pool's connections, which is why it is part of the cache key: reqwest's read timer keeps running while a @@ -86,22 +94,37 @@ def get_pyqwest_transport( diagnostics httpcore used to provide. The SDK's own ``logger`` option is separate and sits above this, on the httpx client.""" key = (proxy, read_timeout) + with _transport_lock: + pool = _pools.get(key) + if pool is None: + pool = SyncHTTPTransport( + tls_include_system_certs=True, + proxy=proxy.to_pyqwest() if proxy is not None else None, + pool_idle_timeout=pool_idle_timeout, + pool_max_idle_per_host=pool_max_idle_per_host, + read_timeout=read_timeout, + # Redirects belong to the httpx client above (which the + # generated clients leave off), not to reqwest. + follow_redirects=False, + ) + _pools[key] = pool + return pool + + +def get_pyqwest_transport( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> ConnectionRetryTransport: + """The shared pool of :func:`get_pool` behind connect-only retries, for + consumers speaking pyqwest natively (the envd RPC clients). Consumers + speaking httpx take :func:`get_httpx_transport`, the adapter over the very + same stack.""" + key = (proxy, read_timeout) + # Resolve the pool before taking the lock: it takes the same one. + pool = get_pool(proxy, read_timeout) with _transport_lock: transport = _transports.get(key) if transport is None: - transport = ConnectionRetryTransport( - SyncHTTPTransport( - tls_include_system_certs=True, - proxy=proxy.to_pyqwest() if proxy is not None else None, - pool_idle_timeout=pool_idle_timeout, - pool_max_idle_per_host=pool_max_idle_per_host, - read_timeout=read_timeout, - # Redirects belong to the httpx client above (which the - # generated clients leave off), not to reqwest. - follow_redirects=False, - ), - max_retries=connection_retries, - ) + transport = ConnectionRetryTransport(pool, max_retries=connection_retries) _transports[key] = transport return transport @@ -109,24 +132,48 @@ def get_pyqwest_transport( def get_httpx_transport( proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None ) -> PyqwestTransport: - """The httpx adapter over the shared pool of + """The httpx adapter over the retrying stack of :func:`get_pyqwest_transport`, for the generated httpx clients (control plane, envd HTTP API, volume content). The adapter holds no state of its own and does not close the pool, so closing an httpx client leaves the pool intact for the other clients on it.""" key = (proxy, read_timeout) - # Resolve the pool before taking the lock: it takes the same one. - pool = get_pyqwest_transport(proxy, read_timeout) + stack = get_pyqwest_transport(proxy, read_timeout) with _transport_lock: transport = _httpx_transports.get(key) if transport is None: - transport = PyqwestTransport(pool) + transport = PyqwestTransport(stack) _httpx_transports[key] = transport return transport +def get_upload_transport( + proxy: Optional[ProxyConfig], read_timeout: Optional[float] = None +) -> PyqwestTransport: + """The httpx adapter over the *same pool* as :func:`get_httpx_transport` + but without the retry layer, for file uploads: envd ``files.write``, volume + ``write_file`` and template context uploads. + + Those are the requests whose body is streamed rather than held in memory, + and the retry middleware makes a body replayable by copying it in full as + it is sent — so uploading a file on the retrying stack would mirror the + whole file in RAM (SDK-332). Skipping the layer, not the pool, keeps the + connection reuse: an upload still travels the sandbox's or the volume + host's pooled connections. What it gives up is the connect retry, which + fires before any of the body was written, so the caller sees the connection + error and can retry the upload itself.""" + key = (proxy, read_timeout) + pool = get_pool(proxy, read_timeout) + with _transport_lock: + transport = _upload_transports.get(key) + if transport is None: + transport = PyqwestTransport(pool) + _upload_transports[key] = transport + return transport + + def get_transport( - config: ConnectionConfig, *, for_streaming: bool = False + config: ConnectionConfig, *, for_streaming: bool = False, for_upload: bool = False ) -> PyqwestTransport: """The shared httpx transport for the control-plane REST API and the envd HTTP API (file transfers, health checks) — one pool serves both, keyed by @@ -137,16 +184,25 @@ def get_transport( long a streamed download may stall without limiting total transfer time. It is fixed per pool — the adapter's per-request timeouts are whole-request deadlines rather than idle bounds — so only streamed - downloads take it, and they get their own pool - (see :func:`get_pyqwest_transport`). + downloads take it, and they get their own pool (see :func:`get_pool`). + + ``for_upload`` selects the same pool without the retry layer, so that a + streamed upload is not copied into memory to keep it replayable + (see :func:`get_upload_transport`). """ - return get_httpx_transport( - proxy_to_config(config.proxy), READ_TIMEOUT if for_streaming else None - ) + proxy = proxy_to_config(config.proxy) + read_timeout = READ_TIMEOUT if for_streaming else None + if for_upload: + return get_upload_transport(proxy, read_timeout) + return get_httpx_transport(proxy, read_timeout) def get_envd_api( - config: ConnectionConfig, base_url: str, *, for_streaming: bool = False + config: ConnectionConfig, + base_url: str, + *, + for_streaming: bool = False, + for_upload: bool = False, ) -> httpx.Client: """An httpx client for a sandbox's envd HTTP API (file transfers, health checks) on the shared transports. The client itself is a cheap stateless @@ -154,7 +210,9 @@ def get_envd_api( is shared and thread-safe.""" return httpx.Client( base_url=base_url, - transport=get_transport(config, for_streaming=for_streaming), + transport=get_transport( + config, for_streaming=for_streaming, for_upload=for_upload + ), headers=config.sandbox_headers, event_hooks=make_logging_event_hooks(config.logger), ) diff --git a/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py b/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py index 725662fad0..838f050b04 100644 --- a/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py +++ b/packages/python-sdk/e2b/sandbox_async/filesystem/filesystem.py @@ -105,6 +105,13 @@ def __init__( self._envd_api_streaming = get_envd_api( connection_config, envd_api_url, for_streaming=True ) + # Uploads go out on a third sibling, on the same pool without the + # connect retries: replaying a request body means copying it, so a + # streamed upload on the retrying transport would be mirrored in + # memory in full (see `get_upload_transport`). + self._envd_api_upload = get_envd_api( + connection_config, envd_api_url, for_upload=True + ) @overload async def read( @@ -405,8 +412,12 @@ async def _upload_file(file): headers["Content-Encoding"] = "gzip" is_streamed = not isinstance(file_data, (str, bytes)) + # Only a streamed body takes the non-retrying client: an + # in-memory one reaches the transport as `bytes`, which the + # retry layer replays without copying anything. + client = self._envd_api_upload if is_streamed else self._envd_api try: - r = await self._envd_api.post( + r = await client.post( ENVD_API_FILES_ROUTE, content=to_upload_body_async(file_data, gzip), headers=headers, @@ -448,19 +459,20 @@ async def _upload_file(file): if len(httpx_files) == 0: return [] + is_streamed = multipart_body_is_streamed(files) + client = self._envd_api_upload if is_streamed else self._envd_api try: - r = await self._envd_api.post( + r = await client.post( ENVD_API_FILES_ROUTE, files=httpx_files, params=params, headers=extra_headers, - # Only a streamed entry drops the deadline: httpx - # forwards binary `IOBase` entries in chunks, while text - # file-like data was buffered by `_to_httpx_file` (httpx - # rejects text-mode objects in multipart). - timeout=( - None if multipart_body_is_streamed(files) else upload_timeout - ), + # Only a streamed entry drops the deadline (and the retries, + # whose replay copy would mirror the body): httpx forwards + # binary `IOBase` entries in chunks, while text file-like + # data was buffered by `_to_httpx_file` (httpx rejects + # text-mode objects in multipart). + timeout=None if is_streamed else upload_timeout, ) except httpx.RemoteProtocolError as e: raise await ahandle_envd_api_transport_exception_with_health( diff --git a/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py b/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py index d3a2ba6463..95e975465a 100644 --- a/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py +++ b/packages/python-sdk/e2b/sandbox_sync/filesystem/filesystem.py @@ -101,6 +101,13 @@ def __init__( self._envd_api_streaming = get_envd_api( connection_config, envd_api_url, for_streaming=True ) + # Uploads go out on a third sibling, on the same pool without the + # connect retries: replaying a request body means copying it, so a + # streamed upload on the retrying transport would be mirrored in + # memory in full (see `get_upload_transport`). + self._envd_api_upload = get_envd_api( + connection_config, envd_api_url, for_upload=True + ) @overload def read( @@ -375,8 +382,12 @@ def write_files( headers["Content-Encoding"] = "gzip" is_streamed = not isinstance(file_data, (str, bytes)) + # Only a streamed body takes the non-retrying client: an + # in-memory one reaches the transport as `bytes`, which the + # retry layer replays without copying anything. + client = self._envd_api_upload if is_streamed else self._envd_api try: - r = self._envd_api.post( + r = client.post( ENVD_API_FILES_ROUTE, content=to_upload_body(file_data, gzip), headers=headers, @@ -412,19 +423,20 @@ def write_files( if len(httpx_files) == 0: return [] + is_streamed = multipart_body_is_streamed(files) + client = self._envd_api_upload if is_streamed else self._envd_api try: - r = self._envd_api.post( + r = client.post( ENVD_API_FILES_ROUTE, files=httpx_files, params=params, headers=extra_headers, - # Only a streamed entry drops the deadline: httpx - # forwards binary `IOBase` entries in chunks, while text - # file-like data was buffered by `_to_httpx_file` (httpx - # rejects text-mode objects in multipart). - timeout=( - None if multipart_body_is_streamed(files) else upload_timeout - ), + # Only a streamed entry drops the deadline (and the retries, + # whose replay copy would mirror the body): httpx forwards + # binary `IOBase` entries in chunks, while text file-like + # data was buffered by `_to_httpx_file` (httpx rejects + # text-mode objects in multipart). + timeout=None if is_streamed else upload_timeout, ) except httpx.RemoteProtocolError as e: raise handle_envd_api_transport_exception_with_health(e, self._envd_api) diff --git a/packages/python-sdk/e2b/template_async/build_api.py b/packages/python-sdk/e2b/template_async/build_api.py index 3d8c8f774e..c1890628f9 100644 --- a/packages/python-sdk/e2b/template_async/build_api.py +++ b/packages/python-sdk/e2b/template_async/build_api.py @@ -4,10 +4,9 @@ from typing import Callable, Optional, List, Union import httpx -from pyqwest import HTTPTransport -from pyqwest.httpx import AsyncPyqwestTransport from e2b.api import handle_api_exception, proxy_to_config +from e2b.api.client_async import get_upload_transport from e2b.io_utils import aiter_io_chunks from e2b.api.client.api.templates import ( post_v3_templates, @@ -128,25 +127,18 @@ async def upload_file( try: size = os.fstat(tar_file.fileno()).st_size + # The archive goes out on the shared connection pool, minus the + # connect retries: those would make the streamed body replayable by + # copying it, mirroring the whole build context in memory + # (SDK-332). Sharing the pool is what makes a per-build transport + # unnecessary. # Through the pyqwest adapter the upload timeout is a # whole-request deadline for the entire transfer, not a per-write # bound as with the httpx transport this replaced. async with httpx.AsyncClient( timeout=httpx.Timeout(upload_timeout), follow_redirects=api_client._follow_redirects, - transport=AsyncPyqwestTransport( - HTTPTransport( - tls_include_system_certs=True, - proxy=( - upload_proxy.to_pyqwest() - if upload_proxy is not None - else None - ), - # Redirects belong to the httpx client above, not to - # reqwest. - follow_redirects=False, - ) - ), + transport=get_upload_transport(upload_proxy), ) as client: # Stream the archive from disk via an async iterator. The # explicit Content-Length suppresses chunked transfer diff --git a/packages/python-sdk/e2b/template_sync/build_api.py b/packages/python-sdk/e2b/template_sync/build_api.py index 0525fa4570..5f9b9dc4dc 100644 --- a/packages/python-sdk/e2b/template_sync/build_api.py +++ b/packages/python-sdk/e2b/template_sync/build_api.py @@ -3,10 +3,9 @@ from typing import Callable, Optional, List, Union import httpx -from pyqwest import SyncHTTPTransport -from pyqwest.httpx import PyqwestTransport from e2b.api import handle_api_exception, proxy_to_config +from e2b.api.client_sync import get_upload_transport from e2b.api.client.api.templates import ( post_v3_templates, get_templates_template_id_files_hash, @@ -124,25 +123,18 @@ def upload_file( file_name, context_path, ignore_patterns, resolve_symlinks, gzip ) try: + # The archive goes out on the shared connection pool, minus the + # connect retries: those would make the streamed body replayable by + # copying it, mirroring the whole build context in memory + # (SDK-332). Sharing the pool is what makes a per-build transport + # unnecessary. # Through the pyqwest adapter the upload timeout is a # whole-request deadline for the entire transfer, not a per-write # bound as with the httpx transport this replaced. with httpx.Client( timeout=httpx.Timeout(upload_timeout), follow_redirects=api_client._follow_redirects, - transport=PyqwestTransport( - SyncHTTPTransport( - tls_include_system_certs=True, - proxy=( - upload_proxy.to_pyqwest() - if upload_proxy is not None - else None - ), - # Redirects belong to the httpx client above, not to - # reqwest. - follow_redirects=False, - ) - ), + transport=get_upload_transport(upload_proxy), ) as client: # httpx streams the archive from disk in chunks and sets # Content-Length from the file size—S3 presigned URLs reject diff --git a/packages/python-sdk/e2b/volume/client_async/__init__.py b/packages/python-sdk/e2b/volume/client_async/__init__.py index 3aed1272a5..ba0164c78f 100644 --- a/packages/python-sdk/e2b/volume/client_async/__init__.py +++ b/packages/python-sdk/e2b/volume/client_async/__init__.py @@ -6,6 +6,7 @@ proxy_to_config, ) from e2b.api.client_async import get_httpx_transport +from e2b.api.client_async import get_upload_transport as get_upload_httpx_transport from e2b.api.metadata import default_headers from e2b.exceptions import AuthenticationException from e2b.volume.client.client import AuthenticatedClient as AsyncVolumeApiClient @@ -25,6 +26,15 @@ def get_streaming_api_client( return _api_client(config, get_streaming_transport(config), **kwargs) +def get_upload_api_client( + config: VolumeConnectionConfig, **kwargs +) -> AsyncVolumeApiClient: + """The client for uploads: the same client on the upload transport, which + leaves the connect retries out so that a streamed body isn't copied into + memory to keep it replayable (see :func:`get_upload_transport`).""" + return _api_client(config, get_upload_transport(config), **kwargs) + + def _api_client( config: VolumeConnectionConfig, transport: AsyncPyqwestTransport, **kwargs ) -> AsyncVolumeApiClient: @@ -76,6 +86,15 @@ def get_transport(config: VolumeConnectionConfig) -> AsyncPyqwestTransport: return get_httpx_transport(proxy_to_config(config.proxy)) +def get_upload_transport(config: VolumeConnectionConfig) -> AsyncPyqwestTransport: + """The transport for uploads: the same pool as :func:`get_transport` minus + the connect retries. The retry layer makes a request body replayable by + copying it in full as it is sent, so a streamed `write_file` on it would be + mirrored in memory (see :func:`e2b.api.client_async.get_upload_transport`). + """ + return get_upload_httpx_transport(proxy_to_config(config.proxy)) + + def get_streaming_transport( config: VolumeConnectionConfig, ) -> AsyncPyqwestTransport: diff --git a/packages/python-sdk/e2b/volume/client_sync/__init__.py b/packages/python-sdk/e2b/volume/client_sync/__init__.py index 57344abbd2..d8d1d78863 100644 --- a/packages/python-sdk/e2b/volume/client_sync/__init__.py +++ b/packages/python-sdk/e2b/volume/client_sync/__init__.py @@ -6,6 +6,7 @@ proxy_to_config, ) from e2b.api.client_sync import get_httpx_transport +from e2b.api.client_sync import get_upload_transport as get_upload_httpx_transport from e2b.api.metadata import default_headers from e2b.exceptions import AuthenticationException from e2b.volume.client.client import AuthenticatedClient as VolumeApiClient @@ -25,6 +26,13 @@ def get_streaming_api_client( return _api_client(config, get_streaming_transport(config), **kwargs) +def get_upload_api_client(config: VolumeConnectionConfig, **kwargs) -> VolumeApiClient: + """The client for uploads: the same client on the upload transport, which + leaves the connect retries out so that a streamed body isn't copied into + memory to keep it replayable (see :func:`get_upload_transport`).""" + return _api_client(config, get_upload_transport(config), **kwargs) + + def _api_client( config: VolumeConnectionConfig, transport: PyqwestTransport, **kwargs ) -> VolumeApiClient: @@ -76,6 +84,15 @@ def get_transport(config: VolumeConnectionConfig) -> PyqwestTransport: return get_httpx_transport(proxy_to_config(config.proxy)) +def get_upload_transport(config: VolumeConnectionConfig) -> PyqwestTransport: + """The transport for uploads: the same pool as :func:`get_transport` minus + the connect retries. The retry layer makes a request body replayable by + copying it in full as it is sent, so a streamed `write_file` on it would be + mirrored in memory (see :func:`e2b.api.client_sync.get_upload_transport`). + """ + return get_upload_httpx_transport(proxy_to_config(config.proxy)) + + def get_streaming_transport(config: VolumeConnectionConfig) -> PyqwestTransport: """The transport for streamed downloads, carrying ``READ_TIMEOUT`` as the idle bound on every read: it resets after each successful read, so it caps diff --git a/packages/python-sdk/e2b/volume/volume_async.py b/packages/python-sdk/e2b/volume/volume_async.py index 7b7c7b62b0..22468b0338 100644 --- a/packages/python-sdk/e2b/volume/volume_async.py +++ b/packages/python-sdk/e2b/volume/volume_async.py @@ -39,6 +39,7 @@ from e2b.volume.client_async import get_api_client as get_volume_api_client from e2b.volume.client_async import ( get_streaming_api_client as get_streaming_volume_api_client, + get_upload_api_client as get_upload_volume_api_client, ) from e2b.volume.connection_config import ( VolumeApiParams, @@ -621,7 +622,10 @@ async def write_file( upload_timeout = VolumeConnectionConfig._get_request_timeout( FILE_TIMEOUT, opts.get("request_timeout") ) - api_client = get_volume_api_client(config) + # An upload streams its body, and the connect retries would make that + # body replayable by copying it in full — the whole file in memory + # (SDK-332) — so it goes out on the non-retrying client instead. + api_client = get_upload_volume_api_client(config) if upload_timeout is not None: api_client = api_client.with_timeout(httpx.Timeout(upload_timeout)) diff --git a/packages/python-sdk/e2b/volume/volume_sync.py b/packages/python-sdk/e2b/volume/volume_sync.py index 4415838fee..ada5b7156d 100644 --- a/packages/python-sdk/e2b/volume/volume_sync.py +++ b/packages/python-sdk/e2b/volume/volume_sync.py @@ -38,6 +38,7 @@ from e2b.volume.client_sync import get_api_client as get_volume_api_client from e2b.volume.client_sync import ( get_streaming_api_client as get_streaming_volume_api_client, + get_upload_api_client as get_upload_volume_api_client, ) from e2b.volume.connection_config import ( VolumeApiParams, @@ -589,7 +590,10 @@ def write_file( upload_timeout = VolumeConnectionConfig._get_request_timeout( FILE_TIMEOUT, opts.get("request_timeout") ) - api_client = get_volume_api_client(config) + # An upload streams its body, and the connect retries would make that + # body replayable by copying it in full — the whole file in memory + # (SDK-332) — so it goes out on the non-retrying client instead. + api_client = get_upload_volume_api_client(config) if upload_timeout is not None: api_client = api_client.with_timeout(httpx.Timeout(upload_timeout)) diff --git a/packages/python-sdk/tests/async/template_async/test_upload_file.py b/packages/python-sdk/tests/async/template_async/test_upload_file.py index a8d2993427..8f95f7c63a 100644 --- a/packages/python-sdk/tests/async/template_async/test_upload_file.py +++ b/packages/python-sdk/tests/async/template_async/test_upload_file.py @@ -7,6 +7,7 @@ import pytest from e2b.api.client.client import AuthenticatedClient +from e2b.api.client_async import get_upload_transport from e2b.template import utils as template_utils from e2b.exceptions import FileUploadException from e2b.template.consts import FILE_UPLOAD_TIMEOUT_SECONDS @@ -128,8 +129,9 @@ async def test_upload_file_leaves_redirects_to_httpx(tmp_path): assert state["paths"] == ["/redirect"] -async def _capture_upload_timeout(tmp_path, request_timeout=None): - """Run upload_file against a local server, capturing the httpx timeout.""" +async def _capture_upload_client_kwargs(tmp_path, request_timeout=None): + """Run upload_file against a local server, capturing how it built its + httpx client.""" (tmp_path / "hello.txt").write_text("hello world") server, thread, state = _make_server() @@ -140,7 +142,7 @@ async def _capture_upload_timeout(tmp_path, request_timeout=None): real_client = httpx.AsyncClient def spy_client(*args, **kwargs): - captured["timeout"] = kwargs.get("timeout") + captured.update(kwargs) return real_client(*args, **kwargs) try: @@ -167,20 +169,29 @@ def spy_client(*args, **kwargs): server.server_close() thread.join(timeout=5) - return captured["timeout"] + return captured async def test_upload_file_defaults_to_one_hour_timeout(tmp_path): # Uploads of large archives need far longer than the 60s general API # timeout, so the default upload timeout is 1 hour (matches the JS SDK). - timeout = await _capture_upload_timeout(tmp_path) - assert timeout == httpx.Timeout(FILE_UPLOAD_TIMEOUT_SECONDS) + captured = await _capture_upload_client_kwargs(tmp_path) + assert captured["timeout"] == httpx.Timeout(FILE_UPLOAD_TIMEOUT_SECONDS) async def test_upload_file_honors_explicit_request_timeout(tmp_path): # An explicitly set request_timeout overrides the 1-hour upload default. - timeout = await _capture_upload_timeout(tmp_path, request_timeout=5.0) - assert timeout == httpx.Timeout(5.0) + captured = await _capture_upload_client_kwargs(tmp_path, request_timeout=5.0) + assert captured["timeout"] == httpx.Timeout(5.0) + + +async def test_upload_file_uses_the_shared_pool_without_retries(tmp_path): + # The archive rides the SDK-wide connection pool instead of a transport + # built per build, but skips its retry layer: replaying a request means + # copying its body, which for a build context is the whole archive in + # memory (SDK-332). + captured = await _capture_upload_client_kwargs(tmp_path) + assert captured["transport"] is get_upload_transport(None) async def test_upload_file_ignores_post_upload_close_failure(tmp_path): diff --git a/packages/python-sdk/tests/sync/template_sync/test_upload_file.py b/packages/python-sdk/tests/sync/template_sync/test_upload_file.py index a08ae9906a..70e22a3f59 100644 --- a/packages/python-sdk/tests/sync/template_sync/test_upload_file.py +++ b/packages/python-sdk/tests/sync/template_sync/test_upload_file.py @@ -7,6 +7,7 @@ import pytest from e2b.api.client.client import AuthenticatedClient +from e2b.api.client_sync import get_upload_transport from e2b.template import utils as template_utils from e2b.exceptions import FileUploadException from e2b.template.consts import FILE_UPLOAD_TIMEOUT_SECONDS @@ -124,8 +125,9 @@ def test_upload_file_leaves_redirects_to_httpx(tmp_path): assert state["paths"] == ["/redirect"] -def _capture_upload_timeout(tmp_path, request_timeout=None): - """Run upload_file against a local server, capturing the httpx timeout.""" +def _capture_upload_client_kwargs(tmp_path, request_timeout=None): + """Run upload_file against a local server, capturing how it built its + httpx client.""" (tmp_path / "hello.txt").write_text("hello world") server, thread, state = _make_server() @@ -136,7 +138,7 @@ def _capture_upload_timeout(tmp_path, request_timeout=None): real_client = httpx.Client def spy_client(*args, **kwargs): - captured["timeout"] = kwargs.get("timeout") + captured.update(kwargs) return real_client(*args, **kwargs) try: @@ -163,20 +165,29 @@ def spy_client(*args, **kwargs): server.server_close() thread.join(timeout=5) - return captured["timeout"] + return captured def test_upload_file_defaults_to_one_hour_timeout(tmp_path): # Uploads of large archives need far longer than the 60s general API # timeout, so the default upload timeout is 1 hour (matches the JS SDK). - timeout = _capture_upload_timeout(tmp_path) - assert timeout == httpx.Timeout(FILE_UPLOAD_TIMEOUT_SECONDS) + captured = _capture_upload_client_kwargs(tmp_path) + assert captured["timeout"] == httpx.Timeout(FILE_UPLOAD_TIMEOUT_SECONDS) def test_upload_file_honors_explicit_request_timeout(tmp_path): # An explicitly set request_timeout overrides the 1-hour upload default. - timeout = _capture_upload_timeout(tmp_path, request_timeout=5.0) - assert timeout == httpx.Timeout(5.0) + captured = _capture_upload_client_kwargs(tmp_path, request_timeout=5.0) + assert captured["timeout"] == httpx.Timeout(5.0) + + +def test_upload_file_uses_the_shared_pool_without_retries(tmp_path): + # The archive rides the SDK-wide connection pool instead of a transport + # built per build, but skips its retry layer: replaying a request means + # copying its body, which for a build context is the whole archive in + # memory (SDK-332). + captured = _capture_upload_client_kwargs(tmp_path) + assert captured["transport"] is get_upload_transport(None) def test_upload_file_ignores_post_upload_close_failure(tmp_path): diff --git a/packages/python-sdk/tests/test_upload_transport.py b/packages/python-sdk/tests/test_upload_transport.py new file mode 100644 index 0000000000..1c58545c3f --- /dev/null +++ b/packages/python-sdk/tests/test_upload_transport.py @@ -0,0 +1,339 @@ +"""File uploads bypass the connect retries so their body isn't copied. + +The retry middleware makes a request body replayable by copying it in full as +it is sent — reqwest reads ahead into the body while connecting, so a connect +error leaves the iterator already started and unrewindable, and the copy is the +only way back. For a `volume.write_file`, an envd `files.write` or a template +context upload that means the whole file in RAM while it is also being streamed +to the wire (SDK-332): a 64 MiB upload peaked at 70 MB. + +So uploads take an httpx adapter over the same pool with the retry layer left +out. They keep the pooled connections and give up the connect retry, which +fires before any of the body was written and so surfaces to the caller intact. +""" + +import os +import tempfile +import threading +import tracemalloc +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +import httpx +import pytest +from packaging.version import Version +from pyqwest import HTTPTransport, SyncHTTPTransport +from transport_caches import reset_transport_caches + +from e2b.api.client_async import ConnectionRetryTransport +from e2b.api.client_async import get_envd_api as get_async_envd_api +from e2b.api.client_async import get_pool as get_async_pool +from e2b.api.client_async import get_transport as get_async_transport +from e2b.api.client_async import get_upload_transport as get_async_upload_transport +from e2b.api.client_sync import ( + ConnectionRetryTransport as SyncConnectionRetryTransport, +) +from e2b.api.client_sync import get_envd_api, get_pool, get_transport +from e2b.api.client_sync import get_upload_transport +from e2b.connection_config import ConnectionConfig +from e2b.io_utils import aiter_io_chunks +from e2b.sandbox_async.filesystem.filesystem import Filesystem as AsyncFilesystem +from e2b.sandbox_sync.filesystem.filesystem import Filesystem + +MIB = 1 << 20 +# Enough that a copy of the body would dwarf everything else the transfer +# allocates, while staying quick over the loopback interface. +UPLOAD_SIZE = 32 * MIB +# The streaming path allocates a chunk at a time (~1 MiB peak measured); the +# retrying transport's copy put this transfer's peak at 39 MB. +MAX_UPLOAD_PEAK = 8 * MIB + +ENVD_URL = "https://49999-sandbox.e2b.app" + + +def test_sync_upload_transport_is_the_same_pool_without_retries(test_api_key): + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key) + + try: + upload = get_upload_transport(None) + retrying = get_transport(config) + + # Same connection pool underneath, so an upload still travels the + # pooled connections; only the retry layer is missing. + assert isinstance(upload._transport, SyncHTTPTransport) + assert upload._transport is get_pool(None) + assert isinstance(retrying._transport, SyncConnectionRetryTransport) + assert retrying._transport._transport is upload._transport + + # Cached like every other transport, per proxy and idle read bound. + assert get_upload_transport(None) is upload + assert get_transport(config, for_upload=True) is upload + assert get_upload_transport(None, 60.0) is not upload + finally: + reset_transport_caches() + + +def test_async_upload_transport_is_the_same_pool_without_retries(test_api_key): + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key) + + try: + upload = get_async_upload_transport(None) + retrying = get_async_transport(config) + + assert isinstance(upload._transport, HTTPTransport) + assert upload._transport is get_async_pool(None) + assert isinstance(retrying._transport, ConnectionRetryTransport) + assert retrying._transport._transport is upload._transport + + assert get_async_upload_transport(None) is upload + assert get_async_transport(config, for_upload=True) is upload + finally: + reset_transport_caches() + + +def test_sync_envd_upload_client_uses_the_upload_transport(test_api_key): + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key, access_token="tok") + + client = get_envd_api(config, ENVD_URL, for_upload=True) + try: + assert client._transport is get_transport(config, for_upload=True) + finally: + client.close() + reset_transport_caches() + + +@pytest.mark.asyncio +async def test_async_envd_upload_client_uses_the_upload_transport(test_api_key): + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key, access_token="tok") + + client = get_async_envd_api(config, ENVD_URL, for_upload=True) + try: + assert client._transport is get_async_transport(config, for_upload=True) + finally: + await client.aclose() + reset_transport_caches() + + +def test_sync_filesystem_writes_go_to_the_upload_client(test_api_key): + # `files.write` posts on `_envd_api_upload`; the retrying client stays for + # the unary calls and health probes, the streaming one for downloads. + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key, access_token="tok") + envd_api = get_envd_api(config, ENVD_URL) + + filesystem = Filesystem(ENVD_URL, Version("0.7.0"), config, envd_api) + try: + assert filesystem._envd_api._transport is get_transport(config) + assert filesystem._envd_api_upload._transport is get_transport( + config, for_upload=True + ) + assert filesystem._envd_api_streaming._transport is get_transport( + config, for_streaming=True + ) + finally: + filesystem._envd_api_upload.close() + filesystem._envd_api_streaming.close() + envd_api.close() + reset_transport_caches() + + +@pytest.mark.asyncio +async def test_async_filesystem_writes_go_to_the_upload_client(test_api_key): + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key, access_token="tok") + envd_api = get_async_envd_api(config, ENVD_URL) + + filesystem = AsyncFilesystem(ENVD_URL, Version("0.7.0"), config, envd_api) + try: + assert filesystem._envd_api._transport is get_async_transport(config) + assert filesystem._envd_api_upload._transport is get_async_transport( + config, for_upload=True + ) + assert filesystem._envd_api_streaming._transport is get_async_transport( + config, for_streaming=True + ) + finally: + await filesystem._envd_api_upload.aclose() + await filesystem._envd_api_streaming.aclose() + await envd_api.aclose() + reset_transport_caches() + + +def _write_response(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, json=[{"name": "a.txt", "type": "file", "path": "/home/user/a.txt"}] + ) + + +def _routing_filesystem(config: ConnectionConfig, kind: type): + """A ``Filesystem`` whose two write clients record which one was used.""" + filesystem = kind(ENVD_URL, Version("0.7.0"), config, None) + used = [] + + def record(name): + def handler(request: httpx.Request) -> httpx.Response: + used.append(name) + return _write_response(request) + + return handler + + client = httpx.AsyncClient if kind is AsyncFilesystem else httpx.Client + filesystem._envd_api = client( + base_url=ENVD_URL, transport=httpx.MockTransport(record("retrying")) + ) + filesystem._envd_api_upload = client( + base_url=ENVD_URL, transport=httpx.MockTransport(record("upload")) + ) + return filesystem, used + + +@pytest.mark.parametrize("use_octet_stream", [True, False]) +def test_sync_only_streamed_writes_skip_the_retries( + test_api_key, tmp_path, use_octet_stream +): + # An in-memory body reaches the transport as `bytes` (or a tiny multipart + # stream), so the retry layer replays it for free and keeps its retries; a + # file-like one is the body that would be copied whole. + config = ConnectionConfig(api_key=test_api_key, access_token="tok") + filesystem, used = _routing_filesystem(config, Filesystem) + path = tmp_path / "a.txt" + path.write_text("streamed") + + try: + filesystem.write("/home/user/a.txt", "text", use_octet_stream=use_octet_stream) + with open(path, "rb") as file: + filesystem.write( + "/home/user/a.txt", file, use_octet_stream=use_octet_stream + ) + + assert used == ["retrying", "upload"] + finally: + filesystem._envd_api.close() + filesystem._envd_api_upload.close() + filesystem._envd_api_streaming.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_octet_stream", [True, False]) +async def test_async_only_streamed_writes_skip_the_retries( + test_api_key, tmp_path, use_octet_stream +): + config = ConnectionConfig(api_key=test_api_key, access_token="tok") + filesystem, used = _routing_filesystem(config, AsyncFilesystem) + path = tmp_path / "a.txt" + path.write_text("streamed") + + try: + await filesystem.write( + "/home/user/a.txt", "text", use_octet_stream=use_octet_stream + ) + with open(path, "rb") as file: + await filesystem.write( + "/home/user/a.txt", file, use_octet_stream=use_octet_stream + ) + + assert used == ["retrying", "upload"] + finally: + await filesystem._envd_api.aclose() + await filesystem._envd_api_upload.aclose() + await filesystem._envd_api_streaming.aclose() + + +class _UploadHandler(BaseHTTPRequestHandler): + """Reads a request body in small blocks — never holding it whole — and + records how much arrived.""" + + protocol_version = "HTTP/1.1" + received = None + + def do_PUT(self): + length = int(self.headers.get("Content-Length", 0)) + read = 0 + while read < length: + block = self.rfile.read(min(1 << 16, length - read)) + if not block: + break + read += len(block) + _UploadHandler.received = (read, self.headers.get("Transfer-Encoding")) + self.send_response(200) + self.send_header("Content-Length", "0") + self.end_headers() + + def log_message(self, *args): + pass + + +@pytest.fixture +def upload_server(): + _UploadHandler.received = None + server = ThreadingHTTPServer(("127.0.0.1", 0), _UploadHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + thread.join() + + +@pytest.fixture +def upload_file_path(): + path = os.path.join(tempfile.mkdtemp(), "context.tar.gz") + with open(path, "wb") as file: + for _ in range(UPLOAD_SIZE // MIB): + file.write(b"x" * MIB) + try: + yield path + finally: + os.remove(path) + + +def test_sync_upload_is_streamed_and_not_mirrored(upload_server, upload_file_path): + # The shape of a volume or template context upload: a file streamed with + # Content-Length framing (S3 presigned URLs reject chunked encoding), and + # nothing on the client holding on to what has already gone out. + reset_transport_caches() + client = httpx.Client(transport=get_upload_transport(None), timeout=None) + + try: + with open(upload_file_path, "rb") as file: + tracemalloc.start() + response = client.put(f"{upload_server}/upload", content=file) + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + assert response.status_code == 200 + assert _UploadHandler.received == (UPLOAD_SIZE, None) + assert peak < MAX_UPLOAD_PEAK + finally: + client.close() + reset_transport_caches() + + +@pytest.mark.asyncio +async def test_async_upload_is_streamed_and_not_mirrored( + upload_server, upload_file_path +): + reset_transport_caches() + client = httpx.AsyncClient(transport=get_async_upload_transport(None), timeout=None) + + try: + with open(upload_file_path, "rb") as file: + tracemalloc.start() + response = await client.put( + f"{upload_server}/upload", + content=aiter_io_chunks(file), + headers={"Content-Length": str(UPLOAD_SIZE)}, + ) + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + + assert response.status_code == 200 + assert _UploadHandler.received == (UPLOAD_SIZE, None) + assert peak < MAX_UPLOAD_PEAK + finally: + await client.aclose() + reset_transport_caches() diff --git a/packages/python-sdk/tests/test_volume_client.py b/packages/python-sdk/tests/test_volume_client.py index 6429daca2e..c24d62e797 100644 --- a/packages/python-sdk/tests/test_volume_client.py +++ b/packages/python-sdk/tests/test_volume_client.py @@ -21,11 +21,15 @@ get_streaming_transport as get_async_streaming_transport, ) from e2b.volume.client_async import get_transport as get_async_transport +from e2b.volume.client_async import ( + get_upload_transport as get_async_upload_transport, +) from e2b.volume.client_sync import get_api_client as get_sync_api_client from e2b.volume.client_sync import ( get_streaming_transport as get_sync_streaming_transport, ) from e2b.volume.client_sync import get_transport as get_sync_transport +from e2b.volume.client_sync import get_upload_transport as get_sync_upload_transport from e2b.volume.connection_config import VolumeConnectionConfig from e2b.volume.volume_async import AsyncVolume from e2b.volume.volume_sync import Volume @@ -118,6 +122,37 @@ def test_volume_transports_are_the_shared_sdk_pools(test_api_key): reset_volume_transports() +def test_volume_uploads_skip_the_connect_retries(test_api_key): + # `write_file` streams its body, and the retry layer would make that body + # replayable by copying it in full — the whole file in memory (SDK-332). It + # goes out on the same pool without that layer, which is also the SDK-wide + # upload transport. + reset_volume_transports() + config = VolumeConnectionConfig(token="vol-token") + api_config = ConnectionConfig(api_key=test_api_key) + + try: + assert get_sync_upload_transport(config) is api_client_sync.get_transport( + api_config, for_upload=True + ) + assert get_async_upload_transport(config) is api_client_async.get_transport( + api_config, for_upload=True + ) + assert get_sync_upload_transport(config) is not get_sync_transport(config) + assert get_async_upload_transport(config) is not get_async_transport(config) + + upload_client = client_sync.get_upload_api_client(config) + try: + assert ( + upload_client.get_httpx_client()._transport + is get_sync_upload_transport(config) + ) + finally: + upload_client.get_httpx_client().close() + finally: + reset_volume_transports() + + def test_sync_transport_is_shared_across_threads(): # pyqwest transports are thread-safe, so one transport (and its pool) # serves all threads — the per-thread caching this replaced is gone. diff --git a/packages/python-sdk/tests/transport_caches.py b/packages/python-sdk/tests/transport_caches.py index 611ca6ab5e..b244640ba5 100644 --- a/packages/python-sdk/tests/transport_caches.py +++ b/packages/python-sdk/tests/transport_caches.py @@ -15,5 +15,7 @@ def reset_transport_caches() -> None: for module in (api_client_sync, api_client_async): + module._pools.clear() module._transports.clear() module._httpx_transports.clear() + module._upload_transports.clear()