11import asyncio
22import threading
3- import weakref
43from typing import Dict , Optional , Tuple , Union
54
65import httpx
1312 AsyncApiClient ,
1413 ProxyConfig ,
1514 connection_retries ,
16- limits ,
15+ make_async_logging_event_hooks ,
1716 pool_idle_timeout ,
1817 pool_max_idle_per_host ,
1918 proxy_to_config ,
2019)
21- from e2b .connection_config import ConnectionConfig , ProxyTypes
22-
23- TransportKey = Tuple [bool , Optional [ProxyTypes ]]
20+ from e2b .connection_config import READ_TIMEOUT , ConnectionConfig
2421
2522
2623def get_api_client (config : ConnectionConfig , ** kwargs ) -> AsyncApiClient :
@@ -69,13 +66,17 @@ def should_retry_response(
6966
7067
7168def retrying_http_transport (
72- proxy : Optional [ProxyConfig ],
69+ proxy : Optional [ProxyConfig ], read_timeout : Optional [ float ] = None
7370) -> ConnectionRetryTransport :
7471 """A fresh pyqwest transport (= its own connection pool) with the SDK's
7572 shared tuning — system CA certs (without which TLS through an
7673 intercepting proxy fails), the httpx-equivalent pool limits, and
77- connect-only retries. The REST API and envd RPC stacks each cache their
78- own instances (pool unification is a follow-up).
74+ connect-only retries. The REST API, envd RPC, and envd HTTP API stacks
75+ each cache their own instances (pool unification is a follow-up).
76+
77+ ``read_timeout`` bounds every read on the transport's connections; see
78+ :func:`get_envd_transport` for when that is (and isn't) appropriate.
79+
7980 Requests are logged by pyqwest itself on the ``pyqwest.access`` and
8081 ``pyqwest`` loggers at ``DEBUG`` (off unless enabled) — the transport-level
8182 diagnostics httpcore used to provide. The SDK's own ``logger`` option is
@@ -86,16 +87,17 @@ def retrying_http_transport(
8687 proxy = proxy .to_pyqwest () if proxy is not None else None ,
8788 pool_idle_timeout = pool_idle_timeout ,
8889 pool_max_idle_per_host = pool_max_idle_per_host ,
90+ read_timeout = read_timeout ,
8991 ),
9092 max_retries = connection_retries ,
9193 )
9294
9395
9496_transport_lock = threading .Lock ()
9597# One transport (= one connection pool) per proxy; None is the direct pool.
96- # pyqwest's I/O runs on its own Rust runtime, so unlike the httpx envd
97- # transports below , the transport is not bound to an event loop and the
98- # cache is process-global rather than per-loop.
98+ # pyqwest's I/O runs on its own Rust runtime, so unlike the httpx transports
99+ # they replaced , the transports are not bound to an event loop and the
100+ # caches are process-global rather than per-loop.
99101_transports : Dict [Optional [ProxyConfig ], "AsyncApiPyqwestTransport" ] = {}
100102
101103
@@ -112,38 +114,54 @@ def get_transport(config: ConnectionConfig) -> "AsyncApiPyqwestTransport":
112114 return transport
113115
114116
115- class AsyncEnvdTransportWithLogger (httpx .AsyncHTTPTransport ):
116- # Keyed weakly by the event loop object itself, not id(loop) — CPython
117- # reuses object ids, so a new loop could otherwise inherit a transport
118- # bound to a previous, closed loop.
119- _instances : weakref .WeakKeyDictionary [
120- asyncio .AbstractEventLoop ,
121- Dict [TransportKey , "AsyncEnvdTransportWithLogger" ],
122- ] = weakref .WeakKeyDictionary ()
123-
124- @property
125- def pool (self ):
126- return self ._pool
117+ # One transport per (proxy, streaming) pair, separate from the REST API
118+ # pools — envd traffic goes to per-sandbox hosts.
119+ _envd_transports : Dict [
120+ Tuple [Optional [ProxyConfig ], bool ], "AsyncApiPyqwestTransport"
121+ ] = {}
127122
128123
129124def get_envd_transport (
130- config : ConnectionConfig , http2 : bool = True
131- ) -> AsyncEnvdTransportWithLogger :
132- loop = asyncio .get_running_loop ()
133- loop_instances = AsyncEnvdTransportWithLogger ._instances .get (loop )
134- if loop_instances is None :
135- loop_instances = {}
136- AsyncEnvdTransportWithLogger ._instances [loop ] = loop_instances
137-
138- key : TransportKey = (http2 , config .proxy )
139- transport = loop_instances .get (key )
140- if transport is None :
141- transport = AsyncEnvdTransportWithLogger (
142- limits = limits ,
143- proxy = config .proxy ,
144- http2 = http2 ,
145- retries = connection_retries ,
146- )
147- loop_instances [key ] = transport
148-
149- return transport
125+ config : ConnectionConfig , * , for_streaming : bool = False
126+ ) -> "AsyncApiPyqwestTransport" :
127+ """The shared pyqwest-backed httpx transports for the envd HTTP API
128+ (file transfers, health checks).
129+
130+ The streaming transport carries ``read_timeout``, the idle bound on
131+ every read: it resets after each successful read, so it caps how long a
132+ streamed download may stall without limiting total transfer time. It is
133+ fixed per transport — the adapter's per-request timeouts are
134+ whole-request deadlines. Only streamed downloads use it: reqwest's read
135+ timer keeps running while a request body is sent and while waiting for
136+ the response head, so on the regular transport it would cut off uploads
137+ and slow unary responses longer than the idle bound (those stay bounded
138+ by their whole-request deadlines instead).
139+ """
140+ proxy = proxy_to_config (config .proxy )
141+ key = (proxy , for_streaming )
142+ with _transport_lock :
143+ transport = _envd_transports .get (key )
144+ if transport is None :
145+ transport = AsyncApiPyqwestTransport (
146+ retrying_http_transport (
147+ proxy ,
148+ read_timeout = READ_TIMEOUT if for_streaming else None ,
149+ )
150+ )
151+ _envd_transports [key ] = transport
152+ return transport
153+
154+
155+ def get_envd_api (
156+ config : ConnectionConfig , base_url : str , * , for_streaming : bool = False
157+ ) -> httpx .AsyncClient :
158+ """An httpx client for a sandbox's envd HTTP API (file transfers, health
159+ checks) on the shared pyqwest transports. The client itself is a cheap
160+ stateless wrapper — one per consumer is fine — while the pooled transport
161+ underneath is shared and loop-independent."""
162+ return httpx .AsyncClient (
163+ base_url = base_url ,
164+ transport = get_envd_transport (config , for_streaming = for_streaming ),
165+ headers = config .sandbox_headers ,
166+ event_hooks = make_async_logging_event_hooks (config .logger ),
167+ )
0 commit comments