diff --git a/.changeset/python-unify-pyqwest-connection-pools.md b/.changeset/python-unify-pyqwest-connection-pools.md new file mode 100644 index 0000000000..eedc77f611 --- /dev/null +++ b/.changeset/python-unify-pyqwest-connection-pools.md @@ -0,0 +1,17 @@ +--- +"@e2b/python-sdk": patch +--- + +Run every persistent HTTP stack in the SDK on one shared pyqwest connection pool +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`, keyed on the three knobs that are fixed +when a pyqwest transport is built — proxy, idle read bound, and HTTP version. +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. No signature changes — +`get_transport` and `get_envd_transport` keep the `http2` parameter restored in +2.39.1, and the two are now the same pool per key rather than two. diff --git a/packages/python-sdk/e2b/api/client_async/__init__.py b/packages/python-sdk/e2b/api/client_async/__init__.py index e8992e581d..678cc29e6d 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,131 +40,163 @@ def should_retry_response( return isinstance(response, ConnectionError) -def retrying_http_transport( +_TransportKey = Tuple[Optional[ProxyConfig], Optional[float], bool] +"""Cache key of the shared transports: proxy, idle read bound, HTTP version. + +All three are fixed when a pyqwest transport is constructed, so each distinct +combination is necessarily its own pool.""" + +_transport_lock = threading.Lock() +# One pyqwest transport — one reqwest connection pool — per key; a `None` proxy +# 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, http2: bool = True, ) -> 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. - - ``http2=False`` pins the transport to HTTP/1.1; see :func:`get_transport` - for when that matters. + """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. + + 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 and ``http2`` + fixes the HTTP version; both are part of the cache key because they are + transport-construction knobs, so one pool cannot serve two values of + either. 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, - # `None` leaves the version to ALPN on TLS connections (HTTP/2 - # against the E2B API) and uses HTTP/1 for plaintext, like the - # http2-enabled httpx transport this replaced. - http_version=None if http2 else HTTPVersion.HTTP1, - # Redirects belong to the httpx client above (which the generated - # clients leave off), not to reqwest. - follow_redirects=False, - ), - max_retries=connection_retries, - ) + key = (proxy, read_timeout, http2) + 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, + # `None` leaves the version to ALPN on TLS connections + # (HTTP/2 against the E2B API and envd) and uses HTTP/1 for + # plaintext, like the http2-enabled httpx transport this + # replaced. + http_version=None if http2 else HTTPVersion.HTTP1, + # 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 -_transport_lock = threading.Lock() -# One transport (= one connection pool) per (proxy, http2) pair; a None proxy -# 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[Tuple[Optional[ProxyConfig], bool], AsyncPyqwestTransport] = {} +def get_httpx_transport( + proxy: Optional[ProxyConfig], + read_timeout: Optional[float] = None, + http2: bool = True, +) -> 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, http2) + # Resolve the pool before taking the lock: it takes the same one. + pool = get_pyqwest_transport(proxy, read_timeout, http2) + with _transport_lock: + transport = _httpx_transports.get(key) + if transport is None: + transport = AsyncPyqwestTransport(pool) + _httpx_transports[key] = transport + return transport def get_transport( - config: ConnectionConfig, http2: bool = True + config: ConnectionConfig, http2: bool = True, *, for_streaming: bool = False ) -> 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. + """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 TLS connections ALPN negotiates the HTTP + version (HTTP/2 against the E2B API), like the http2-enabled httpx + transport this replaced. ``http2=False`` returns a separate transport (its own pool) pinned to HTTP/1.1. That matters for a server that reacts to a client going away: HTTP/2 multiplexes requests over one connection, so abandoning a request only resets its stream and the server may never notice, while HTTP/1.1's one-connection-per-request closes the connection and the server observes - the disconnect.""" - proxy = proxy_to_config(config.proxy) - key = (proxy, http2) - with _transport_lock: - transport = _transports.get(key) - if transport is None: - transport = AsyncPyqwestTransport( - retrying_http_transport(proxy, http2=http2) - ) - _transports[key] = transport - return transport - - -# One transport per (proxy, http2, streaming) triple, separate from the REST -# API pools — envd traffic goes to per-sandbox hosts. -_envd_transports: Dict[ - Tuple[Optional[ProxyConfig], bool, bool], AsyncPyqwestTransport -] = {} + the disconnect. + + ``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`). + """ + return get_httpx_transport( + proxy_to_config(config.proxy), + READ_TIMEOUT if for_streaming else None, + http2, + ) def get_envd_transport( config: ConnectionConfig, http2: bool = True, *, 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). - - ``http2=False`` pins the transport to HTTP/1.1 — see + """The envd HTTP API's transport, which is :func:`get_transport` — the two + now share one pool per key, since reqwest pools per host and envd RPC and + the envd HTTP API hit the same sandbox host anyway. + + Kept only as a backward-compatible alias of :func:`get_transport` for any + external importer of the older public name; nothing inside the SDK calls it + (``e2b-code-interpreter`` imports :func:`get_transport` directly). Prefer :func:`get_transport`. + + :deprecated: Use :func:`get_transport` instead; will be removed in the next + major version. """ - proxy = proxy_to_config(config.proxy) - key = (proxy, http2, 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, - http2=http2, - ) - ) - _envd_transports[key] = transport - return transport + return get_transport(config, http2, for_streaming=for_streaming) 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 698b7cae12..73c279acec 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,126 +40,162 @@ def should_retry_response( return isinstance(response, ConnectionError) -def retrying_http_transport( +_TransportKey = Tuple[Optional[ProxyConfig], Optional[float], bool] +"""Cache key of the shared transports: proxy, idle read bound, HTTP version. + +All three are fixed when a pyqwest transport is constructed, so each distinct +combination is necessarily its own pool.""" + +_transport_lock = threading.Lock() +# One pyqwest transport — one reqwest connection pool — per key; a `None` proxy +# 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, http2: bool = True, ) -> 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. - - ``http2=False`` pins the transport to HTTP/1.1; see :func:`get_transport` - for when that matters. + """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. + + 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 and ``http2`` + fixes the HTTP version; both are part of the cache key because they are + transport-construction knobs, so one pool cannot serve two values of + either. 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, - # `None` leaves the version to ALPN on TLS connections (HTTP/2 - # against the E2B API) and uses HTTP/1 for plaintext, like the - # http2-enabled httpx transport this replaced. - http_version=None if http2 else HTTPVersion.HTTP1, - # Redirects belong to the httpx client above (which the generated - # clients leave off), not to reqwest. - follow_redirects=False, - ), - max_retries=connection_retries, - ) + key = (proxy, read_timeout, http2) + 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, + # `None` leaves the version to ALPN on TLS connections + # (HTTP/2 against the E2B API and envd) and uses HTTP/1 for + # plaintext, like the http2-enabled httpx transport this + # replaced. + http_version=None if http2 else HTTPVersion.HTTP1, + # 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 -_transport_lock = threading.Lock() -# One transport (= one connection pool) per (proxy, http2) pair; a None proxy -# 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[Tuple[Optional[ProxyConfig], bool], PyqwestTransport] = {} +def get_httpx_transport( + proxy: Optional[ProxyConfig], + read_timeout: Optional[float] = None, + http2: bool = True, +) -> 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, http2) + # Resolve the pool before taking the lock: it takes the same one. + pool = get_pyqwest_transport(proxy, read_timeout, http2) + with _transport_lock: + transport = _httpx_transports.get(key) + if transport is None: + transport = PyqwestTransport(pool) + _httpx_transports[key] = transport + return transport -def get_transport(config: ConnectionConfig, http2: bool = True) -> 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. +def get_transport( + config: ConnectionConfig, http2: bool = True, *, for_streaming: 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 + the connection's proxy. For TLS connections ALPN negotiates the HTTP + version (HTTP/2 against the E2B API), like the http2-enabled httpx + transport this replaced. ``http2=False`` returns a separate transport (its own pool) pinned to HTTP/1.1. That matters for a server that reacts to a client going away: HTTP/2 multiplexes requests over one connection, so abandoning a request only resets its stream and the server may never notice, while HTTP/1.1's one-connection-per-request closes the connection and the server observes - the disconnect.""" - proxy = proxy_to_config(config.proxy) - key = (proxy, http2) - with _transport_lock: - transport = _transports.get(key) - if transport is None: - transport = PyqwestTransport(retrying_http_transport(proxy, http2=http2)) - _transports[key] = transport - return transport - - -# One transport per (proxy, http2, streaming) triple, separate from the REST -# API pools — envd traffic goes to per-sandbox hosts. -_envd_transports: Dict[Tuple[Optional[ProxyConfig], bool, bool], PyqwestTransport] = {} + the disconnect. + + ``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`). + """ + return get_httpx_transport( + proxy_to_config(config.proxy), + READ_TIMEOUT if for_streaming else None, + http2, + ) def get_envd_transport( config: ConnectionConfig, http2: bool = True, *, 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). - - ``http2=False`` pins the transport to HTTP/1.1 — see + """The envd HTTP API's transport, which is :func:`get_transport` — the two + now share one pool per key, since reqwest pools per host and envd RPC and + the envd HTTP API hit the same sandbox host anyway. + + Kept only as a backward-compatible alias of :func:`get_transport` for any + external importer of the older public name; nothing inside the SDK calls it + (``e2b-code-interpreter`` imports :func:`get_transport` directly). Prefer :func:`get_transport`. + + :deprecated: Use :func:`get_transport` instead; will be removed in the next + major version. """ - proxy = proxy_to_config(config.proxy) - key = (proxy, http2, 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, - http2=http2, - ) - ) - _envd_transports[key] = transport - return transport + return get_transport(config, http2, for_streaming=for_streaming) 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/connection_config.py b/packages/python-sdk/e2b/connection_config.py index 803b2af64f..bad04d6eeb 100644 --- a/packages/python-sdk/e2b/connection_config.py +++ b/packages/python-sdk/e2b/connection_config.py @@ -26,6 +26,11 @@ # resets on each chunk, so it never limits total transfer time — only a # fully stalled stream. Matches the previous default stream idle timeout # (the request timeout). +# +# Kept equal to `e2b.volume.connection_config.READ_TIMEOUT` on purpose: the +# read bound is part of the transport cache key, so the volume streaming pool +# and the sandbox-filesystem streaming pool are the same reqwest pool only +# while the two constants agree. Change one and they silently split in two. READ_TIMEOUT: float = 60.0 # 60 seconds KEEPALIVE_PING_INTERVAL_SEC = 50 # 50 seconds 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/e2b/volume/connection_config.py b/packages/python-sdk/e2b/volume/connection_config.py index c7de6ca3bd..905badebee 100644 --- a/packages/python-sdk/e2b/volume/connection_config.py +++ b/packages/python-sdk/e2b/volume/connection_config.py @@ -19,6 +19,11 @@ # aborted when no bytes at all arrive for this long. It resets on each chunk, # so it never limits total transfer time — only a fully stalled connection. # Matches the JS SDK's default stream idle timeout (REQUEST_TIMEOUT_MS). +# +# Kept equal to `e2b.connection_config.READ_TIMEOUT` on purpose: the read bound +# is part of the transport cache key, so the volume streaming pool shares the +# sandbox-filesystem streaming pool only while the two constants agree. Change +# one and they silently split into two reqwest pools. READ_TIMEOUT: float = 60.0 # 60 seconds diff --git a/packages/python-sdk/tests/envd_frame_server.py b/packages/python-sdk/tests/envd_frame_server.py index df7f33447b..ebc3330580 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,148 @@ 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() + # Set by the test in "drop" mode once it has consumed the event written + # before the fault, so the RST cannot race the response head: an + # immediate RST lets hyper fail the in-flight request with the + # connection error instead of yielding the event it already received. + self.drop_when = threading.Event() + + 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: + # Only tear the connection down once the client has read the + # event just written: an immediate RST races the response + # head, and hyper then fails the request itself instead of + # yielding the event. + self.drop_when.wait(5) + # 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 c4a04ce991..0d9f89bc79 100644 --- a/packages/python-sdk/tests/test_api_client_transport.py +++ b/packages/python-sdk/tests/test_api_client_transport.py @@ -9,30 +9,26 @@ import httpx import pytest -from pyqwest import HTTPVersion +from pyqwest import HTTPVersion, Request, SyncRequest 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 +import e2b.api.client_async as api_client_async +import e2b.api.client_sync as api_client_sync +from e2b.api import pool_idle_timeout, pool_max_idle_per_host, proxy_to_config 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_pyqwest_transport as get_async_pyqwest_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_pyqwest_transport as get_sync_pyqwest_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() +from e2b.connection_config import READ_TIMEOUT, ConnectionConfig def run_in_worker_thread(fn): @@ -41,7 +37,7 @@ def run_in_worker_thread(fn): 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", @@ -57,11 +53,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", @@ -84,13 +80,13 @@ 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_transports_keyed_by_http_version(test_api_key): # The HTTP version is part of the cache key: without it, whichever caller # asked second would get a transport pinned to the other version. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) proxied_config = ConnectionConfig( api_key=test_api_key, @@ -105,7 +101,10 @@ def test_sync_transports_keyed_by_http_version(test_api_key): assert http1 is not negotiated assert envd_http1 is not envd_negotiated - assert envd_http1 is not http1 + # The envd HTTP API draws from the same pool as the control plane, per + # version — `get_envd_transport` is `get_transport` under another name. + assert envd_negotiated is negotiated + assert envd_http1 is http1 # Each version still has one pool per proxy, and repeat calls with the # same arguments reuse it. assert get_sync_transport(proxied_config, http2=False) not in ( @@ -120,7 +119,7 @@ def test_sync_transports_keyed_by_http_version(test_api_key): is not envd_http1 ) finally: - reset_sync_api_transports() + reset_transport_caches() def test_sync_transports_pass_http_version_to_pyqwest(test_api_key, monkeypatch): @@ -128,29 +127,65 @@ def test_sync_transports_pass_http_version_to_pyqwest(test_api_key, monkeypatch) # API), `HTTP1` pins HTTP/1.1. Which version was negotiated is only # observable over TLS — the local echo server is plaintext, where both # settings speak HTTP/1 — so assert what reaches the pyqwest transport. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) captured = [] - build_transport = client_sync.SyncHTTPTransport + build_transport = api_client_sync.SyncHTTPTransport def record(**kwargs): captured.append(kwargs["http_version"]) return build_transport(**kwargs) - monkeypatch.setattr(client_sync, "SyncHTTPTransport", record) + monkeypatch.setattr(api_client_sync, "SyncHTTPTransport", record) try: get_sync_transport(config) get_sync_transport(config, http2=False) - get_sync_envd_transport(config, http2=False) + # A third pool: same version as the call above, different idle bound. + # (`get_envd_transport(config, http2=False)` would be a cache hit and + # build nothing, since it shares the control plane's pool.) + get_sync_envd_transport(config, http2=False, for_streaming=True) assert captured == [None, HTTPVersion.HTTP1, HTTPVersion.HTTP1] finally: - reset_sync_api_transports() + reset_transport_caches() + + +def test_sync_transport_passes_pool_tuning_to_pyqwest(test_api_key, monkeypatch): + # The tuning that keeps a sandbox on one reused connection has to reach the + # pyqwest constructor: dropping any of it (`pool_max_idle_per_host=0`, no + # system CA certs, `follow_redirects=True`) would leave every identity and + # frame-level test green while a sandbox redialed on every request or TLS + # broke through an intercepting proxy. The identity assertions only prove + # one pool is reused, not how it was built. + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key) + captured = {} + build_transport = api_client_sync.SyncHTTPTransport + + def record(**kwargs): + captured.update(kwargs) + return build_transport(**kwargs) + + monkeypatch.setattr(api_client_sync, "SyncHTTPTransport", record) + + try: + # The streaming pool is the one carrying the idle read bound, so it + # pins `read_timeout` reaching the constructor as well. + get_sync_transport(config, for_streaming=True) + + assert captured["tls_include_system_certs"] is True + assert captured["proxy"] is None + assert captured["pool_idle_timeout"] == pool_idle_timeout + assert captured["pool_max_idle_per_host"] == pool_max_idle_per_host + assert captured["read_timeout"] == READ_TIMEOUT + assert captured["follow_redirects"] is False + finally: + 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) @@ -160,11 +195,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) @@ -174,34 +209,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) + streaming_transport = get_sync_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 - ) + 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) client = get_sync_envd_api(config, "https://sandbox.e2b.app") @@ -209,23 +240,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) @@ -237,12 +266,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", @@ -258,12 +287,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", @@ -279,12 +308,12 @@ 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 async def test_async_transports_keyed_by_http_version(test_api_key): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) try: @@ -295,7 +324,10 @@ async def test_async_transports_keyed_by_http_version(test_api_key): assert http1 is not negotiated assert envd_http1 is not envd_negotiated - assert envd_http1 is not http1 + # The envd HTTP API draws from the same pool as the control plane, per + # version — `get_envd_transport` is `get_transport` under another name. + assert envd_negotiated is negotiated + assert envd_http1 is http1 assert get_async_transport(config, http2=False) is http1 assert get_async_transport(config) is negotiated assert get_async_envd_transport(config, http2=False) is envd_http1 @@ -304,30 +336,57 @@ async def test_async_transports_keyed_by_http_version(test_api_key): is not envd_http1 ) finally: - reset_async_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_transports_pass_http_version_to_pyqwest(test_api_key, monkeypatch): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) captured = [] - build_transport = client_async.HTTPTransport + build_transport = api_client_async.HTTPTransport def record(**kwargs): captured.append(kwargs["http_version"]) return build_transport(**kwargs) - monkeypatch.setattr(client_async, "HTTPTransport", record) + monkeypatch.setattr(api_client_async, "HTTPTransport", record) try: get_async_transport(config) get_async_transport(config, http2=False) - get_async_envd_transport(config, http2=False) + # A third pool: same version as the call above, different idle bound. + get_async_envd_transport(config, http2=False, for_streaming=True) assert captured == [None, HTTPVersion.HTTP1, HTTPVersion.HTTP1] finally: - reset_async_api_transports() + reset_transport_caches() + + +@pytest.mark.asyncio +async def test_async_transport_passes_pool_tuning_to_pyqwest(test_api_key, monkeypatch): + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key) + captured = {} + build_transport = api_client_async.HTTPTransport + + def record(**kwargs): + captured.update(kwargs) + return build_transport(**kwargs) + + monkeypatch.setattr(api_client_async, "HTTPTransport", record) + + try: + get_async_transport(config, for_streaming=True) + + assert captured["tls_include_system_certs"] is True + assert captured["proxy"] is None + assert captured["pool_idle_timeout"] == pool_idle_timeout + assert captured["pool_max_idle_per_host"] == pool_max_idle_per_host + assert captured["read_timeout"] == READ_TIMEOUT + assert captured["follow_redirects"] is False + finally: + reset_transport_caches() @pytest.mark.asyncio @@ -336,7 +395,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) @@ -354,45 +413,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) + streaming_transport = get_async_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 - ) + 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) 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): @@ -460,7 +515,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( @@ -478,14 +533,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() @@ -503,11 +558,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() @@ -522,13 +577,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() @@ -544,12 +599,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() @@ -563,12 +618,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() @@ -583,14 +638,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() @@ -608,12 +663,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() @@ -631,13 +686,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() @@ -647,14 +702,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() @@ -664,7 +719,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( @@ -672,7 +727,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() @@ -682,14 +737,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() @@ -699,13 +754,13 @@ 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_http1_transport_round_trips(test_api_key, echo_server, caplog): # The HTTP/1.1-pinned transport is functional, not just configured: pinning # a version reqwest can't use for a request would fail at connect time. - reset_sync_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) client = httpx.Client( base_url=echo_server, transport=get_sync_transport(config, http2=False) @@ -723,12 +778,12 @@ def test_sync_http1_transport_round_trips(test_api_key, echo_server, caplog): assert f'GET {echo_server}/sandboxes "HTTP/1.0 200 OK"' in caplog.text finally: client.close() - reset_sync_api_transports() + reset_transport_caches() @pytest.mark.asyncio async def test_async_http1_transport_round_trips(test_api_key, echo_server): - reset_async_api_transports() + reset_transport_caches() config = ConnectionConfig(api_key=test_api_key) client = httpx.AsyncClient( base_url=echo_server, transport=get_async_transport(config, http2=False) @@ -741,7 +796,7 @@ async def test_async_http1_transport_round_trips(test_api_key, echo_server): assert response.json()["path"] == "/sandboxes" finally: await client.aclose() - reset_async_api_transports() + reset_transport_caches() def test_sync_transport_sends_multipart_bodies(test_api_key, echo_server): @@ -750,11 +805,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))]) @@ -762,4 +815,62 @@ 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() + + +def test_sync_closing_one_client_leaves_the_shared_pool_open(test_api_key, echo_server): + # Every stack draws on one pool now, so a close reaching it would take the + # whole process' HTTP down with it: pyqwest pools are closable + # (`SyncHTTPTransport.close`) and each httpx client holds the same cached + # adapter over one. The adapter forwards neither `close()` nor the + # context-manager exit the generated clients call, so closing one client + # must leave the others — and the pool the envd RPC stack talks to + # directly — working. + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) + api_httpx = get_sync_api_client(config).get_httpx_client() + envd_api = get_sync_envd_api(config, echo_server) + pool = get_sync_pyqwest_transport(proxy_to_config(config.proxy)) + + try: + assert api_httpx._transport is envd_api._transport + assert api_httpx.request("GET", "/sandboxes").status_code == 200 + + api_httpx.close() + + assert envd_api.get("/health").status_code == 200 + rpc_response = pool.execute_sync(SyncRequest("GET", f"{echo_server}/health")) + try: + assert rpc_response.status == 200 + finally: + rpc_response.close() + finally: + envd_api.close() + reset_transport_caches() + + +@pytest.mark.asyncio +async def test_async_closing_one_client_leaves_the_shared_pool_open( + test_api_key, echo_server +): + reset_transport_caches() + config = ConnectionConfig(api_key=test_api_key, api_url=echo_server) + api_httpx = get_async_api_client(config).get_async_httpx_client() + envd_api = get_async_envd_api(config, echo_server) + pool = get_async_pyqwest_transport(proxy_to_config(config.proxy)) + + try: + assert api_httpx._transport is envd_api._transport + assert (await api_httpx.request("GET", "/sandboxes")).status_code == 200 + + await api_httpx.aclose() + + assert (await envd_api.get("/health")).status_code == 200 + rpc_response = await pool.execute(Request("GET", f"{echo_server}/health")) + try: + assert rpc_response.status == 200 + finally: + await rpc_response.aclose() + finally: + await envd_api.aclose() + 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..f548e3dd6b --- /dev/null +++ b/packages/python-sdk/tests/test_shared_transport_pool.py @@ -0,0 +1,167 @@ +"""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, server) -> ConnectError: + """Read the first event, then the failure the server injects after it.""" + assert_stdout_event(next(events)) + # Release a "drop" server's RST only now that the event has been consumed, + # so the connection reset cannot race the response head (harmless for the + # "reset" server, which never waits on it). + server.drop_when.set() + 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, server) -> ConnectError: + assert_stdout_event(await events.__anext__()) + server.drop_when.set() + 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()), + server, + ) + 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()), + server, + ) + # 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()), + server, + ) + 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()), + server, + ) + 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..0ff5ce5f55 --- /dev/null +++ b/packages/python-sdk/tests/transport_caches.py @@ -0,0 +1,21 @@ +"""Reset the SDK's process-global pyqwest transport caches. + +Every persistent HTTP stack in the SDK — control-plane REST, envd HTTP API, +envd RPC, volume content — draws its connection pool from +``e2b.api.client_sync``/``client_async``, so one helper clears them all +(template uploads deliberately build their own non-retrying transport inline +and are not cached here). 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()