Prevent replaying streams after auth read errors

Author: tnsardesai

src/kernel/_base_client.py

@@ -772,6 +772,9 @@ def _calculate_retry_timeout(
         timeout = sleep_seconds * jitter
         return timeout if timeout >= 0 else 0
 
+    def _should_retry_on_connection_error(self, _request: httpx.Request) -> bool:
+        return True
+
     def _should_retry(self, response: httpx.Response) -> bool:
         # Note: this is not a standard header
         should_retry_header = response.headers.get("x-should-retry")
@@ -1026,7 +1029,7 @@ def request(
             except Exception as err:
                 log.debug("Encountered Exception", 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,
@@ -1610,7 +1613,7 @@ async def request(
             except Exception as err:
                 log.debug("Encountered Exception", 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

@@ -45,6 +45,7 @@
     should_retry_stale_direct_vm_auth,
     install_stale_direct_vm_auth_eviction,
     maybe_evict_browser_route_from_response,
+    should_retry_direct_vm_connection_error,
     install_async_stale_direct_vm_auth_eviction,
     maybe_populate_browser_route_cache_from_response,
 )
@@ -375,6 +376,10 @@ def _prepare_options(self, options: Any) -> Any:
     def _prepare_request(self, request: httpx.Request) -> None:
         strip_direct_vm_auth(request, cache=self.browser_route_cache)
 
+    @override
+    def _should_retry_on_connection_error(self, request: httpx.Request) -> bool:
+        return should_retry_direct_vm_connection_error(request)
+
     @override
     def _should_retry(self, response: httpx.Response) -> bool:
         if is_stale_direct_vm_auth_response(response):
@@ -767,6 +772,10 @@ async def _prepare_options(self, options: Any) -> Any:
     async def _prepare_request(self, request: httpx.Request) -> None:
         strip_direct_vm_auth(request, cache=self.browser_route_cache)
 
+    @override
+    def _should_retry_on_connection_error(self, request: httpx.Request) -> bool:
+        return should_retry_direct_vm_connection_error(request)
+
     @override
     def _should_retry(self, response: httpx.Response) -> bool:
         if is_stale_direct_vm_auth_response(response):

src/kernel/lib/browser_routing/routing.py

@@ -33,6 +33,7 @@ class BrowserRoutingConfig:
 
 
 _EVICTION_HOOK_CACHE_ATTR = "_kernel_browser_route_cache"
+_STALE_DIRECT_VM_AUTH_REQUEST_EXTENSION = "kernel_stale_direct_vm_auth"
 
 
 _BROWSER_ROUTE_CACHEABLE_PATH = re.compile(r"^/(?:v\d+/)?browsers(?:/[^/]+)?/?$")
@@ -220,6 +221,7 @@ def install_stale_direct_vm_auth_eviction(client: httpx.Client, *, cache: Browse
 
     def evict(response: httpx.Response) -> None:
         if is_stale_direct_vm_auth_response(response):
+            response.request.extensions[_STALE_DIRECT_VM_AUTH_REQUEST_EXTENSION] = True
             maybe_evict_browser_route_from_response(response, cache=cache)
 
     setattr(evict, _EVICTION_HOOK_CACHE_ATTR, cache)
@@ -234,6 +236,7 @@ def install_async_stale_direct_vm_auth_eviction(client: httpx.AsyncClient, *, ca
 
     async def evict(response: httpx.Response) -> None:
         if is_stale_direct_vm_auth_response(response):
+            response.request.extensions[_STALE_DIRECT_VM_AUTH_REQUEST_EXTENSION] = True
             maybe_evict_browser_route_from_response(response, cache=cache)
 
     setattr(evict, _EVICTION_HOOK_CACHE_ATTR, cache)
@@ -246,6 +249,19 @@ def _has_eviction_hook(hooks: list[Any], cache: BrowserRouteCache) -> bool:
     return any(getattr(hook, _EVICTION_HOOK_CACHE_ATTR, None) is cache for hook in hooks)
 
 
+def should_retry_direct_vm_connection_error(request: httpx.Request) -> bool:
+    """Prevent generic retries from replaying a stream after stale VM auth.
+
+    httpx may raise while eagerly reading the error response body, before the
+    SDK can pass the response to `_should_retry`. The response hook records the
+    known stale-auth status on the request so this connection-error path can
+    apply the same body replay check.
+    """
+    if not request.extensions.get(_STALE_DIRECT_VM_AUTH_REQUEST_EXTENSION):
+        return True
+    return direct_vm_request_body_is_replayable(request)
+
+
 def should_retry_stale_direct_vm_auth(response: httpx.Response) -> bool:
     """Whether a stale direct-to-VM auth failure can be retried on the control plane.
 

tests/test_browser_routing.py

@@ -50,6 +50,10 @@ def _skip_retry_sleep(_self: object, **_kwargs: object) -> None:
     return None
 
 
+async def _skip_async_retry_sleep(_self: object, **_kwargs: object) -> None:
+    return None
+
+
 class _UnseekableFile(io.RawIOBase):
     """A file-like upload body that cannot be rewound, e.g. a pipe.
 
@@ -1347,6 +1351,98 @@ async def handle_request(request: httpx.Request) -> httpx.Response:
     assert requests[1].headers.get("Authorization") == f"Bearer {api_key}"
 
 
+def test_stale_direct_vm_auth_body_read_failure_does_not_retry_unreplayable_write(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False)
+    monkeypatch.setattr("kernel._base_client.SyncAPIClient._sleep_for_retry", _skip_retry_sleep)
+    requests: list[tuple[httpx.URL, bytes]] = []
+
+    def handle_request(request: httpx.Request) -> httpx.Response:
+        body = b"".join(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(201)
+
+    class Transport(httpx.BaseTransport):
+        @override
+        def handle_request(self, request: httpx.Request) -> httpx.Response:
+            return handle_request(request)
+
+    http_client = httpx.Client(transport=Transport())
+    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):
+            client.browsers.fs.write_file("sess-1", _UnseekableFile(b"payload"), path="/tmp/x")
+
+        assert requests == [
+            (
+                httpx.URL("http://browser-session.test/browser/kernel/fs/write_file?path=%2Ftmp%2Fx&jwt=token-abc"),
+                b"payload",
+            )
+        ]
+        assert client.browser_route_cache.get("sess-1") is None
+
+        client.browsers.fs.write_file("sess-1", b"next", path="/tmp/x")
+
+    assert requests[1] == (httpx.URL(f"{base_url}/browsers/sess-1/fs/write_file?path=%2Ftmp%2Fx"), b"next")
+
+
+@pytest.mark.asyncio
+async def test_async_stale_direct_vm_auth_body_read_failure_does_not_retry_unreplayable_write(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False)
+    monkeypatch.setattr("kernel._base_client.AsyncAPIClient._sleep_for_retry", _skip_async_retry_sleep)
+    requests: list[tuple[httpx.URL, bytes]] = []
+
+    async def handle_request(request: httpx.Request) -> httpx.Response:
+        body = b"".join([chunk async for chunk in 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(201)
+
+    async def payload() -> AsyncIterator[bytes]:
+        yield b"payload"
+
+    class Transport(httpx.AsyncBaseTransport):
+        @override
+        async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+            return await handle_request(request)
+
+    http_client = httpx.AsyncClient(transport=Transport())
+    async with AsyncKernel(
+        base_url=base_url,
+        api_key=api_key,
+        http_client=http_client,
+        _strict_response_validation=True,
+    ) as client:
+        route = browser_route_from_browser(_fake_browser())
+        assert route is not None
+        client.browser_route_cache.set(route)
+        with pytest.raises(APIConnectionError):
+            await client.browsers.fs.write_file("sess-1", cast(Any, payload()), path="/tmp/x")
+
+        assert requests == [
+            (
+                httpx.URL("http://browser-session.test/browser/kernel/fs/write_file?path=%2Ftmp%2Fx&jwt=token-abc"),
+                b"payload",
+            )
+        ]
+        assert client.browser_route_cache.get("sess-1") is None
+
+        await client.browsers.fs.write_file("sess-1", b"next", path="/tmp/x")
+
+    assert requests[1] == (httpx.URL(f"{base_url}/browsers/sess-1/fs/write_file?path=%2Ftmp%2Fx"), b"next")
+
+
 def test_copied_client_registers_one_route_eviction_hook() -> None:
     with Kernel(base_url=base_url, api_key=api_key, _strict_response_validation=True) as client:
         copied = client.copy(api_key="sk-456")