Evict a stale direct-to-VM route before the error body is read

httpx reads the body of a non-streamed response inside send(), so a 401 or
403 whose body read fails surfaces as a connection error and never reaches
the status-error path. The route eviction ran after that read, which left a
dead JWT cached and wedged every later call for the session.

Move eviction into an httpx response event hook, which runs once the status
is known and before any body is read, for both the sync and async clients.
`_should_retry` now only decides whether replaying the body on the control
plane is safe.

Author: tnsardesai

src/kernel/_client.py

@@ -43,7 +43,9 @@
     browser_routing_config_from_env,
     is_stale_direct_vm_auth_response,
     should_retry_stale_direct_vm_auth,
+    install_stale_direct_vm_auth_eviction,
     maybe_evict_browser_route_from_response,
+    install_async_stale_direct_vm_auth_eviction,
     maybe_populate_browser_route_cache_from_response,
 )
 
@@ -204,6 +206,7 @@ def __init__(
         )
         self.browser_route_cache = _browser_route_cache or BrowserRouteCache()
         self._browser_routing = browser_routing_config_from_env()
+        install_stale_direct_vm_auth_eviction(self._client, cache=self.browser_route_cache)
 
     @cached_property
     def deployments(self) -> DeploymentsResource:
@@ -364,28 +367,15 @@ 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)
 
-    def _evict_stale_direct_vm_route(self, response: httpx.Response) -> None:
-        maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache)
-
     @override
     def _should_retry(self, response: httpx.Response) -> bool:
         if is_stale_direct_vm_auth_response(response):
-            self._evict_stale_direct_vm_route(response)
-            # The route is evicted either way; only retry when the body can be
-            # rebuilt, otherwise the caller sees the original auth failure and a
-            # later call goes to the control plane.
+            # The route was already evicted by the response hook; retry only when
+            # the body can be rebuilt, otherwise the caller sees the original auth
+            # failure and a later call goes to the control plane.
             return should_retry_stale_direct_vm_auth(response)
         return super()._should_retry(response)
 
-    @override
-    def _make_status_error_from_response(self, response: httpx.Response) -> APIStatusError:
-        # `_should_retry` never runs when the request has no retries left, so this
-        # is the only place a stale direct-to-VM route gets evicted before the
-        # error surfaces to the caller.
-        if is_stale_direct_vm_auth_response(response):
-            self._evict_stale_direct_vm_route(response)
-        return super()._make_status_error_from_response(response)
-
     @override
     def _process_response(
         self,
@@ -602,6 +592,7 @@ def __init__(
         )
         self.browser_route_cache = _browser_route_cache or BrowserRouteCache()
         self._browser_routing = browser_routing_config_from_env()
+        install_async_stale_direct_vm_auth_eviction(self._client, cache=self.browser_route_cache)
 
     @cached_property
     def deployments(self) -> AsyncDeploymentsResource:
@@ -762,28 +753,15 @@ 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)
 
-    def _evict_stale_direct_vm_route(self, response: httpx.Response) -> None:
-        maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache)
-
     @override
     def _should_retry(self, response: httpx.Response) -> bool:
         if is_stale_direct_vm_auth_response(response):
-            self._evict_stale_direct_vm_route(response)
-            # The route is evicted either way; only retry when the body can be
-            # rebuilt, otherwise the caller sees the original auth failure and a
-            # later call goes to the control plane.
+            # The route was already evicted by the response hook; retry only when
+            # the body can be rebuilt, otherwise the caller sees the original auth
+            # failure and a later call goes to the control plane.
             return should_retry_stale_direct_vm_auth(response)
         return super()._should_retry(response)
 
