diff --git a/.changeset/python-pyqwest-volume-content.md b/.changeset/python-pyqwest-volume-content.md new file mode 100644 index 0000000000..b878e8f7b6 --- /dev/null +++ b/.changeset/python-pyqwest-volume-content.md @@ -0,0 +1,19 @@ +--- +"@e2b/python-sdk": minor +--- + +Move the volume content client (`Volume`/`AsyncVolume` file operations) onto +[`pyqwest`](https://pypi.org/project/pyqwest/) via its httpx-compatible +transport adapter, the same stack the REST API client uses. The connection +pool is shared process-wide per proxy instead of one pool per thread (sync) +or per event loop (async), and connection-establishment failures are retried +with backoff (`E2B_CONNECTION_RETRIES`, default 3), as before. + +For streamed volume reads (`Volume.read_file(format="stream")`), a stalled +stream is by default bounded by a transport-wide idle read timeout of +60 seconds that resets on every chunk (still surfaced as +`httpx.ReadTimeout`; matches the JS SDK's default stream idle timeout). +`AsyncVolume.read_file` keeps honoring an explicit `stream_idle_timeout` +per read (including `0` to disable); the sync client ignores it — it cannot +interrupt a blocking read. Passing `request_timeout` to a streamed read now +bounds the whole transfer rather than individual socket operations. diff --git a/packages/python-sdk/e2b/volume/client_async/__init__.py b/packages/python-sdk/e2b/volume/client_async/__init__.py index 27ee13d608..db11da292c 100644 --- a/packages/python-sdk/e2b/volume/client_async/__init__.py +++ b/packages/python-sdk/e2b/volume/client_async/__init__.py @@ -1,28 +1,27 @@ -import asyncio -import os -import weakref +import threading from typing import Dict, Optional import httpx -from httpx import Limits - -from e2b.api import connection_retries, make_async_logging_event_hooks +from pyqwest import HTTPTransport + +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 AsyncApiPyqwestTransport, ConnectionRetryTransport from e2b.api.metadata import default_headers -from e2b.connection_config import ProxyTypes from e2b.exceptions import AuthenticationException from e2b.volume.client.client import AuthenticatedClient as AsyncVolumeApiClient -from e2b.volume.connection_config import VolumeConnectionConfig - -limits = Limits( - max_keepalive_connections=int(os.getenv("E2B_MAX_KEEPALIVE_CONNECTIONS") or "20"), - max_connections=int(os.getenv("E2B_MAX_CONNECTIONS") or "2000"), - keepalive_expiry=int(os.getenv("E2B_KEEPALIVE_EXPIRY") or "300"), -) - -TransportKey = Optional[ProxyTypes] +from e2b.volume.connection_config import READ_TIMEOUT, VolumeConnectionConfig -def get_api_client(config: VolumeConnectionConfig, **kwargs) -> AsyncVolumeApiClient: +def get_api_client( + config: VolumeConnectionConfig, *, for_streaming: bool = False, **kwargs +) -> AsyncVolumeApiClient: if config.access_token is None: raise AuthenticationException( "Volume token is required for volume content operations. " @@ -49,42 +48,53 @@ def get_api_client(config: VolumeConnectionConfig, **kwargs) -> AsyncVolumeApiCl httpx_args={ # The proxy lives in the cached transport; passing `proxy` here too # would mount a fresh, never-closed proxy transport per client. - "transport": get_transport(config), + "transport": get_transport(config, for_streaming=for_streaming), "event_hooks": make_async_logging_event_hooks(config.logger), }, **kwargs, ) -class AsyncTransportWithLogger(httpx.AsyncHTTPTransport): - # Keyed weakly by the event loop object itself, not id(loop) — CPython - # reuses object ids, so a new loop could otherwise inherit a transport - # bound to a previous, closed loop. - _instances: weakref.WeakKeyDictionary[ - asyncio.AbstractEventLoop, - Dict[TransportKey, "AsyncTransportWithLogger"], - ] = weakref.WeakKeyDictionary() - - @property - def pool(self): - return self._pool - - -def get_transport(config: VolumeConnectionConfig) -> AsyncTransportWithLogger: - loop = asyncio.get_running_loop() - loop_instances = AsyncTransportWithLogger._instances.get(loop) - if loop_instances is None: - loop_instances = {} - AsyncTransportWithLogger._instances[loop] = loop_instances - - key: TransportKey = config.proxy - transport = loop_instances.get(key) - if transport is None: - transport = AsyncTransportWithLogger( - limits=limits, - proxy=config.proxy, - retries=connection_retries, - ) - loop_instances[key] = transport - - return transport +_transport_lock = threading.Lock() +# One transport (= one connection pool) per (proxy, streaming) pair; 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], bool], AsyncApiPyqwestTransport] = {} + + +def get_transport( + config: VolumeConnectionConfig, *, for_streaming: bool = False +) -> AsyncApiPyqwestTransport: + """The shared pyqwest-backed httpx transports for volume content API calls. + + 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, and the sync adapter does not bound body reads + at all. 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). + """ + proxy = proxy_to_config(config.proxy) + key = (proxy, for_streaming) + with _transport_lock: + transport = _transports.get(key) + if transport is None: + transport = AsyncApiPyqwestTransport( + 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 if for_streaming else None, + ), + max_retries=connection_retries, + ) + ) + _transports[key] = transport + return transport diff --git a/packages/python-sdk/e2b/volume/client_sync/__init__.py b/packages/python-sdk/e2b/volume/client_sync/__init__.py index 288f7b4983..862b10d156 100644 --- a/packages/python-sdk/e2b/volume/client_sync/__init__.py +++ b/packages/python-sdk/e2b/volume/client_sync/__init__.py @@ -1,27 +1,27 @@ -import os import threading from typing import Dict, Optional import httpx -from httpx import Limits +from pyqwest import SyncHTTPTransport -from e2b.api import connection_retries, make_logging_event_hooks +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 ApiPyqwestTransport, ConnectionRetryTransport from e2b.api.metadata import default_headers -from e2b.connection_config import ProxyTypes from e2b.exceptions import AuthenticationException from e2b.volume.client.client import AuthenticatedClient as VolumeApiClient -from e2b.volume.connection_config import VolumeConnectionConfig - -limits = Limits( - max_keepalive_connections=int(os.getenv("E2B_MAX_KEEPALIVE_CONNECTIONS") or "20"), - max_connections=int(os.getenv("E2B_MAX_CONNECTIONS") or "2000"), - keepalive_expiry=int(os.getenv("E2B_KEEPALIVE_EXPIRY") or "300"), -) - -TransportKey = Optional[ProxyTypes] +from e2b.volume.connection_config import READ_TIMEOUT, VolumeConnectionConfig -def get_api_client(config: VolumeConnectionConfig, **kwargs) -> VolumeApiClient: +def get_api_client( + config: VolumeConnectionConfig, *, for_streaming: bool = False, **kwargs +) -> VolumeApiClient: if config.access_token is None: raise AuthenticationException( "Volume token is required for volume content operations. " @@ -48,35 +48,52 @@ def get_api_client(config: VolumeConnectionConfig, **kwargs) -> VolumeApiClient: httpx_args={ # The proxy lives in the cached transport; passing `proxy` here too # would mount a fresh, never-closed proxy transport per client. - "transport": get_transport(config), + "transport": get_transport(config, for_streaming=for_streaming), "event_hooks": make_logging_event_hooks(config.logger), }, **kwargs, ) -class TransportWithLogger(httpx.HTTPTransport): - _thread_local = threading.local() +_transport_lock = threading.Lock() +# One transport (= one connection pool) per (proxy, streaming) pair; 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], bool], ApiPyqwestTransport] = {} - @property - def pool(self): - return self._pool +def get_transport( + config: VolumeConnectionConfig, *, for_streaming: bool = False +) -> ApiPyqwestTransport: + """The shared pyqwest-backed httpx transports for volume content API calls. -def get_transport(config: VolumeConnectionConfig) -> TransportWithLogger: - instances: Dict[TransportKey, TransportWithLogger] = getattr( - TransportWithLogger._thread_local, "instances", {} - ) - key: TransportKey = config.proxy - cached = instances.get(key) - if cached is not None: - return cached - - transport = TransportWithLogger( - limits=limits, - proxy=config.proxy, - retries=connection_retries, - ) - instances[key] = transport - TransportWithLogger._thread_local.instances = instances - return transport + 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, and the sync adapter does not bound body reads + at all. 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). + """ + proxy = proxy_to_config(config.proxy) + key = (proxy, for_streaming) + with _transport_lock: + transport = _transports.get(key) + if transport is None: + transport = ApiPyqwestTransport( + 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 if for_streaming else None, + ), + max_retries=connection_retries, + ) + ) + _transports[key] = transport + return transport diff --git a/packages/python-sdk/e2b/volume/connection_config.py b/packages/python-sdk/e2b/volume/connection_config.py index 2e902208f4..c7de6ca3bd 100644 --- a/packages/python-sdk/e2b/volume/connection_config.py +++ b/packages/python-sdk/e2b/volume/connection_config.py @@ -15,6 +15,12 @@ # bounds each chunk by the request timeout and leaves the total to the server.) FILE_TIMEOUT: float = 3600.0 # 1 hour +# Idle bound for every read on the volume content transports: the transfer is +# 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). +READ_TIMEOUT: float = 60.0 # 60 seconds + class VolumeApiParams(TypedDict, total=False): """ diff --git a/packages/python-sdk/e2b/volume/volume_async.py b/packages/python-sdk/e2b/volume/volume_async.py index 2e86ce4389..891788a59d 100644 --- a/packages/python-sdk/e2b/volume/volume_async.py +++ b/packages/python-sdk/e2b/volume/volume_async.py @@ -1,3 +1,5 @@ +import asyncio + from typing import AsyncIterator, IO, List, Literal, Optional, Union, cast, overload from http import HTTPStatus @@ -464,10 +466,12 @@ async def read_file( :param path: Path to the file :param format: Format of the file content—`text` by default :param stream_idle_timeout: Idle timeout in **seconds** for a streamed - read (`format="stream"`)—abort if no chunk arrives within this - window while reading. Resets on every chunk, so it bounds a stalled - stream without limiting total transfer time. Defaults to the request - timeout; pass `0` to disable. + read (`format="stream"`)—abort with `httpx.ReadTimeout` if the + response head or the next chunk doesn't arrive within this + window. Resets on every chunk, so it bounds a stalled stream + without limiting total transfer time. Defaults to the + transport-wide idle read timeout (60 seconds); pass `0` to + disable. :param opts: Connection options :return: File content as string, bytes, or async iterator of bytes @@ -481,38 +485,67 @@ async def read_file( ) if format == "stream": - # The request timeout bounds connection setup, not total transfer; - # consuming the body must not be killed by it. httpx's per-chunk - # `read` timeout becomes the idle-read timeout for the body - # (defaults to the request timeout), bounding a stalled stream - # without limiting total transfer time. Pass `0` to disable. - # Mirrors the sandbox files stream path. - idle_timeout = ( - timeout if stream_idle_timeout is None else stream_idle_timeout + # Through the pyqwest adapter a per-request timeout is a + # whole-request deadline that would kill long downloads, so a + # streamed read is sent with one only when the caller set + # `request_timeout` explicitly (making it the total-transfer + # deadline). + stream_timeout = VolumeConnectionConfig._get_request_timeout( + None, opts.get("request_timeout") + ) + # By default a stalled stream is bounded by the streaming + # transport's idle read timeout (see `get_transport`). An + # explicit `stream_idle_timeout` is applied per read with + # `wait_for` instead, on the regular transport — so values above + # the transport bound aren't capped by it and `0` disables idle + # bounding entirely. (The sync client can't interrupt a blocking + # read, so only the async flavor honors the per-call value.) + stream_client = get_volume_api_client( + config, for_streaming=stream_idle_timeout is None ) - stream_timeout = httpx.Timeout(timeout, read=idle_timeout or None) + + async def read_bounded(awaitable): + if not stream_idle_timeout: + return await awaitable + return await asyncio.wait_for(awaitable, stream_idle_timeout) async def stream_file() -> AsyncIterator[bytes]: - async with api_client.get_async_httpx_client().stream( - method="GET", - url=f"/volumecontent/{self._volume_id}/file", - params=params, - timeout=stream_timeout, - ) as response: - if response.status_code == 404: - raise NotFoundException(f"Path {path} not found") - - if response.status_code >= 300: - api_response = Response( - status_code=HTTPStatus(response.status_code), - content=await response.aread(), - headers=response.headers, - parsed=None, - ) - raise handle_api_exception(api_response, VolumeException) - - async for chunk in response.aiter_bytes(): - yield chunk + # pyqwest raises the builtin TimeoutError; keep the httpx + # exception the streamed-read contract established. + try: + stream_cm = stream_client.get_async_httpx_client().stream( + method="GET", + url=f"/volumecontent/{self._volume_id}/file", + params=params, + timeout=stream_timeout, + ) + response = await read_bounded(stream_cm.__aenter__()) + try: + if response.status_code == 404: + raise NotFoundException(f"Path {path} not found") + + if response.status_code >= 300: + api_response = Response( + status_code=HTTPStatus(response.status_code), + content=await response.aread(), + headers=response.headers, + parsed=None, + ) + raise handle_api_exception(api_response, VolumeException) + + chunks = response.aiter_bytes() + while True: + try: + chunk = await read_bounded(chunks.__anext__()) + except StopAsyncIteration: + break + yield chunk + finally: + await stream_cm.__aexit__(None, None, None) + # asyncio.TimeoutError is distinct from the builtin until 3.11, + # and wait_for and the adapter's whole-request deadline raise it. + except (TimeoutError, asyncio.TimeoutError) as e: + raise httpx.ReadTimeout(str(e)) from e return stream_file() diff --git a/packages/python-sdk/e2b/volume/volume_sync.py b/packages/python-sdk/e2b/volume/volume_sync.py index 5ca28d9dde..191f063e5e 100644 --- a/packages/python-sdk/e2b/volume/volume_sync.py +++ b/packages/python-sdk/e2b/volume/volume_sync.py @@ -462,11 +462,11 @@ def read_file( :param path: Path to the file :param format: Format of the file content—`text` by default - :param stream_idle_timeout: Idle timeout in **seconds** for a streamed - read (`format="stream"`)—abort if no chunk arrives within this - window while reading. Resets on every chunk, so it bounds a stalled - stream without limiting total transfer time. Defaults to the request - timeout; pass `0` to disable. + :param stream_idle_timeout: Ignored — the sync client cannot + interrupt a blocking read. A stalled streamed read is bounded by + a transport-wide idle read timeout instead (60 seconds), which + resets on every chunk. (`AsyncVolume.read_file` honors this + parameter.) :param opts: Connection options :return: File content as string, bytes, or iterator of bytes @@ -480,37 +480,46 @@ def read_file( ) if format == "stream": - # The request timeout bounds connection setup, not total transfer; - # consuming the body must not be killed by it. httpx's per-chunk - # `read` timeout becomes the idle-read timeout for the body - # (defaults to the request timeout), bounding a stalled stream - # without limiting total transfer time. Pass `0` to disable. - # Mirrors the sandbox files stream path. - idle_timeout = ( - timeout if stream_idle_timeout is None else stream_idle_timeout + # Through the pyqwest adapter a per-request timeout is a + # whole-request deadline that would kill long downloads, so a + # streamed read is sent with one only when the caller set + # `request_timeout` explicitly (making it the total-transfer + # deadline). A stalled stream is instead bounded by the + # transport-wide idle read timeout (see `get_transport`), which + # resets on every chunk without limiting total transfer time. + stream_timeout = VolumeConnectionConfig._get_request_timeout( + None, opts.get("request_timeout") ) - stream_timeout = httpx.Timeout(timeout, read=idle_timeout or None) + # The streaming transport carries the idle read timeout; the + # regular one must not (it would cut off slow uploads and + # responses), so streamed reads get their own client. + stream_client = get_volume_api_client(config, for_streaming=True) def stream_file() -> Iterator[bytes]: - with api_client.get_httpx_client().stream( - method="GET", - url=f"/volumecontent/{self._volume_id}/file", - params=params, - timeout=stream_timeout, - ) as response: - if response.status_code == 404: - raise NotFoundException(f"Path {path} not found") - - if response.status_code >= 300: - api_response = Response( - status_code=HTTPStatus(response.status_code), - content=response.read(), - headers=response.headers, - parsed=None, - ) - raise handle_api_exception(api_response, VolumeException) - - yield from response.iter_bytes() + # pyqwest raises the builtin TimeoutError; keep the httpx + # exception the streamed-read contract established. + try: + with stream_client.get_httpx_client().stream( + method="GET", + url=f"/volumecontent/{self._volume_id}/file", + params=params, + timeout=stream_timeout, + ) as response: + if response.status_code == 404: + raise NotFoundException(f"Path {path} not found") + + if response.status_code >= 300: + api_response = Response( + status_code=HTTPStatus(response.status_code), + content=response.read(), + headers=response.headers, + parsed=None, + ) + raise handle_api_exception(api_response, VolumeException) + + yield from response.iter_bytes() + except TimeoutError as e: + raise httpx.ReadTimeout(str(e)) from e return stream_file() diff --git a/packages/python-sdk/tests/test_volume_client.py b/packages/python-sdk/tests/test_volume_client.py index c3e7c231c1..ec8af11742 100644 --- a/packages/python-sdk/tests/test_volume_client.py +++ b/packages/python-sdk/tests/test_volume_client.py @@ -1,21 +1,28 @@ import asyncio -import gc +import socket import threading +import time +from typing import List import httpx import pytest +from pyqwest.httpx import AsyncPyqwestTransport, PyqwestTransport +import e2b.volume.client_async as client_async +import e2b.volume.client_sync as client_sync from e2b.exceptions import AuthenticationException -from e2b.volume.client_async import ( - AsyncTransportWithLogger as AsyncVolumeTransport, - get_api_client as get_async_api_client, - get_transport as get_async_transport, -) -from e2b.volume.client_sync import ( - get_api_client as get_sync_api_client, - get_transport as get_sync_transport, -) +from e2b.volume.client_async import get_api_client as get_async_api_client +from e2b.volume.client_async import get_transport as get_async_transport +from e2b.volume.client_sync import get_api_client as get_sync_api_client +from e2b.volume.client_sync import get_transport as get_sync_transport from e2b.volume.connection_config import VolumeConnectionConfig +from e2b.volume.volume_async import AsyncVolume +from e2b.volume.volume_sync import Volume + + +def reset_volume_transports(): + client_sync._transports.clear() + client_async._transports.clear() def test_sync_client_requires_volume_token(monkeypatch): @@ -61,85 +68,267 @@ async def run(): def test_sync_transport_is_cached_per_proxy(): + reset_volume_transports() config = VolumeConnectionConfig(token="vol-token") proxied = VolumeConnectionConfig(token="vol-token", proxy="http://127.0.0.1:8080") - transport_a = get_sync_transport(config) - transport_b = get_sync_transport(config) - transport_c = get_sync_transport(proxied) + try: + transport_a = get_sync_transport(config) + transport_b = get_sync_transport(config) + transport_c = get_sync_transport(proxied) - assert transport_a is transport_b - assert transport_a is not transport_c + assert isinstance(transport_a, PyqwestTransport) + assert transport_a is transport_b + assert transport_a is not transport_c + finally: + reset_volume_transports() -def test_sync_transport_is_not_shared_across_threads(): +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. + reset_volume_transports() config = VolumeConnectionConfig(token="vol-token") - main_transport = get_sync_transport(config) - result = {} + try: + main_transport = get_sync_transport(config) - def worker(): - result["transport"] = get_sync_transport(config) + result = {} - thread = threading.Thread(target=worker) - thread.start() - thread.join() + def worker(): + result["transport"] = get_sync_transport(config) - assert result["transport"] is not main_transport + thread = threading.Thread(target=worker) + thread.start() + thread.join() + assert result["transport"] is main_transport + finally: + reset_volume_transports() -def test_async_transport_is_cached_per_event_loop(): + +def test_async_transport_is_shared_across_loops(): + # pyqwest's I/O runs on its own Rust runtime, so the transport is not + # bound to an event loop — the per-loop caching this replaced is gone. + reset_volume_transports() config = VolumeConnectionConfig(token="vol-token") proxied = VolumeConnectionConfig(token="vol-token", proxy="http://127.0.0.1:8080") async def get_transports(): return get_async_transport(config), get_async_transport(config) - async def get_proxied_transport(): - return get_async_transport(proxied) - - loop_a = asyncio.new_event_loop() - loop_b = asyncio.new_event_loop() try: - transport_a1, transport_a2 = loop_a.run_until_complete(get_transports()) - transport_b1, _ = loop_b.run_until_complete(get_transports()) - proxied_a = loop_a.run_until_complete(get_proxied_transport()) + transport_a1, transport_a2 = asyncio.run(get_transports()) + transport_b1, _ = asyncio.run(get_transports()) + proxied_transport = get_async_transport(proxied) - # Same loop reuses the transport, another loop gets its own + assert isinstance(transport_a1, AsyncPyqwestTransport) assert transport_a1 is transport_a2 - assert transport_a1 is not transport_b1 + assert transport_a1 is transport_b1 - # Different proxy gets its own transport even on the same loop - assert proxied_a is not transport_a1 + # Different proxy still gets its own transport. + assert proxied_transport is not transport_a1 finally: - loop_a.close() - loop_b.close() + reset_volume_transports() + + +CHUNK = b"x" * 1024 + + +def _start_volume_file_server( + chunk_delays: List[float], ttfb_delay: float = 0.0 +) -> str: + """One-shot HTTP server streaming a chunked volume-file body, sleeping + ``ttfb_delay`` before the response head and ``chunk_delays[i]`` before + sending chunk ``i``. Returns its base URL.""" + sock = socket.socket() + sock.bind(("127.0.0.1", 0)) + sock.listen(1) + port = sock.getsockname()[1] + + def serve(): + try: + conn, _ = sock.accept() + while b"\r\n\r\n" not in conn.recv(65536): + pass + time.sleep(ttfb_delay) + conn.sendall( + b"HTTP/1.1 200 OK\r\n" + b"Content-Type: application/octet-stream\r\n" + b"Transfer-Encoding: chunked\r\n\r\n" + ) + for delay in chunk_delays: + time.sleep(delay) + conn.sendall(f"{len(CHUNK):x}\r\n".encode() + CHUNK + b"\r\n") + conn.sendall(b"0\r\n\r\n") + conn.close() + except OSError: + pass + finally: + sock.close() + + threading.Thread(target=serve, daemon=True).start() + return f"http://127.0.0.1:{port}" + + +@pytest.fixture +def short_read_timeout(monkeypatch): + """Rebuild the volume transports with a short idle read timeout.""" + reset_volume_transports() + monkeypatch.setattr(client_sync, "READ_TIMEOUT", 0.3) + monkeypatch.setattr(client_async, "READ_TIMEOUT", 0.3) + yield 0.3 + reset_volume_transports() + + +def test_sync_stream_survives_transfers_longer_than_read_timeout(short_read_timeout): + # The transport read timeout is an idle bound that resets on every chunk: + # a healthy stream whose total duration exceeds it must complete. + # `stream_idle_timeout` is ignored in the sync client (it cannot + # interrupt a blocking read) — a value shorter than every chunk gap must + # not abort the stream. + api_url = _start_volume_file_server([0.15] * 4) + volume = Volume(volume_id="v1", name="test", token="vol-token") + + stream = volume.read_file( + "file.bin", format="stream", stream_idle_timeout=0.01, api_url=api_url + ) + assert b"".join(stream) == CHUNK * 4 -def test_async_transport_not_reused_across_sequential_loops(): - AsyncVolumeTransport._instances.clear() - config = VolumeConnectionConfig(token="vol-token") +def test_sync_stream_stall_raises_read_timeout(short_read_timeout): + # A mid-body stall longer than the idle read timeout surfaces as + # httpx.ReadTimeout (pyqwest's builtin TimeoutError is remapped). + api_url = _start_volume_file_server([0.0, 5.0]) + volume = Volume(volume_id="v1", name="test", token="vol-token") + + stream = volume.read_file("file.bin", format="stream", api_url=api_url) + received = [next(iter(stream))] + with pytest.raises(httpx.ReadTimeout): + for chunk in stream: + received.append(chunk) + assert received == [CHUNK] + + +def test_async_stream_survives_transfers_longer_than_read_timeout(short_read_timeout): + api_url = _start_volume_file_server([0.15] * 4) + volume = AsyncVolume(volume_id="v1", name="test", token="vol-token") + + async def run(): + stream = await volume.read_file("file.bin", format="stream", api_url=api_url) + return b"".join([chunk async for chunk in stream]) + + assert asyncio.run(run()) == CHUNK * 4 + + +def test_async_stream_stall_raises_read_timeout(short_read_timeout): + api_url = _start_volume_file_server([0.0, 5.0]) + volume = AsyncVolume(volume_id="v1", name="test", token="vol-token") + + async def run(): + stream = await volume.read_file("file.bin", format="stream", api_url=api_url) + received = [await stream.__anext__()] + with pytest.raises(httpx.ReadTimeout): + async for chunk in stream: + received.append(chunk) + return received + + assert asyncio.run(run()) == [CHUNK] + - async def get_transport(): - return get_async_transport(config) +def test_async_explicit_stream_idle_timeout_aborts_stall(): + # An explicit stream_idle_timeout is honored per read with wait_for + # (like the JS SDK's streamIdleTimeoutMs) — no transport rebuild needed. + reset_volume_transports() + api_url = _start_volume_file_server([0.0, 5.0]) + volume = AsyncVolume(volume_id="v1", name="test", token="vol-token") + + async def run(): + stream = await volume.read_file( + "file.bin", format="stream", stream_idle_timeout=0.3, api_url=api_url + ) + received = [await stream.__anext__()] + with pytest.raises(httpx.ReadTimeout): + async for chunk in stream: + received.append(chunk) + return received - loop_a = asyncio.new_event_loop() try: - transport_a = loop_a.run_until_complete(get_transport()) + assert asyncio.run(run()) == [CHUNK] finally: - loop_a.close() - del loop_a - gc.collect() + reset_volume_transports() + + +def test_async_explicit_stream_idle_timeout_above_transport_bound(short_read_timeout): + # An explicit value larger than the transport's idle read timeout must + # not be capped by it: explicit values run on the regular transport. + api_url = _start_volume_file_server([short_read_timeout * 1.5] * 3) + volume = AsyncVolume(volume_id="v1", name="test", token="vol-token") + + async def run(): + stream = await volume.read_file( + "file.bin", format="stream", stream_idle_timeout=5.0, api_url=api_url + ) + return b"".join([chunk async for chunk in stream]) + + assert asyncio.run(run()) == CHUNK * 3 + + +def test_async_stream_idle_timeout_zero_disables_idle_bound(short_read_timeout): + # `stream_idle_timeout=0` disables idle bounding entirely — a stall + # longer than the transport's idle read timeout must not abort. + api_url = _start_volume_file_server([0.0, short_read_timeout * 3]) + volume = AsyncVolume(volume_id="v1", name="test", token="vol-token") - # The cache entry dies with the loop, so a later loop can never inherit - # a transport bound to a closed loop, even when CPython reuses the dead - # loop's object id. - assert len(AsyncVolumeTransport._instances) == 0 + async def run(): + stream = await volume.read_file( + "file.bin", format="stream", stream_idle_timeout=0, api_url=api_url + ) + return b"".join([chunk async for chunk in stream]) + + assert asyncio.run(run()) == CHUNK * 2 + + +def test_stream_transport_is_separate_from_regular_transport(): + # reqwest's read timer keeps running while a request body is sent and + # while waiting for the response head, so the idle read timeout lives on + # a dedicated streaming transport — putting it on the shared one would + # cut off uploads and slow unary responses longer than the idle bound. + reset_volume_transports() + config = VolumeConnectionConfig(token="vol-token") - loop_b = asyncio.new_event_loop() try: - transport_b = loop_b.run_until_complete(get_transport()) + regular = get_sync_transport(config) + streaming = get_sync_transport(config, for_streaming=True) + assert regular is not streaming + assert get_sync_transport(config) is regular + assert get_sync_transport(config, for_streaming=True) is streaming + + async_regular = get_async_transport(config) + async_streaming = get_async_transport(config, for_streaming=True) + assert async_regular is not async_streaming finally: - loop_b.close() + reset_volume_transports() + + +def test_sync_non_stream_read_survives_response_slower_than_idle_timeout( + short_read_timeout, +): + # Non-streamed requests go through the regular transport, which has no + # idle read timeout: a server that takes longer than the streaming idle + # bound to start responding must not be cut off. + api_url = _start_volume_file_server([0.0], ttfb_delay=short_read_timeout * 3) + volume = Volume(volume_id="v1", name="test", token="vol-token") + + assert volume.read_file("file.bin", format="bytes", api_url=api_url) == CHUNK + + +def test_sync_stream_response_head_is_bounded_by_idle_timeout(short_read_timeout): + # For streamed reads the idle read timeout also bounds waiting for the + # response head (like the JS SDK's handshake timeout on stream start). + api_url = _start_volume_file_server([0.0], ttfb_delay=5.0) + volume = Volume(volume_id="v1", name="test", token="vol-token") - assert transport_b is not transport_a + stream = volume.read_file("file.bin", format="stream", api_url=api_url) + with pytest.raises(httpx.ReadTimeout): + next(iter(stream))