Preserve stream replayability before request hooks

Author: tnsardesai

src/kernel/_base_client.py

@@ -1015,7 +1015,7 @@ def request(
             except httpx.TimeoutException as err:
                 log.debug("Encountered httpx.TimeoutException", exc_info=True)
 
-                if remaining_retries > 0:
+                if remaining_retries > 0 and self._should_retry_on_connection_error(request):
                     self._sleep_for_retry(
                         retries_taken=retries_taken,
                         max_retries=max_retries,
@@ -1599,7 +1599,7 @@ async def request(
             except httpx.TimeoutException as err:
                 log.debug("Encountered httpx.TimeoutException", exc_info=True)
 
-                if remaining_retries > 0:
+                if remaining_retries > 0 and self._should_retry_on_connection_error(request):
                     await self._sleep_for_retry(
                         retries_taken=retries_taken,
                         max_retries=max_retries,

src/kernel/_client.py

@@ -38,7 +38,7 @@
 from .lib.browser_routing.routing import (
     BrowserRouteCache,
     BrowserRoutingConfig,
-    strip_direct_vm_auth,
+    prepare_direct_vm_request,
     rewrite_direct_vm_options,
     browser_routing_config_from_env,
     is_stale_direct_vm_auth_response,
@@ -374,7 +374,7 @@ def _prepare_options(self, options: Any) -> Any:
 
     @override
     def _prepare_request(self, request: httpx.Request) -> None:
-        strip_direct_vm_auth(request, cache=self.browser_route_cache)
+        prepare_direct_vm_request(request, cache=self.browser_route_cache)
 
     @override
     def _should_retry_on_connection_error(self, request: httpx.Request) -> bool:
@@ -770,7 +770,7 @@ async def _prepare_options(self, options: Any) -> Any:
 
     @override
     async def _prepare_request(self, request: httpx.Request) -> None:
-        strip_direct_vm_auth(request, cache=self.browser_route_cache)
+        prepare_direct_vm_request(request, cache=self.browser_route_cache)
 
     @override
     def _should_retry_on_connection_error(self, request: httpx.Request) -> bool:

src/kernel/lib/browser_routing/routing.py

@@ -34,6 +34,7 @@ class BrowserRoutingConfig:
 
 _EVICTION_HOOK_CACHE_ATTR = "_kernel_browser_route_cache"
 _STALE_DIRECT_VM_AUTH_REQUEST_EXTENSION = "kernel_stale_direct_vm_auth"
+_DIRECT_VM_BODY_REPLAYABLE_REQUEST_EXTENSION = "kernel_direct_vm_body_replayable"
 
 
 _BROWSER_ROUTE_CACHEABLE_PATH = re.compile(r"^/(?:v\d+/)?browsers(?:/[^/]+)?/?$")
@@ -276,12 +277,18 @@ def should_retry_stale_direct_vm_auth(response: httpx.Response) -> bool:
 
 
 def direct_vm_request_body_is_replayable(request: httpx.Request) -> bool:
+    replayable = request.extensions.get(_DIRECT_VM_BODY_REPLAYABLE_REQUEST_EXTENSION)
+    if isinstance(replayable, bool):
+        return replayable
+    return _classify_direct_vm_request_body_replayability(request)
+
+
+def _classify_direct_vm_request_body_replayability(request: httpx.Request) -> bool:
     try:
         _ = request.content
     except httpx.RequestNotRead:
         pass
     else:
-        # httpx already buffered the body, so rebuilding it yields the same bytes.
         return True
 
     # httpx encodes multipart bodies as a stream of fields it re-renders per attempt.
@@ -410,10 +417,15 @@ def rewrite_direct_vm_options(
     return rewritten
 
 
-def strip_direct_vm_auth(request: httpx.Request, *, cache: BrowserRouteCache) -> None:
+def prepare_direct_vm_request(request: httpx.Request, *, cache: BrowserRouteCache) -> None:
     raw = str(request.url)
     for route in cache.values():
         if raw.startswith(route.base_url.rstrip("/") + "/"):
+            # Request hooks and custom auth can buffer this request while consuming
+            # the original body that the SDK would use to build a retry.
+            request.extensions[_DIRECT_VM_BODY_REPLAYABLE_REQUEST_EXTENSION] = (
+                _classify_direct_vm_request_body_replayability(request)
+            )
             request.headers.pop("Authorization", None)
             return
 

tests/test_browser_routing.py

@@ -14,6 +14,7 @@
 from kernel import (
     Kernel,
     AsyncKernel,
+    APITimeoutError,
     APIConnectionError,
     AuthenticationError,
     InternalServerError,
@@ -1269,15 +1270,21 @@ def test_indexed_multipart_body_flattens_only_given_values() -> None:
 class _FailingSyncStream(httpx.SyncByteStream):
     """A response body that fails while it is being read."""
 
+    def __init__(self, error_type: type[httpx.TransportError] = httpx.ReadError) -> None:
+        self.error_type = error_type
+
     @override
     def __iter__(self) -> Iterator[bytes]:
-        raise httpx.ReadError("connection reset while reading the error body")
+        raise self.error_type("connection failed while reading the error body")
 
 
 class _FailingAsyncStream(httpx.AsyncByteStream):
+    def __init__(self, error_type: type[httpx.TransportError] = httpx.ReadError) -> None:
+        self.error_type = error_type
+
     @override
     async def __aiter__(self) -> AsyncIterator[bytes]:
-        raise httpx.ReadError("connection reset while reading the error body")
+        raise self.error_type("connection failed while reading the error body")
         yield b""  # pragma: no cover - unreachable, keeps this an async generator
 
 
@@ -1351,8 +1358,14 @@ async def handle_request(request: httpx.Request) -> httpx.Response:
     assert requests[1].headers.get("Authorization") == f"Bearer {api_key}"
 
 
+@pytest.mark.parametrize(
+    ("error_type", "expected_error"),
+    [(httpx.ReadError, APIConnectionError), (httpx.ReadTimeout, APITimeoutError)],
+)
 def test_stale_direct_vm_auth_body_read_failure_does_not_retry_unreplayable_write(
     monkeypatch: pytest.MonkeyPatch,
+    error_type: type[httpx.TransportError],
+    expected_error: type[Exception],
 ) -> None:
     monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False)
     monkeypatch.setattr("kernel._base_client.SyncAPIClient._sleep_for_retry", _skip_retry_sleep)
@@ -1362,23 +1375,26 @@ def handle_request(request: httpx.Request) -> httpx.Response:
         body = b"".join(cast(Iterator[bytes], request.stream))
         requests.append((request.url, body))
         if "browser-session.test" in str(request.url):
-            return httpx.Response(401, stream=_FailingSyncStream(), headers={"content-type": "text/plain"})
+            return httpx.Response(401, stream=_FailingSyncStream(error_type), headers={"content-type": "text/plain"})
         return httpx.Response(201)
 
+    def read_request(request: httpx.Request) -> None:
+        request.read()
+
     class Transport(httpx.BaseTransport):
         @override
         def handle_request(self, request: httpx.Request) -> httpx.Response:
             return handle_request(request)
 
-    http_client = httpx.Client(transport=Transport())
+    http_client = httpx.Client(transport=Transport(), event_hooks={"request": [read_request]})
     with Kernel(
         base_url=base_url,
         api_key=api_key,
         http_client=http_client,
         _strict_response_validation=True,
     ) as client:
         _cache_browser(client)
-        with pytest.raises(APIConnectionError):
+        with pytest.raises(expected_error):
             client.browsers.fs.write_file("sess-1", _UnseekableFile(b"payload"), path="/tmp/x")
 
         assert requests == [
@@ -1395,8 +1411,14 @@ def handle_request(self, request: httpx.Request) -> httpx.Response:
 
 
 @pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("error_type", "expected_error"),
+    [(httpx.ReadError, APIConnectionError), (httpx.ReadTimeout, APITimeoutError)],
+)
 async def test_async_stale_direct_vm_auth_body_read_failure_does_not_retry_unreplayable_write(
     monkeypatch: pytest.MonkeyPatch,
+    error_type: type[httpx.TransportError],
+    expected_error: type[Exception],
 ) -> None:
     monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False)
     monkeypatch.setattr("kernel._base_client.AsyncAPIClient._sleep_for_retry", _skip_async_retry_sleep)
@@ -1406,9 +1428,12 @@ async def handle_request(request: httpx.Request) -> httpx.Response:
         body = b"".join([chunk async for chunk in cast(AsyncIterator[bytes], request.stream)])
         requests.append((request.url, body))
         if "browser-session.test" in str(request.url):
-            return httpx.Response(403, stream=_FailingAsyncStream(), headers={"content-type": "text/plain"})
+            return httpx.Response(403, stream=_FailingAsyncStream(error_type), headers={"content-type": "text/plain"})
         return httpx.Response(201)
 
+    async def read_request(request: httpx.Request) -> None:
+        await request.aread()
+
     async def payload() -> AsyncIterator[bytes]:
         yield b"payload"
 
@@ -1417,7 +1442,7 @@ class Transport(httpx.AsyncBaseTransport):
         async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
             return await handle_request(request)
 
-    http_client = httpx.AsyncClient(transport=Transport())
+    http_client = httpx.AsyncClient(transport=Transport(), event_hooks={"request": [read_request]})
     async with AsyncKernel(
         base_url=base_url,
         api_key=api_key,
@@ -1427,7 +1452,7 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
         route = browser_route_from_browser(_fake_browser())
         assert route is not None
         client.browser_route_cache.set(route)
-        with pytest.raises(APIConnectionError):
+        with pytest.raises(expected_error):
             await client.browsers.fs.write_file("sess-1", cast(Any, payload()), path="/tmp/x")
 
         assert requests == [