-    @override
-    def _make_status_error_from_response(self, response: httpx.Response) -> APIStatusError:
-        # `_should_retry` never runs when the request has no retries left, so this
-        # is the only place a stale direct-to-VM route gets evicted before the
-        # error surfaces to the caller.
-        if is_stale_direct_vm_auth_response(response):
-            self._evict_stale_direct_vm_route(response)
-        return super()._make_status_error_from_response(response)
-
     @override
     async def _process_response(
         self,

src/kernel/lib/browser_routing/routing.py

@@ -32,6 +32,9 @@ class BrowserRoutingConfig:
     subresources: tuple[str, ...] = field(default_factory=tuple)
 
 
+_EVICTION_HOOK_CACHE_ATTR = "_kernel_browser_route_cache"
+
+
 _BROWSER_ROUTE_CACHEABLE_PATH = re.compile(r"^/(?:v\d+/)?browsers(?:/[^/]+)?/?$")
 _BROWSER_DELETE_BY_ID_PATH = re.compile(r"^/(?:v\d+/)?browsers/([^/]+)/?$")
 _BROWSER_POOL_ACQUIRE_PATH = re.compile(r"^/(?:v\d+/)?browser_pools/[^/]+/acquire/?$")
@@ -198,6 +201,48 @@ def is_stale_direct_vm_auth_response(response: httpx.Response) -> bool:
     return bool(response.request.url.params.get("jwt"))
 
 
+def install_stale_direct_vm_auth_eviction(client: httpx.Client, *, cache: BrowserRouteCache) -> None:
+    """Evict stale direct-to-VM routes as soon as the response status is known.
+
+    httpx reads the body of a non-streamed response inside `send()`, so a caller
+    that only inspects the returned response never learns the status of a 401/403
+    whose body read fails — the read error surfaces from `send()` instead and the
+    dead route would stay cached, wedging every later call for that session. A
+    response event hook runs after the status is known and before any body is
+    read, which keeps eviction independent of the body.
+    """
+    hooks = client.event_hooks.setdefault("response", [])
+    if _has_eviction_hook(hooks, cache):
+        return
+
+    def evict(response: httpx.Response) -> None:
+        if is_stale_direct_vm_auth_response(response):
+            maybe_evict_browser_route_from_response(response, cache=cache)
+
+    setattr(evict, _EVICTION_HOOK_CACHE_ATTR, cache)
+    hooks.append(evict)
+
+
+def install_async_stale_direct_vm_auth_eviction(client: httpx.AsyncClient, *, cache: BrowserRouteCache) -> None:
+    """Async counterpart of `install_stale_direct_vm_auth_eviction`."""
+    hooks = client.event_hooks.setdefault("response", [])
+    if _has_eviction_hook(hooks, cache):
+        return
+
+    async def evict(response: httpx.Response) -> None:
+        if is_stale_direct_vm_auth_response(response):
+            maybe_evict_browser_route_from_response(response, cache=cache)
+
+    setattr(evict, _EVICTION_HOOK_CACHE_ATTR, cache)
+    hooks.append(evict)
+
+
+def _has_eviction_hook(hooks: list[Any], cache: BrowserRouteCache) -> bool:
+    # A copied client shares both the httpx client and the route cache, so the
+    # hook is registered once per cache instead of once per client.
+    return any(getattr(hook, _EVICTION_HOOK_CACHE_ATTR, None) is cache for hook in hooks)
+
+
 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

@@ -3,7 +3,7 @@
 import io
 import os
 import asyncio
-from typing import Any, AsyncIterator, cast
+from typing import Any, Iterator, AsyncIterator, cast
 from pathlib import Path
 from typing_extensions import override
 
@@ -14,6 +14,7 @@
 from kernel import (
     Kernel,
     AsyncKernel,
+    APIConnectionError,
     AuthenticationError,
     InternalServerError,
     PermissionDeniedError,
@@ -1259,3 +1260,96 @@ def test_indexed_multipart_body_flattens_only_given_values() -> None:
         "files[1][dest_path]": "/tmp/two",
         "flag": True,
     }
+
+
+class _FailingSyncStream(httpx.SyncByteStream):
+    """A response body that fails while it is being read."""
+
+    @override
+    def __iter__(self) -> Iterator[bytes]:
+        raise httpx.ReadError("connection reset while reading the error body")
+
+
+class _FailingAsyncStream(httpx.AsyncByteStream):
+    @override
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        raise httpx.ReadError("connection reset while reading the error body")
+        yield b""  # pragma: no cover - unreachable, keeps this an async generator
+
+
+def test_stale_direct_vm_jwt_evicts_route_when_error_body_read_fails(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False)
+    requests: list[httpx.Request] = []
+
+    def handle_request(request: httpx.Request) -> httpx.Response:
+        requests.append(request)
+        if "browser-session.test" in str(request.url):
+            return httpx.Response(401, stream=_FailingSyncStream(), headers={"content-type": "text/plain"})
+        return httpx.Response(200, content=b"png", headers={"content-type": "image/png"})
+
+    http_client = httpx.Client(transport=httpx.MockTransport(handle_request))
+    with Kernel(
+        base_url=base_url,
+        api_key=api_key,
+        max_retries=0,
+        http_client=http_client,
+        _strict_response_validation=True,
+    ) as client:
+        _cache_browser(client)
+        with pytest.raises(APIConnectionError):
+            client.browsers.computer.capture_screenshot("sess-1")
+        # The status was known before the body read failed, so the dead route is gone.
+        assert client.browser_route_cache.get("sess-1") is None
+
+        client.browsers.computer.capture_screenshot("sess-1")
+
+    assert str(requests[0].url).startswith("http://browser-session.test/browser/kernel/computer/screenshot")
+    assert requests[1].url == httpx.URL(f"{base_url}/browsers/sess-1/computer/screenshot")
+    assert requests[1].url.params.get("jwt") is None
+    assert requests[1].headers.get("Authorization") == f"Bearer {api_key}"
+
+
+@pytest.mark.asyncio
+async def test_async_stale_direct_vm_jwt_evicts_route_when_error_body_read_fails(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False)
+    requests: list[httpx.Request] = []
+
+    async def handle_request(request: httpx.Request) -> httpx.Response:
+        requests.append(request)
+        if "browser-session.test" in str(request.url):
+            return httpx.Response(403, stream=_FailingAsyncStream(), headers={"content-type": "text/plain"})
+        return httpx.Response(200, content=b"png", headers={"content-type": "image/png"})
+
+    http_client = httpx.AsyncClient(transport=httpx.MockTransport(handle_request))
+    async with AsyncKernel(
+        base_url=base_url,
+        api_key=api_key,
+        max_retries=0,
+        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.computer.capture_screenshot("sess-1")
+        assert client.browser_route_cache.get("sess-1") is None
+
+        await client.browsers.computer.capture_screenshot("sess-1")
+
+    assert str(requests[0].url).startswith("http://browser-session.test/browser/kernel/computer/screenshot")
+    assert requests[1].url == httpx.URL(f"{base_url}/browsers/sess-1/computer/screenshot")
+    assert requests[1].url.params.get("jwt") is None
+    assert requests[1].headers.get("Authorization") == f"Bearer {api_key}"
+
+
+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")
+        assert copied.browser_route_cache is client.browser_route_cache
+        hooks = client._client.event_hooks["response"]  # pyright: ignore[reportPrivateUsage]
+        assert len(hooks) == 1