diff --git a/src/runloop_api_client/_base_client.py b/src/runloop_api_client/_base_client.py index 88e0bbb3b..da359063d 100644 --- a/src/runloop_api_client/_base_client.py +++ b/src/runloop_api_client/_base_client.py @@ -8,6 +8,7 @@ import asyncio import inspect import logging +import secrets import weakref import platform import warnings @@ -24,6 +25,7 @@ Generic, Mapping, TypeVar, + Callable, Iterable, Iterator, Optional, @@ -78,8 +80,11 @@ DEFAULT_MAX_RETRIES, INITIAL_RETRY_DELAY, RAW_RESPONSE_HEADER, + DEFAULT_API_POOL_SHARDS, OVERRIDE_CAST_TO_HEADER, DEFAULT_CONNECTION_LIMITS, + DEFAULT_TRANSFER_POOL_SHARDS, + DEFAULT_BACKGROUND_POOL_SHARDS, ) from ._streaming import Stream, SSEDecoder, AsyncStream, SSEBytesDecoder from ._exceptions import ( @@ -98,6 +103,10 @@ # when the last user releases it. # The async transport is keyed by event loop because connections bind to the # loop that created them and cannot be reused across asyncio.run() calls. +# +# httpx/httpcore use a single H2 connection per transport, so each workload +# category uses a sharded Registry of SharedTransport instances. Per-client +# round-robin spreads requests across shards. _pool_lock = threading.Lock() @@ -134,6 +143,35 @@ def close(self) -> None: if should_close: self._transport.close() + class Registry: + """Process-global map of shard index → shared H2 transport.""" + + def __init__(self) -> None: + self._shards: dict[int, _SharedTransport] = {} + + def acquire(self, shard: int) -> _SharedTransport: + with _pool_lock: + existing = self._shards.get(shard) + if existing is not None and existing.acquire(): + return existing + transport = _SharedTransport( + httpx.HTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True), + ) + self._shards[shard] = transport + return transport + + def take_all(self) -> list[_SharedTransport]: + """Test-only: remove and return all transports for fixture cleanup.""" + with _pool_lock: + values = list(self._shards.values()) + self._shards.clear() + return values + + def shard_ids(self) -> set[int]: + """Test-only: which shard indexes currently hold a shared transport.""" + with _pool_lock: + return set(self._shards) + class _SharedAsyncTransport(httpx.AsyncBaseTransport): """Async refcounted wrapper: delegates to a real async transport.""" @@ -168,11 +206,150 @@ async def aclose(self) -> None: if should_close: await self._transport.aclose() + class Registry: + """Per-event-loop map of shard index → shared async H2 transport.""" + + def __init__(self) -> None: + self._by_loop: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, dict[int, _SharedAsyncTransport]] = ( + weakref.WeakKeyDictionary() + ) + + def acquire(self, loop: asyncio.AbstractEventLoop, shard: int) -> _SharedAsyncTransport: + with _pool_lock: + bucket = self._by_loop.get(loop) + if bucket is None: + bucket = {} + self._by_loop[loop] = bucket + existing = bucket.get(shard) + if existing is not None and existing.acquire(): + return existing + transport = _SharedAsyncTransport( + httpx.AsyncHTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True), + ) + bucket[shard] = transport + return transport + + def clear(self) -> None: + """Test-only: drop all per-loop shard maps (does not close transports).""" + with _pool_lock: + self._by_loop.clear() + + +# Process-global sharded registries (one ≈ H2 connection per shard index). +_shared_sync_api_transports = _SharedTransport.Registry() +_shared_sync_background_transports = _SharedTransport.Registry() +_shared_sync_transfer_transports = _SharedTransport.Registry() +_shared_async_api_transports = _SharedAsyncTransport.Registry() +_shared_async_background_transports = _SharedAsyncTransport.Registry() +_shared_async_transfer_transports = _SharedAsyncTransport.Registry() + +_BACKGROUND_PATH_SUFFIXES = ("/wait_for_status",) +_TRANSFER_PATH_SUFFIXES = ("/upload_file", "/download_file") + + +def _is_background_path(path: str) -> bool: + return path.endswith(_BACKGROUND_PATH_SUFFIXES) + + +def _is_transfer_path(path: str) -> bool: + return path.endswith(_TRANSFER_PATH_SUFFIXES) + + +class _SyncClientPool: + """Per-SDK-client round-robin over sharded httpx.Client wrappers.""" + + def __init__( + self, + *, + shards: int, + shared: bool, + registry: _SharedTransport.Registry | None, + make_client: Callable[[httpx.BaseTransport | None], httpx.Client], + ) -> None: + if shards < 1: + raise ValueError("shards must be >= 1") + self.shards = shards + self._shared = shared + self._registry = registry + self._make_client = make_client + self._clients: dict[int, httpx.Client] = {} + self._lock = threading.Lock() + self._next = secrets.randbelow(shards) + + def ensure(self, shard: int) -> httpx.Client: + existing = self._clients.get(shard) + if existing is not None: + return existing + with self._lock: + existing = self._clients.get(shard) + if existing is not None: + return existing + transport: httpx.BaseTransport | None = None + if self._shared and self._registry is not None: + transport = self._registry.acquire(shard) + client = self._make_client(transport) + self._clients[shard] = client + return client + + def next_client(self) -> httpx.Client: + with self._lock: + shard = self._next % self.shards + self._next += 1 + return self.ensure(shard) + + def close(self) -> None: + for client in self._clients.values(): + client.close() + self._clients.clear() + + +class _AsyncClientPool: + """Per-SDK-client round-robin over sharded httpx.AsyncClient wrappers.""" + + def __init__( + self, + *, + shards: int, + shared: bool, + registry: _SharedAsyncTransport.Registry | None, + make_client: Callable[[httpx.AsyncBaseTransport | None], httpx.AsyncClient], + ) -> None: + if shards < 1: + raise ValueError("shards must be >= 1") + self.shards = shards + self._shared = shared + self._registry = registry + self._make_client = make_client + self._clients: dict[int, httpx.AsyncClient] = {} + self._next = secrets.randbelow(shards) + + def ensure(self, shard: int) -> httpx.AsyncClient: + existing = self._clients.get(shard) + if existing is not None: + return existing + transport: httpx.AsyncBaseTransport | None = None + if self._shared and self._registry is not None: + try: + loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() + except RuntimeError: + loop = None + if loop is not None: + transport = self._registry.acquire(loop, shard) + client = self._make_client(transport) + self._clients[shard] = client + return client + + def next_client(self) -> httpx.AsyncClient: + # Single-threaded event loop: counter bump needs no lock when there is no await. + shard = self._next % self.shards + self._next += 1 + return self.ensure(shard) + + async def aclose(self) -> None: + for client in self._clients.values(): + await client.aclose() + self._clients.clear() -_shared_sync_transport: _SharedTransport | None = None -_shared_async_transports: weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, _SharedAsyncTransport] = ( - weakref.WeakKeyDictionary() -) # TODO: make base page type vars covariant SyncPageT = TypeVar("SyncPageT", bound="BaseSyncPage[Any]") @@ -929,8 +1106,15 @@ def __del__(self) -> None: class SyncAPIClient(BaseClient[httpx.Client, Stream[Any]]): _client: httpx.Client + _api_pool: _SyncClientPool | None + _background_pool: _SyncClientPool | None + _transfer_pool: _SyncClientPool | None _default_stream_cls: type[Stream[Any]] | None = None _uses_shared_pool: bool + _isolate_workload_pools: bool + _api_pool_shards: int + _background_pool_shards: int + _transfer_pool_shards: int _closed: bool def __init__( @@ -945,6 +1129,9 @@ def __init__( custom_query: Mapping[str, object] | None = None, _strict_response_validation: bool, shared_http_pool: bool = True, + api_pool_shards: int = DEFAULT_API_POOL_SHARDS, + background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, + transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, ) -> None: if not is_given(timeout): # if the user passed in a custom http client with a non-default @@ -964,6 +1151,13 @@ def __init__( f"Invalid `http_client` argument; Expected an instance of `httpx.Client` but got {type(http_client)}" ) + if api_pool_shards < 1: + raise ValueError("api_pool_shards must be >= 1") + if background_pool_shards < 1: + raise ValueError("background_pool_shards must be >= 1") + if transfer_pool_shards < 1: + raise ValueError("transfer_pool_shards must be >= 1") + super().__init__( version=version, # cast to a valid type because mypy doesn't understand our type narrowing @@ -976,29 +1170,72 @@ def __init__( ) self._closed = False + self._api_pool_shards = api_pool_shards + self._background_pool_shards = background_pool_shards + self._transfer_pool_shards = transfer_pool_shards + # Custom http_client owns the full transport stack; don't invent sibling pools. + self._isolate_workload_pools = http_client is None if http_client is not None: self._client = http_client self._uses_shared_pool = False - elif shared_http_pool: - global _shared_sync_transport - with _pool_lock: - if _shared_sync_transport is None or not _shared_sync_transport.acquire(): - _shared_sync_transport = _SharedTransport( - httpx.HTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True), - ) - self._client = SyncHttpxClientWrapper( - base_url=base_url, - timeout=cast(Timeout, timeout), - transport=_shared_sync_transport, - ) - self._uses_shared_pool = True - else: - self._client = SyncHttpxClientWrapper( - base_url=base_url, - timeout=cast(Timeout, timeout), + self._api_pool = None + self._background_pool = None + self._transfer_pool = None + return + + self._uses_shared_pool = shared_http_pool + + def make_client(transport: httpx.BaseTransport | None) -> httpx.Client: + timeout_ = cast(Timeout, self.timeout) + if transport is not None: + return SyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout_, + transport=transport, + ) + return SyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout_, ) - self._uses_shared_pool = False + + api_registry = _shared_sync_api_transports if shared_http_pool else None + bg_registry = _shared_sync_background_transports if shared_http_pool else None + xfer_registry = _shared_sync_transfer_transports if shared_http_pool else None + + self._api_pool = _SyncClientPool( + shards=api_pool_shards, + shared=shared_http_pool, + registry=api_registry, + make_client=make_client, + ) + self._background_pool = _SyncClientPool( + shards=background_pool_shards, + shared=shared_http_pool, + registry=bg_registry, + make_client=make_client, + ) + self._transfer_pool = _SyncClientPool( + shards=transfer_pool_shards, + shared=shared_http_pool, + registry=xfer_registry, + make_client=make_client, + ) + # Eager primary client (shard 0) for lifecycle / cookie-jar compatibility. + self._client = self._api_pool.ensure(0) + + def _get_client_for_path(self, path: str) -> httpx.Client: + if not self._isolate_workload_pools or self._api_pool is None: + return self._client + assert self._background_pool is not None and self._transfer_pool is not None + if _is_background_path(path): + return self._background_pool.next_client() + if _is_transfer_path(path): + return self._transfer_pool.next_client() + return self._api_pool.next_client() + + def _send_client_for_request(self, request: httpx.Request) -> httpx.Client: + return self._get_client_for_path(request.url.path) def is_closed(self) -> bool: return self._closed or self._client.is_closed @@ -1013,7 +1250,14 @@ def close(self) -> None: if self._closed: return self._closed = True - self._client.close() + if self._api_pool is not None: + # Closes _client (api shard 0) along with any other API shards. + self._api_pool.close() + assert self._background_pool is not None and self._transfer_pool is not None + self._background_pool.close() + self._transfer_pool.close() + else: + self._client.close() def __enter__(self: _T) -> _T: return self @@ -1114,7 +1358,7 @@ def request( response = None try: - response = self._client.send( + response = self._send_client_for_request(request).send( request, stream=stream or self._should_stream_response_body(request=request), **kwargs, @@ -1561,8 +1805,15 @@ def __del__(self) -> None: class AsyncAPIClient(BaseClient[httpx.AsyncClient, AsyncStream[Any]]): _client: httpx.AsyncClient + _api_pool: _AsyncClientPool | None + _background_pool: _AsyncClientPool | None + _transfer_pool: _AsyncClientPool | None _default_stream_cls: type[AsyncStream[Any]] | None = None _uses_shared_pool: bool + _isolate_workload_pools: bool + _api_pool_shards: int + _background_pool_shards: int + _transfer_pool_shards: int _closed: bool def __init__( @@ -1577,6 +1828,9 @@ def __init__( custom_headers: Mapping[str, str] | None = None, custom_query: Mapping[str, object] | None = None, shared_http_pool: bool = True, + api_pool_shards: int = DEFAULT_API_POOL_SHARDS, + background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, + transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, ) -> None: if not is_given(timeout): # if the user passed in a custom http client with a non-default @@ -1596,6 +1850,13 @@ def __init__( f"Invalid `http_client` argument; Expected an instance of `httpx.AsyncClient` but got {type(http_client)}" ) + if api_pool_shards < 1: + raise ValueError("api_pool_shards must be >= 1") + if background_pool_shards < 1: + raise ValueError("background_pool_shards must be >= 1") + if transfer_pool_shards < 1: + raise ValueError("transfer_pool_shards must be >= 1") + super().__init__( version=version, base_url=base_url, @@ -1608,43 +1869,79 @@ def __init__( ) self._closed = False + self._api_pool_shards = api_pool_shards + self._background_pool_shards = background_pool_shards + self._transfer_pool_shards = transfer_pool_shards + # Custom http_client owns the full transport stack; don't invent sibling pools. + self._isolate_workload_pools = http_client is None if http_client is not None: self._client = http_client self._uses_shared_pool = False - elif shared_http_pool: + self._api_pool = None + self._background_pool = None + self._transfer_pool = None + return + + # Async shared transports require a running loop; without one, fall back to private. + can_share = shared_http_pool + if shared_http_pool: try: - loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() + asyncio.get_running_loop() except RuntimeError: - loop = None - if loop is not None: - with _pool_lock: - existing = _shared_async_transports.get(loop) - if existing is not None and existing.acquire(): - transport: _SharedAsyncTransport = existing - else: - transport = _SharedAsyncTransport( - httpx.AsyncHTTPTransport(limits=DEFAULT_CONNECTION_LIMITS, http2=True), - ) - _shared_async_transports[loop] = transport - self._client = AsyncHttpxClientWrapper( - base_url=base_url, - timeout=cast(Timeout, timeout), + can_share = False + + self._uses_shared_pool = can_share + + def make_client(transport: httpx.AsyncBaseTransport | None) -> httpx.AsyncClient: + timeout_ = cast(Timeout, self.timeout) + if transport is not None: + return AsyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout_, transport=transport, ) - self._uses_shared_pool = True - else: - self._client = AsyncHttpxClientWrapper( - base_url=base_url, - timeout=cast(Timeout, timeout), - ) - self._uses_shared_pool = False - else: - self._client = AsyncHttpxClientWrapper( - base_url=base_url, - timeout=cast(Timeout, timeout), + return AsyncHttpxClientWrapper( + base_url=self._base_url, + timeout=timeout_, ) - self._uses_shared_pool = False + + api_registry = _shared_async_api_transports if can_share else None + bg_registry = _shared_async_background_transports if can_share else None + xfer_registry = _shared_async_transfer_transports if can_share else None + + self._api_pool = _AsyncClientPool( + shards=api_pool_shards, + shared=can_share, + registry=api_registry, + make_client=make_client, + ) + self._background_pool = _AsyncClientPool( + shards=background_pool_shards, + shared=can_share, + registry=bg_registry, + make_client=make_client, + ) + self._transfer_pool = _AsyncClientPool( + shards=transfer_pool_shards, + shared=can_share, + registry=xfer_registry, + make_client=make_client, + ) + self._client = self._api_pool.ensure(0) + + def _get_client_for_path(self, path: str) -> httpx.AsyncClient: + if not self._isolate_workload_pools or self._api_pool is None: + return self._client + assert self._background_pool is not None and self._transfer_pool is not None + if _is_background_path(path): + return self._background_pool.next_client() + if _is_transfer_path(path): + return self._transfer_pool.next_client() + return self._api_pool.next_client() + + def _send_client_for_request(self, request: httpx.Request) -> httpx.AsyncClient: + return self._get_client_for_path(request.url.path) def is_closed(self) -> bool: return self._closed or self._client.is_closed @@ -1659,7 +1956,13 @@ async def close(self) -> None: if self._closed: return self._closed = True - await self._client.aclose() + if self._api_pool is not None: + await self._api_pool.aclose() + assert self._background_pool is not None and self._transfer_pool is not None + await self._background_pool.aclose() + await self._transfer_pool.aclose() + else: + await self._client.aclose() async def __aenter__(self: _T) -> _T: return self @@ -1765,7 +2068,7 @@ async def request( response = None try: - response = await self._client.send( + response = await self._send_client_for_request(request).send( request, stream=stream or self._should_stream_response_body(request=request), **kwargs, diff --git a/src/runloop_api_client/_client.py b/src/runloop_api_client/_client.py index ed052ec87..bf32e59d6 100644 --- a/src/runloop_api_client/_client.py +++ b/src/runloop_api_client/_client.py @@ -26,6 +26,7 @@ ) from ._compat import cached_property from ._version import __version__ +from ._constants import DEFAULT_API_POOL_SHARDS, DEFAULT_TRANSFER_POOL_SHARDS, DEFAULT_BACKGROUND_POOL_SHARDS from ._streaming import Stream as Stream, AsyncStream as AsyncStream from ._exceptions import RunloopError, APIStatusError from ._base_client import ( @@ -96,6 +97,11 @@ def __init__( # Enables HTTP/2 multiplexing and avoids ConnectTimeout storms under high concurrency. # Set to False to create a private connection pool (old behavior). shared_http_pool: bool = True, + # Sharded H2 pools by workload (API / long-polls / transfers). Each shard ≈ + # one H2 connection; requests round-robin per client from a random offset. + api_pool_shards: int = DEFAULT_API_POOL_SHARDS, + background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, + transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, # Enable or disable schema validation for data returned by the API. # When enabled an error APIResponseValidationError is raised # if the API responds with invalid data for the expected schema. @@ -142,6 +148,9 @@ def __init__( custom_query=default_query, _strict_response_validation=_strict_response_validation, shared_http_pool=shared_http_pool, + api_pool_shards=api_pool_shards, + background_pool_shards=background_pool_shards, + transfer_pool_shards=transfer_pool_shards, ) self._idempotency_header = "x-request-id" @@ -284,6 +293,9 @@ def copy( timeout: float | Timeout | None | NotGiven = not_given, http_client: httpx.Client | None = None, shared_http_pool: bool | None = None, + api_pool_shards: int | None = None, + background_pool_shards: int | None = None, + transfer_pool_shards: int | None = None, max_retries: int | NotGiven = not_given, default_headers: Mapping[str, str] | None = None, set_default_headers: Mapping[str, str] | None = None, @@ -325,6 +337,13 @@ def copy( timeout=self.timeout if isinstance(timeout, NotGiven) else timeout, http_client=http_client, shared_http_pool=resolved_shared, + api_pool_shards=(api_pool_shards if api_pool_shards is not None else self._api_pool_shards), + background_pool_shards=( + background_pool_shards if background_pool_shards is not None else self._background_pool_shards + ), + transfer_pool_shards=( + transfer_pool_shards if transfer_pool_shards is not None else self._transfer_pool_shards + ), max_retries=max_retries if is_given(max_retries) else self.max_retries, default_headers=headers, default_query=params, @@ -390,6 +409,11 @@ def __init__( # Enables HTTP/2 multiplexing and avoids ConnectTimeout storms under high concurrency. # Set to False to create a private connection pool (old behavior). shared_http_pool: bool = True, + # Sharded H2 pools by workload (API / long-polls / transfers). Each shard ≈ + # one H2 connection; requests round-robin per client from a random offset. + api_pool_shards: int = DEFAULT_API_POOL_SHARDS, + background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, + transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, # Enable or disable schema validation for data returned by the API. # When enabled an error APIResponseValidationError is raised # if the API responds with invalid data for the expected schema. @@ -436,6 +460,9 @@ def __init__( custom_query=default_query, _strict_response_validation=_strict_response_validation, shared_http_pool=shared_http_pool, + api_pool_shards=api_pool_shards, + background_pool_shards=background_pool_shards, + transfer_pool_shards=transfer_pool_shards, ) self._idempotency_header = "x-request-id" @@ -578,6 +605,9 @@ def copy( timeout: float | Timeout | None | NotGiven = not_given, http_client: httpx.AsyncClient | None = None, shared_http_pool: bool | None = None, + api_pool_shards: int | None = None, + background_pool_shards: int | None = None, + transfer_pool_shards: int | None = None, max_retries: int | NotGiven = not_given, default_headers: Mapping[str, str] | None = None, set_default_headers: Mapping[str, str] | None = None, @@ -619,6 +649,13 @@ def copy( timeout=self.timeout if isinstance(timeout, NotGiven) else timeout, http_client=http_client, shared_http_pool=resolved_shared, + api_pool_shards=(api_pool_shards if api_pool_shards is not None else self._api_pool_shards), + background_pool_shards=( + background_pool_shards if background_pool_shards is not None else self._background_pool_shards + ), + transfer_pool_shards=( + transfer_pool_shards if transfer_pool_shards is not None else self._transfer_pool_shards + ), max_retries=max_retries if is_given(max_retries) else self.max_retries, default_headers=headers, default_query=params, diff --git a/src/runloop_api_client/_constants.py b/src/runloop_api_client/_constants.py index 83506afe3..6061b7708 100644 --- a/src/runloop_api_client/_constants.py +++ b/src/runloop_api_client/_constants.py @@ -10,6 +10,13 @@ DEFAULT_MAX_RETRIES = 5 DEFAULT_CONNECTION_LIMITS = httpx.Limits(max_connections=20, max_keepalive_connections=10) +# Separate H2 connection pools by workload. Each "shard" is its own httpx +# client / transport (≈ one H2 connection). Defaults sized for Jetty's ~128 +# streams/connection: enough API/wait concurrency without oversized transfer. +DEFAULT_API_POOL_SHARDS = 8 +DEFAULT_BACKGROUND_POOL_SHARDS = 16 +DEFAULT_TRANSFER_POOL_SHARDS = 2 + INITIAL_RETRY_DELAY = 1.0 MAX_RETRY_DELAY = 60.0 diff --git a/src/runloop_api_client/sdk/async_.py b/src/runloop_api_client/sdk/async_.py index 6319e02bd..a23f35e2d 100644 --- a/src/runloop_api_client/sdk/async_.py +++ b/src/runloop_api_client/sdk/async_.py @@ -42,6 +42,7 @@ from .._client import DEFAULT_MAX_RETRIES, AsyncRunloop from ._helpers import detect_content_type from .async_axon import AsyncAxon +from .._constants import DEFAULT_API_POOL_SHARDS, DEFAULT_TRANSFER_POOL_SHARDS, DEFAULT_BACKGROUND_POOL_SHARDS from .async_agent import AsyncAgent from .async_devbox import AsyncDevbox from .async_scorer import AsyncScorer @@ -1329,6 +1330,9 @@ def __init__( default_headers: Mapping[str, str] | None = None, default_query: Mapping[str, object] | None = None, http_client: httpx.AsyncClient | None = None, + api_pool_shards: int = DEFAULT_API_POOL_SHARDS, + background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, + transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, ) -> None: """Configure the asynchronous SDK wrapper. @@ -1346,6 +1350,12 @@ def __init__( :type default_query: Mapping[str, object] | None, optional :param http_client: Custom ``httpx.AsyncClient`` instance to reuse, defaults to None :type http_client: httpx.AsyncClient | None, optional + :param api_pool_shards: H2 shards for short RPCs (round-robin), defaults to 8 + :type api_pool_shards: int, optional + :param background_pool_shards: H2 shards for long-polls (round-robin), defaults to 16 + :type background_pool_shards: int, optional + :param transfer_pool_shards: H2 shards for upload/download (round-robin), defaults to 2 + :type transfer_pool_shards: int, optional """ self.api = AsyncRunloop( bearer_token=bearer_token, @@ -1355,6 +1365,9 @@ def __init__( default_headers=default_headers, default_query=default_query, http_client=http_client, + api_pool_shards=api_pool_shards, + background_pool_shards=background_pool_shards, + transfer_pool_shards=transfer_pool_shards, ) self.agent = AsyncAgentOps(self.api) diff --git a/src/runloop_api_client/sdk/sync.py b/src/runloop_api_client/sdk/sync.py index 7d3afcb24..ce10d1457 100644 --- a/src/runloop_api_client/sdk/sync.py +++ b/src/runloop_api_client/sdk/sync.py @@ -50,6 +50,7 @@ from .benchmark import Benchmark from .blueprint import Blueprint from .mcp_config import McpConfig +from .._constants import DEFAULT_API_POOL_SHARDS, DEFAULT_TRANSFER_POOL_SHARDS, DEFAULT_BACKGROUND_POOL_SHARDS from .gateway_config import GatewayConfig from .network_policy import NetworkPolicy from .storage_object import StorageObject @@ -1354,6 +1355,9 @@ def __init__( default_headers: Mapping[str, str] | None = None, default_query: Mapping[str, object] | None = None, http_client: httpx.Client | None = None, + api_pool_shards: int = DEFAULT_API_POOL_SHARDS, + background_pool_shards: int = DEFAULT_BACKGROUND_POOL_SHARDS, + transfer_pool_shards: int = DEFAULT_TRANSFER_POOL_SHARDS, ) -> None: """Configure the synchronous SDK wrapper. @@ -1371,6 +1375,12 @@ def __init__( :type default_query: Mapping[str, object] | None, optional :param http_client: Custom ``httpx.Client`` instance to reuse, defaults to None :type http_client: httpx.Client | None, optional + :param api_pool_shards: H2 shards for short RPCs (round-robin), defaults to 8 + :type api_pool_shards: int, optional + :param background_pool_shards: H2 shards for long-polls (round-robin), defaults to 16 + :type background_pool_shards: int, optional + :param transfer_pool_shards: H2 shards for upload/download (round-robin), defaults to 2 + :type transfer_pool_shards: int, optional """ self.api = Runloop( bearer_token=bearer_token, @@ -1380,6 +1390,9 @@ def __init__( default_headers=default_headers, default_query=default_query, http_client=http_client, + api_pool_shards=api_pool_shards, + background_pool_shards=background_pool_shards, + transfer_pool_shards=transfer_pool_shards, ) self.agent = AgentOps(self.api) diff --git a/tests/test_h2_bulkhead_integration.py b/tests/test_h2_bulkhead_integration.py new file mode 100644 index 000000000..581241a31 --- /dev/null +++ b/tests/test_h2_bulkhead_integration.py @@ -0,0 +1,418 @@ +"""TLS+ALPN HTTP/2 integration tests for workload bulkhead pools. + +Proves wire behavior (not just Python object identity): +- Blocking waits and creates land on different server connections. +- Create completes promptly even when background streams are saturated. +- Upload body stall does not delay create on the API pool. +""" + +from __future__ import annotations + +import ssl +import json +import time +import asyncio +import tempfile +import threading +import subprocess +from typing import Any, Iterator, cast +from pathlib import Path + +import httpx +import pytest + +h2 = pytest.importorskip("h2") +import h2.config # noqa: E402 +import h2.events # noqa: E402 +import h2.settings # noqa: E402 +import h2.connection # noqa: E402 +import h2.exceptions # noqa: E402 + +import runloop_api_client._base_client as _base_mod +from runloop_api_client import AsyncRunloop +from runloop_api_client._base_client import make_request_options + +pytestmark = pytest.mark.timeout(30) + + +def _clear_pool_state() -> None: + old: list[_base_mod._SharedTransport] = [] + old.extend(_base_mod._shared_sync_api_transports.take_all()) + old.extend(_base_mod._shared_sync_background_transports.take_all()) + old.extend(_base_mod._shared_sync_transfer_transports.take_all()) + _base_mod._shared_async_api_transports.clear() + _base_mod._shared_async_background_transports.clear() + _base_mod._shared_async_transfer_transports.clear() + for transport in old: + try: + transport._transport.close() + except Exception: + pass + + +@pytest.fixture(autouse=True) +def _reset_pools() -> Iterator[None]: # pyright: ignore[reportUnusedFunction] + _clear_pool_state() + yield + _clear_pool_state() + + +def _ensure_certs(dir_path: Path) -> tuple[Path, Path]: + cert = dir_path / "cert.pem" + key = dir_path / "key.pem" + if not cert.exists(): + subprocess.run( + [ + "openssl", + "req", + "-x509", + "-newkey", + "rsa:2048", + "-keyout", + str(key), + "-out", + str(cert), + "-days", + "1", + "-nodes", + "-subj", + "/CN=localhost", + ], + check=True, + capture_output=True, + ) + return cert, key + + +class _H2BulkheadServer: + """Minimal TLS+ALPN h2 server that records per-connection request timing.""" + + def __init__(self, *, max_concurrent_streams: int = 2, upload_stall_s: float = 0.0) -> None: + self.max_concurrent_streams = max_concurrent_streams + self.upload_stall_s = upload_stall_s + self.observations: list[dict[str, Any]] = [] + self.active_waits = 0 + self.active_uploads = 0 + self.release_upload_credit = asyncio.Event() + self._conn_seq = 0 + self._conn_lock = threading.Lock() + self._server: asyncio.Server | None = None + self._cert_dir = Path(tempfile.mkdtemp(prefix="h2-bulkhead-")) + self.cert, self.key = _ensure_certs(self._cert_dir) + self.host = "127.0.0.1" + self.port = 0 + + def _next_conn_id(self) -> int: + with self._conn_lock: + self._conn_seq += 1 + return self._conn_seq + + async def start(self) -> str: + ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER) + ctx.load_cert_chain(str(self.cert), str(self.key)) + ctx.set_alpn_protocols(["h2"]) + self._server = await asyncio.start_server(self._handle, self.host, 0, ssl=ctx) + sockets = self._server.sockets + assert sockets + self.port = int(sockets[0].getsockname()[1]) + return f"https://{self.host}:{self.port}" + + async def stop(self) -> None: + if self._server is not None: + self._server.close() + await self._server.wait_closed() + self._server = None + + async def _handle(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: + conn_id = self._next_conn_id() + config = h2.config.H2Configuration(client_side=False, header_encoding="utf-8") + conn = h2.connection.H2Connection(config=config) + conn.initiate_connection() + conn.update_settings( + { + h2.settings.SettingCodes.MAX_CONCURRENT_STREAMS: self.max_concurrent_streams, + } + ) + writer.write(conn.data_to_send()) + await writer.drain() + + streams: dict[int, dict[str, Any]] = {} + write_lock = asyncio.Lock() + + async def flush_withheld_credit() -> None: + """When the test releases the stall, return WINDOW_UPDATE so the upload can finish.""" + await self.release_upload_credit.wait() + for _ in range(200): + async with write_lock: + for sid, st in list(streams.items()): + owed = int(st.get("flow_controlled") or 0) + if owed: + conn.acknowledge_received_data(owed, sid) + st["flow_controlled"] = 0 + to_send = conn.data_to_send() + if to_send: + writer.write(to_send) + await writer.drain() + if self.active_uploads == 0 and not any(int(st.get("flow_controlled") or 0) for st in streams.values()): + return + await asyncio.sleep(0.05) + + credit_task = asyncio.create_task(flush_withheld_credit()) + + async def respond(stream_id: int, st: dict[str, Any]) -> None: + path = st["path"] + method = st["method"] + started = st["started"] + hold_s = 0.0 + if path.endswith("/wait_for_status"): + hold_s = float(st.get("hold_s") or 2.0) + if hold_s > 0: + await asyncio.sleep(hold_s) + + if st.get("flow_controlled"): + async with write_lock: + conn.acknowledge_received_data(st["flow_controlled"], stream_id) + writer.write(conn.data_to_send()) + await writer.drain() + st["flow_controlled"] = 0 + + finished = time.perf_counter() + self.observations.append( + { + "conn_id": conn_id, + "path": path, + "method": method, + "started": started, + "finished": finished, + "duration_s": finished - started, + "body_bytes": len(st.get("body") or b""), + } + ) + if path.endswith("/wait_for_status"): + self.active_waits = max(0, self.active_waits - 1) + if path.endswith("/upload_file"): + self.active_uploads = max(0, self.active_uploads - 1) + + if path == "/v1/devboxes" and method == "POST": + payload = {"id": "dbx_test", "status": "running"} + else: + payload = {"ok": True} + raw = json.dumps(payload).encode() + headers = [ + (":status", "200"), + ("content-type", "application/json"), + ("content-length", str(len(raw))), + ] + async with write_lock: + conn.send_headers(stream_id, headers) + conn.send_data(stream_id, raw, end_stream=True) + writer.write(conn.data_to_send()) + await writer.drain() + streams.pop(stream_id, None) + + try: + while True: + data = await reader.read(65535) + if not data: + break + try: + events = conn.receive_data(data) + except h2.exceptions.ProtocolError: + break + for raw_event in events: + # h2's Event hierarchy is poorly typed under pyright strict. + event = cast(Any, raw_event) + if isinstance(raw_event, h2.events.RequestReceived): + st = streams.setdefault( + event.stream_id, + { + "headers": [], + "body": bytearray(), + "method": "GET", + "path": "/", + "started": time.perf_counter(), + "flow_controlled": 0, + "hold_s": 2.0, + }, + ) + st["headers"] = list(event.headers) + for k, v in event.headers: + if k == ":method": + st["method"] = v + elif k == ":path": + st["path"] = v.split("?", 1)[0] + elif k.lower() == "x-hold-seconds": + try: + st["hold_s"] = float(v) + except ValueError: + pass + if str(st["path"]).endswith("/wait_for_status"): + self.active_waits += 1 + if str(st["path"]).endswith("/upload_file"): + self.active_uploads += 1 + elif isinstance(raw_event, h2.events.DataReceived): + st = streams.setdefault(event.stream_id, {"body": bytearray(), "flow_controlled": 0}) + st["body"].extend(event.data) + stall_upload = self.upload_stall_s > 0 and str(st.get("path", "")).endswith("/upload_file") + if stall_upload and not self.release_upload_credit.is_set(): + st["flow_controlled"] = st.get("flow_controlled", 0) + event.flow_controlled_length + else: + conn.acknowledge_received_data(event.flow_controlled_length, event.stream_id) + elif isinstance(raw_event, h2.events.StreamEnded): + st = streams.get(event.stream_id) + if st is not None: + asyncio.create_task(respond(event.stream_id, st)) + elif isinstance(raw_event, h2.events.StreamReset): + streams.pop(event.stream_id, None) + async with write_lock: + to_send = conn.data_to_send() + if to_send: + writer.write(to_send) + await writer.drain() + except (ConnectionResetError, BrokenPipeError, asyncio.CancelledError): + pass + finally: + credit_task.cancel() + try: + writer.close() + await writer.wait_closed() + except Exception: + pass + + +async def _wait_until(predicate: Any, *, timeout_s: float = 5.0) -> None: + deadline = time.perf_counter() + timeout_s + while time.perf_counter() < deadline: + if predicate(): + return + await asyncio.sleep(0.02) + raise AssertionError("condition not met before timeout") + + +@pytest.mark.asyncio +async def test_waits_and_create_use_different_h2_connections(monkeypatch: pytest.MonkeyPatch) -> None: + # Force trust of the ephemeral self-signed cert used by the local h2 server. + orig_transport_init = httpx.AsyncHTTPTransport.__init__ + + def _transport_init(self: Any, *args: Any, **kwargs: Any) -> None: + kwargs["verify"] = False + orig_transport_init(self, *args, **kwargs) + + monkeypatch.setattr(httpx.AsyncHTTPTransport, "__init__", _transport_init) + + server = _H2BulkheadServer(max_concurrent_streams=2) + base = await server.start() + try: + client = AsyncRunloop( + bearer_token="test-token", + base_url=base, + shared_http_pool=False, + background_pool_shards=1, + transfer_pool_shards=1, + max_retries=0, + ) + try: + wait_headers = {"X-Hold-Seconds": "2.5"} + + async def wait_call() -> object: + return await client.post( + "/v1/devboxes/dbx_shared/wait_for_status", + body={"statuses": ["running"], "timeout_seconds": 5}, + options=make_request_options(extra_headers=wait_headers), + cast_to=object, + ) + + wait_tasks = [asyncio.create_task(wait_call()) for _ in range(2)] + await _wait_until(lambda: server.active_waits >= 2, timeout_s=5.0) + + t0 = time.perf_counter() + create_result = await client.post( + "/v1/devboxes", + body={"name": "bulkhead-create"}, + cast_to=object, + ) + create_s = time.perf_counter() - t0 + + assert isinstance(create_result, dict) + assert create_result["id"] == "dbx_test" + assert create_s < 0.75, f"create blocked by waits: {create_s:.3f}s" + + await asyncio.gather(*wait_tasks) + + wait_obs = [o for o in server.observations if o["path"].endswith("/wait_for_status")] + create_obs = [o for o in server.observations if o["path"] == "/v1/devboxes"] + assert len(wait_obs) == 2 + assert len(create_obs) == 1 + wait_conns = {o["conn_id"] for o in wait_obs} + create_conn = create_obs[0]["conn_id"] + assert create_conn not in wait_conns, f"create conn {create_conn} overlapped wait conns {wait_conns}" + finally: + await client.close() + finally: + await server.stop() + + +@pytest.mark.asyncio +async def test_upload_stall_does_not_block_create(monkeypatch: pytest.MonkeyPatch) -> None: + orig_transport_init = httpx.AsyncHTTPTransport.__init__ + + def _transport_init(self: Any, *args: Any, **kwargs: Any) -> None: + kwargs["verify"] = False + orig_transport_init(self, *args, **kwargs) + + monkeypatch.setattr(httpx.AsyncHTTPTransport, "__init__", _transport_init) + + server = _H2BulkheadServer(max_concurrent_streams=100, upload_stall_s=2.5) + base = await server.start() + try: + client = AsyncRunloop( + bearer_token="test-token", + base_url=base, + shared_http_pool=False, + background_pool_shards=1, + transfer_pool_shards=1, + max_retries=0, + ) + try: + payload = b"x" * (256 * 1024) + + async def upload() -> object: + return await client.post( + "/v1/devboxes/dbx_up/upload_file", + content=payload, + cast_to=object, + options=make_request_options( + extra_headers={"content-type": "application/octet-stream"}, + ), + ) + + upload_task = asyncio.create_task(upload()) + await _wait_until(lambda: server.active_uploads >= 1, timeout_s=5.0) + + t0 = time.perf_counter() + create_result = await client.post( + "/v1/devboxes", + body={"name": "during-upload"}, + cast_to=object, + ) + create_s = time.perf_counter() - t0 + + assert isinstance(create_result, dict) + assert create_result["id"] == "dbx_test" + assert create_s < 0.75, f"create blocked by upload stall: {create_s:.3f}s" + + # Unblock the stalled upload body so the test can finish cleanly. + server.release_upload_credit.set() + # Nudge the connection by allowing the client to finish sending. + await asyncio.wait_for(upload_task, timeout=10.0) + + upload_obs = [o for o in server.observations if o["path"].endswith("/upload_file")] + create_obs = [o for o in server.observations if o["path"] == "/v1/devboxes"] + assert len(upload_obs) == 1 + assert len(create_obs) == 1 + assert create_obs[0]["conn_id"] != upload_obs[0]["conn_id"] + assert create_obs[0]["finished"] < upload_obs[0]["finished"] + finally: + await client.close() + finally: + await server.stop() diff --git a/tests/test_shared_pool.py b/tests/test_shared_pool.py index 4220f8ba9..0300d3625 100644 --- a/tests/test_shared_pool.py +++ b/tests/test_shared_pool.py @@ -28,13 +28,16 @@ def _reset_shared_pool() -> Iterator[None]: # pyright: ignore[reportUnusedFunct def _clear_pool_state() -> None: - with _base_mod._pool_lock: - old_sync = _base_mod._shared_sync_transport - _base_mod._shared_sync_transport = None - _base_mod._shared_async_transports.clear() - if old_sync is not None: + old: list[_base_mod._SharedTransport] = [] + old.extend(_base_mod._shared_sync_api_transports.take_all()) + old.extend(_base_mod._shared_sync_background_transports.take_all()) + old.extend(_base_mod._shared_sync_transfer_transports.take_all()) + _base_mod._shared_async_api_transports.clear() + _base_mod._shared_async_background_transports.clear() + _base_mod._shared_async_transfer_transports.clear() + for transport in old: try: - old_sync._transport.close() + transport._transport.close() except Exception: pass diff --git a/tests/test_transfer_client.py b/tests/test_transfer_client.py new file mode 100644 index 000000000..b7fb81c9c --- /dev/null +++ b/tests/test_transfer_client.py @@ -0,0 +1,243 @@ +"""Tests for sharded H2 background + transfer bulkhead pools.""" + +from __future__ import annotations + +import os +import re +import threading +from typing import Any, Iterator +from pathlib import Path + +import httpx +import pytest + +import runloop_api_client._base_client as _base_mod +from runloop_api_client import Runloop, AsyncRunloop +from runloop_api_client._base_client import ( + _is_transfer_path, + _is_background_path, +) + +base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010") +bearer_token = "My Bearer Token" + +# Invariant: every long-lived / bulk-body RPC in the resources tree must match +# these suffixes or it silently shares the API H2 pool. Update both the +# classification helpers and this list when adding endpoints. +_KNOWN_BACKGROUND_PATHS = ( + "/v1/devboxes/dbx_1/wait_for_status", + "/v1/devboxes/dbx_1/executions/ex_1/wait_for_status", +) +_KNOWN_TRANSFER_PATHS = ( + "/v1/devboxes/dbx_1/upload_file", + "/v1/devboxes/dbx_1/download_file", +) + + +@pytest.fixture(autouse=True) +def _reset_shared_pool() -> Iterator[None]: # pyright: ignore[reportUnusedFunction] + _clear_pool_state() + yield + _clear_pool_state() + + +def _clear_pool_state() -> None: + old: list[_base_mod._SharedTransport] = [] + old.extend(_base_mod._shared_sync_api_transports.take_all()) + old.extend(_base_mod._shared_sync_background_transports.take_all()) + old.extend(_base_mod._shared_sync_transfer_transports.take_all()) + _base_mod._shared_async_api_transports.clear() + _base_mod._shared_async_background_transports.clear() + _base_mod._shared_async_transfer_transports.clear() + for transport in old: + try: + transport._transport.close() + except Exception: + pass + + +def _make_client(**kwargs: Any) -> Runloop: + kwargs.setdefault("base_url", base_url) + kwargs.setdefault("bearer_token", bearer_token) + return Runloop(**kwargs) + + +def test_path_classification() -> None: + assert _is_background_path("/v1/devboxes/dbx_1/wait_for_status") is True + assert _is_background_path("/v1/devboxes/dbx_1/executions/ex_1/wait_for_status") is True + assert _is_transfer_path("/v1/devboxes/dbx_1/upload_file") is True + assert _is_transfer_path("/v1/devboxes/dbx_1/download_file") is True + assert _is_background_path("/v1/devboxes") is False + assert _is_transfer_path("/v1/objects/obj_1/download") is False + + +def test_api_background_transfer_use_distinct_transports() -> None: + client = _make_client(shared_http_pool=True, api_pool_shards=2, background_pool_shards=2, transfer_pool_shards=2) + try: + api_req = httpx.Request("POST", f"{base_url}/v1/devboxes") + wait_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_a/wait_for_status") + upload_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_a/upload_file") + + api = client._get_client_for_path(api_req.url.path) + wait = client._get_client_for_path(wait_req.url.path) + upload = client._get_client_for_path(upload_req.url.path) + + assert wait is not api + assert upload is not api + assert wait is not upload + assert wait._transport is not api._transport # type: ignore[attr-defined] + assert upload._transport is not api._transport # type: ignore[attr-defined] + assert wait._transport is not upload._transport # type: ignore[attr-defined] + finally: + client.close() + + +def test_round_robin_spreads_requests_across_shards() -> None: + client = _make_client(background_pool_shards=2, transfer_pool_shards=2) + assert client._background_pool is not None and client._transfer_pool is not None + try: + wait_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status") + w0 = client._get_client_for_path(wait_req.url.path) + w1 = client._get_client_for_path(wait_req.url.path) + w2 = client._get_client_for_path(wait_req.url.path) + assert w0 is not w1 + assert w0 is w2 + assert set(client._background_pool._clients) == {0, 1} + + upload_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/upload_file") + t0 = client._get_client_for_path(upload_req.url.path) + t1 = client._get_client_for_path(upload_req.url.path) + assert t0 is not t1 + assert set(client._transfer_pool._clients) == {0, 1} + assert _base_mod._shared_sync_background_transports.shard_ids() == {0, 1} + assert _base_mod._shared_sync_transfer_transports.shard_ids() == {0, 1} + finally: + client.close() + + +def test_many_clients_first_requests_use_both_shards() -> None: + clients = [_make_client(background_pool_shards=2) for _ in range(40)] + try: + path = "/v1/devboxes/dbx_1/wait_for_status" + seen_transports = {id(c._get_client_for_path(path)._transport) for c in clients} # type: ignore[attr-defined] + assert len(seen_transports) == 2 + assert _base_mod._shared_sync_background_transports.shard_ids() == {0, 1} + finally: + for c in clients: + c.close() + + +def test_custom_http_client_skips_isolation() -> None: + custom = httpx.Client() + client = _make_client(http_client=custom) + try: + assert client._isolate_workload_pools is False + assert client._get_client_for_path("/v1/devboxes/dbx_1/wait_for_status") is custom + assert client._background_pool is None + finally: + client.close() + custom.close() + + +def test_round_robin_is_per_client_while_transports_are_shared() -> None: + c1 = _make_client(background_pool_shards=2) + c2 = _make_client(background_pool_shards=2) + assert c1._background_pool is not None and c2._background_pool is not None + try: + path = "/v1/devboxes/dbx_shared/wait_for_status" + start1 = c1._background_pool._next + start2 = c2._background_pool._next + t1 = c1._get_client_for_path(path) + t2 = c2._get_client_for_path(path) + assert t1 is not t2 + if start1 % 2 == start2 % 2: + assert t1._transport is t2._transport # type: ignore[attr-defined] + else: + assert t1._transport is not t2._transport # type: ignore[attr-defined] + assert c1._background_pool._next == start1 + 1 + assert c2._background_pool._next == start2 + 1 + + t1b = c1._get_client_for_path(path) + assert t1b is not t1 + assert c1._background_pool._next == start1 + 2 + assert c2._background_pool._next == start2 + 1 + finally: + c1.close() + c2.close() + + +@pytest.mark.asyncio +async def test_async_bulkheads() -> None: + client = AsyncRunloop( + base_url=base_url, + bearer_token=bearer_token, + shared_http_pool=True, + background_pool_shards=2, + transfer_pool_shards=2, + ) + try: + wait = client._get_client_for_path("/v1/devboxes/dbx_1/executions/ex_1/wait_for_status") + upload = client._get_client_for_path("/v1/devboxes/dbx_1/upload_file") + assert wait is not client._client + assert upload is not client._client + assert wait is not upload + finally: + await client.close() + + +def test_concurrent_background_client_init_is_singleton() -> None: + client = _make_client(shared_http_pool=False, background_pool_shards=1) + assert client._background_pool is not None + barrier = threading.Barrier(16) + results: list[httpx.Client] = [] + errors: list[BaseException] = [] + + def worker() -> None: + try: + barrier.wait() + results.append(client._background_pool.ensure(0)) # type: ignore[union-attr] + except BaseException as exc: # noqa: BLE001 + errors.append(exc) + + threads = [threading.Thread(target=worker) for _ in range(16)] + for t in threads: + t.start() + for t in threads: + t.join() + + assert not errors + assert len(results) == 16 + assert len({id(c) for c in results}) == 1 + client.close() + + +def test_known_long_lived_paths_are_classified() -> None: + for path in _KNOWN_BACKGROUND_PATHS: + assert _is_background_path(path), path + assert not _is_transfer_path(path), path + for path in _KNOWN_TRANSFER_PATHS: + assert _is_transfer_path(path), path + assert not _is_background_path(path), path + + +def test_resources_tree_long_lived_ops_match_classifier() -> None: + """Fail if a new wait_for_status / upload_file / download_file path is added + without being covered by the bulkhead suffixes. + """ + root = Path(__file__).resolve().parents[1] / "src" / "runloop_api_client" / "resources" + text = "\n".join(p.read_text() for p in root.rglob("*.py")) + # Paths appear as f-strings or path_template("/v1/.../wait_for_status", ...) + found_wait = set(re.findall(r'(/v1/[^"\s]*wait_for_status)', text)) + found_upload = set(re.findall(r'(/v1/[^"\s]*upload_file)', text)) + found_download = set(re.findall(r'(/v1/[^"\s]*download_file)', text)) + + assert found_wait, "expected wait_for_status paths in resources/" + assert found_upload, "expected upload_file paths in resources/" + assert found_download, "expected download_file paths in resources/" + + for path in found_wait: + concrete = path.replace("{id}", "dbx_1").replace("{devbox_id}", "dbx_1").replace("{execution_id}", "ex_1") + assert _is_background_path(concrete), path + for path in found_upload | found_download: + concrete = path.replace("{id}", "dbx_1") + assert _is_transfer_path(concrete), path