fix: round-robin H2 bulkhead shards per SDK client

Replace CRC32 resource affinity with per-client counters so concurrent waits/uploads spread across shards instead of pinning a busy resource to one connection.

Co-authored-by: Cursor <cursoragent@cursor.com>

Author: Reflex

src/runloop_api_client/_base_client.py

@@ -4,7 +4,6 @@
 import json
 import time
 import uuid
-import zlib
 import email
 import asyncio
 import inspect
@@ -179,7 +178,8 @@ async def aclose(self) -> None:
 
 # Sharded H2 bulkheads: long-polls and file transfers stay off the API control-plane
 # connection. Each shard index maps to its own shared transport (≈ one H2 connection).
-# Removable once httpcore respects stream capacity when opening connections.
+# Per-client round-robin spreads concurrent requests across shards. Removable once
+# httpcore respects stream capacity when opening connections.
 _shared_sync_background_transports: dict[int, _SharedTransport] = {}
 _shared_sync_transfer_transports: dict[int, _SharedTransport] = {}
 _shared_async_background_transports: weakref.WeakKeyDictionary[
@@ -201,30 +201,6 @@ def _is_transfer_path(path: str) -> bool:
     return path.endswith(_TRANSFER_PATH_SUFFIXES)
 
 
-def _pool_affinity_key(path: str) -> str:
-    """Pick a stable resource id from the URL for shard routing."""
-    parts = [p for p in path.split("/") if p]
-    try:
-        if "executions" in parts:
-            idx = parts.index("executions")
-            if idx + 1 < len(parts):
-                return parts[idx + 1]
-        if "devboxes" in parts:
-            idx = parts.index("devboxes")
-            if idx + 1 < len(parts):
-                return parts[idx + 1]
-    except ValueError:
-        pass
-    return path
-
-
-def _shard_index(key: str, shards: int) -> int:
-    if shards <= 1:
-        return 0
-    # crc32 is stable across processes (unlike PYTHONHASHSEED-randomized hash()).
-    return zlib.crc32(key.encode("utf-8")) % shards
-
-
 def _acquire_shared_sync_transport(bucket: dict[int, _SharedTransport], shard: int) -> _SharedTransport:
     with _pool_lock:
         existing = bucket.get(shard)
@@ -1019,6 +995,8 @@ class SyncAPIClient(BaseClient[httpx.Client, Stream[Any]]):
     _isolate_workload_pools: bool
     _background_pool_shards: int
     _transfer_pool_shards: int
+    _background_next: int
+    _transfer_next: int
     _closed: bool
 
     def __init__(
@@ -1076,6 +1054,8 @@ def __init__(
         self._bulkhead_lock = threading.Lock()
         self._background_pool_shards = background_pool_shards
         self._transfer_pool_shards = transfer_pool_shards
+        self._background_next = 0
+        self._transfer_next = 0
         # Custom http_client owns the full transport stack; don't invent sibling pools.
         self._isolate_workload_pools = http_client is None
 
@@ -1151,15 +1131,27 @@ def _ensure_transfer_client(self, shard: int) -> httpx.Client:
             self._transfer_clients[shard] = client
             return client
 
+    def _next_background_client(self) -> httpx.Client:
+        # Select under the lock; ensure afterward so _ensure_* can take the same lock.
+        with self._bulkhead_lock:
+            shard = self._background_next % self._background_pool_shards
+            self._background_next += 1
+        return self._ensure_background_client(shard)
+
+    def _next_transfer_client(self) -> httpx.Client:
+        with self._bulkhead_lock:
+            shard = self._transfer_next % self._transfer_pool_shards
+            self._transfer_next += 1
+        return self._ensure_transfer_client(shard)
+
     def _send_client_for_request(self, request: httpx.Request) -> httpx.Client:
         if not self._isolate_workload_pools:
             return self._client
         path = request.url.path
-        key = _pool_affinity_key(path)
         if _is_background_path(path):
-            return self._ensure_background_client(_shard_index(key, self._background_pool_shards))
+            return self._next_background_client()
         if _is_transfer_path(path):
-            return self._ensure_transfer_client(_shard_index(key, self._transfer_pool_shards))
+            return self._next_transfer_client()
         return self._client
 
     def is_closed(self) -> bool:
@@ -1736,6 +1728,8 @@ class AsyncAPIClient(BaseClient[httpx.AsyncClient, AsyncStream[Any]]):
     _isolate_workload_pools: bool
     _background_pool_shards: int
     _transfer_pool_shards: int
+    _background_next: int
+    _transfer_next: int
     _closed: bool
 
     def __init__(
@@ -1792,6 +1786,8 @@ def __init__(
         self._transfer_clients = {}
         self._background_pool_shards = background_pool_shards
         self._transfer_pool_shards = transfer_pool_shards
+        self._background_next = 0
+        self._transfer_next = 0
         # Custom http_client owns the full transport stack; don't invent sibling pools.
         self._isolate_workload_pools = http_client is None
 
@@ -1877,15 +1873,25 @@ def _ensure_transfer_client(self, shard: int) -> httpx.AsyncClient:
         self._transfer_clients[shard] = client
         return client
 
+    def _next_background_client(self) -> httpx.AsyncClient:
+        # Single-threaded event loop: counter bump needs no lock when there is no await.
+        shard = self._background_next % self._background_pool_shards
+        self._background_next += 1
+        return self._ensure_background_client(shard)
+
+    def _next_transfer_client(self) -> httpx.AsyncClient:
+        shard = self._transfer_next % self._transfer_pool_shards
+        self._transfer_next += 1
+        return self._ensure_transfer_client(shard)
+
     def _send_client_for_request(self, request: httpx.Request) -> httpx.AsyncClient:
         if not self._isolate_workload_pools:
             return self._client
         path = request.url.path
-        key = _pool_affinity_key(path)
         if _is_background_path(path):
-            return self._ensure_background_client(_shard_index(key, self._background_pool_shards))
+            return self._next_background_client()
         if _is_transfer_path(path):
-            return self._ensure_transfer_client(_shard_index(key, self._transfer_pool_shards))
+            return self._next_transfer_client()
         return self._client
 
     def is_closed(self) -> bool:

src/runloop_api_client/_client.py

@@ -98,7 +98,7 @@ def __init__(
         # Set to False to create a private connection pool (old behavior).
         shared_http_pool: bool = True,
         # Separate H2 connection pools for long-polls (/wait_for_status) and file
-        # transfers. Each shard ≈ one H2 connection, selected by hash(resource_id).
+        # transfers. Each shard ≈ one H2 connection; requests round-robin per client.
         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.
@@ -406,7 +406,7 @@ def __init__(
         # Set to False to create a private connection pool (old behavior).
         shared_http_pool: bool = True,
         # Separate H2 connection pools for long-polls (/wait_for_status) and file
-        # transfers. Each shard ≈ one H2 connection, selected by hash(resource_id).
+        # transfers. Each shard ≈ one H2 connection; requests round-robin per client.
         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.

src/runloop_api_client/sdk/async_.py

@@ -1349,9 +1349,9 @@ 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 background_pool_shards: H2 shards for long-polls, defaults to 2
+        :param background_pool_shards: H2 shards for long-polls (round-robin), defaults to 2
         :type background_pool_shards: int, optional
-        :param transfer_pool_shards: H2 shards for upload/download, defaults to 2
+        :param transfer_pool_shards: H2 shards for upload/download (round-robin), defaults to 2
         :type transfer_pool_shards: int, optional
         """
         self.api = AsyncRunloop(

src/runloop_api_client/sdk/sync.py

@@ -1374,9 +1374,9 @@ 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 background_pool_shards: H2 shards for long-polls, defaults to 2
+        :param background_pool_shards: H2 shards for long-polls (round-robin), defaults to 2
         :type background_pool_shards: int, optional
-        :param transfer_pool_shards: H2 shards for upload/download, defaults to 2
+        :param transfer_pool_shards: H2 shards for upload/download (round-robin), defaults to 2
         :type transfer_pool_shards: int, optional
         """
         self.api = Runloop(

tests/test_transfer_client.py

@@ -14,9 +14,7 @@
 import runloop_api_client._base_client as _base_mod
 from runloop_api_client import Runloop, AsyncRunloop
 from runloop_api_client._base_client import (
-    _shard_index,
     _is_transfer_path,
-    _pool_affinity_key,
     _is_background_path,
 )
 
@@ -77,15 +75,6 @@ def test_path_classification() -> None:
     assert _is_transfer_path("/v1/objects/obj_1/download") is False
 
 
-def test_affinity_key_and_shard() -> None:
-    assert _pool_affinity_key("/v1/devboxes/dbx_1/upload_file") == "dbx_1"
-    assert _pool_affinity_key("/v1/devboxes/dbx_1/wait_for_status") == "dbx_1"
-    assert _pool_affinity_key("/v1/devboxes/dbx_1/executions/ex_9/wait_for_status") == "ex_9"
-    assert _shard_index("dbx_1", 1) == 0
-    assert _shard_index("dbx_1", 2) in (0, 1)
-    assert _shard_index("dbx_1", 2) == _shard_index("dbx_1", 2)
-
-
 def test_api_background_transfer_use_distinct_transports() -> None:
     client = _make_client(shared_http_pool=True, background_pool_shards=2, transfer_pool_shards=2)
     try:
@@ -108,26 +97,24 @@ def test_api_background_transfer_use_distinct_transports() -> None:
         client.close()
 
 
-def test_same_resource_affinity_uses_same_shard() -> None:
+def test_round_robin_spreads_requests_across_shards() -> None:
     client = _make_client(background_pool_shards=2, transfer_pool_shards=2)
     try:
-        w1 = client._send_client_for_request(httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status"))
-        w2 = client._send_client_for_request(httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status"))
-        assert w1 is w2
-    finally:
-        client.close()
-
-
-def test_different_resources_can_land_on_different_shards() -> None:
-    client = _make_client(background_pool_shards=2)
-    try:
-        seen: set[int] = set()
-        for i in range(40):
-            req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_{i}/wait_for_status")
-            c = client._send_client_for_request(req)
-            seen.add(id(c._transport))  # type: ignore[attr-defined]
-        assert len(seen) == 2
+        wait_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/wait_for_status")
+        w0 = client._send_client_for_request(wait_req)
+        w1 = client._send_client_for_request(wait_req)
+        w2 = client._send_client_for_request(wait_req)
+        assert w0 is not w1
+        assert w0 is w2
+        assert set(client._background_clients) == {0, 1}
+
+        upload_req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_same/upload_file")
+        t0 = client._send_client_for_request(upload_req)
+        t1 = client._send_client_for_request(upload_req)
+        assert t0 is not t1
+        assert set(client._transfer_clients) == {0, 1}
         assert len(_base_mod._shared_sync_background_transports) == 2
+        assert len(_base_mod._shared_sync_transfer_transports) == 2
     finally:
         client.close()
 
@@ -145,15 +132,24 @@ def test_custom_http_client_skips_isolation() -> None:
         custom.close()
 
 
-def test_shards_shared_across_sdk_clients() -> None:
+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)
     try:
         req = httpx.Request("POST", f"{base_url}/v1/devboxes/dbx_shared/wait_for_status")
+        # Each SDK client has its own counter; both start at shard 0.
         t1 = c1._send_client_for_request(req)
         t2 = c2._send_client_for_request(req)
         assert t1 is not t2
         assert t1._transport is t2._transport  # type: ignore[attr-defined]
+        assert c1._background_next == 1
+        assert c2._background_next == 1
+
+        # Advancing one client does not affect the other.
+        t1b = c1._send_client_for_request(req)
+        assert t1b is not t1
+        assert c1._background_next == 2
+        assert c2._background_next == 1
     finally:
         c1.close()
         c2.close()