Prevent retries of unreplayable VM bodies

Author: tnsardesai

src/kernel/_client.py

@@ -47,6 +47,7 @@
     maybe_evict_browser_route_from_response,
     should_retry_direct_vm_connection_error,
     install_async_stale_direct_vm_auth_eviction,
+    direct_vm_request_body_is_known_unreplayable,
     maybe_populate_browser_route_cache_from_response,
 )
 
@@ -382,6 +383,8 @@ def _should_retry_on_connection_error(self, request: httpx.Request) -> bool:
 
     @override
     def _should_retry(self, response: httpx.Response) -> bool:
+        if direct_vm_request_body_is_known_unreplayable(response.request):
+            return False
         if is_stale_direct_vm_auth_response(response):
             # The route was already evicted by the response hook; retry only when
             # the body can be rebuilt, otherwise the caller sees the original auth
@@ -778,6 +781,8 @@ def _should_retry_on_connection_error(self, request: httpx.Request) -> bool:
 
     @override
     def _should_retry(self, response: httpx.Response) -> bool:
+        if direct_vm_request_body_is_known_unreplayable(response.request):
+            return False
         if is_stale_direct_vm_auth_response(response):
             # The route was already evicted by the response hook; retry only when
             # the body can be rebuilt, otherwise the caller sees the original auth

src/kernel/lib/browser_routing/routing.py

@@ -251,13 +251,9 @@ def _has_eviction_hook(hooks: list[Any], cache: BrowserRouteCache) -> bool:
 
 
 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.
-    """
+    """Prevent connection retries from replaying an unreplayable VM request body."""
+    if direct_vm_request_body_is_known_unreplayable(request):
+        return False
     if not request.extensions.get(_STALE_DIRECT_VM_AUTH_REQUEST_EXTENSION):
         return True
     return direct_vm_request_body_is_replayable(request)
@@ -276,6 +272,10 @@ def should_retry_stale_direct_vm_auth(response: httpx.Response) -> bool:
     return direct_vm_request_body_is_replayable(response.request)
 
 
+def direct_vm_request_body_is_known_unreplayable(request: httpx.Request) -> bool:
+    return request.extensions.get(_DIRECT_VM_BODY_REPLAYABLE_REQUEST_EXTENSION) is False
+
+
 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):

tests/test_browser_routing.py

@@ -1468,6 +1468,106 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
     assert requests[1] == (httpx.URL(f"{base_url}/browsers/sess-1/fs/write_file?path=%2Ftmp%2Fx"), b"next")
 
 
+@pytest.mark.parametrize(
+    ("failure_type", "expected_error"),
+    [
+        (httpx.ReadError, APIConnectionError),
+        (httpx.ReadTimeout, APITimeoutError),
+        (None, InternalServerError),
+    ],
+)
+def test_direct_vm_failure_does_not_retry_unreplayable_write(
+    monkeypatch: pytest.MonkeyPatch,
+    failure_type: type[httpx.TransportError] | None,
+    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)
+    requests: list[tuple[httpx.URL, bytes]] = []
+
+    class Transport(httpx.BaseTransport):
+        @override
+        def handle_request(self, request: httpx.Request) -> httpx.Response:
+            body = b"".join(cast(Iterator[bytes], request.stream))
+            requests.append((request.url, body))
+            if failure_type is not None:
+                raise failure_type("connection failed after sending the request", request=request)
+            return httpx.Response(500, json={"error": "boom"})
+
+    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(expected_error):
+            client.browsers.fs.write_file("sess-1", _UnseekableFile(b"payload"), path="/tmp/x")
+
+        assert client.browser_route_cache.get("sess-1") is not None
+
+    assert requests == [
+        (
+            httpx.URL("http://browser-session.test/browser/kernel/fs/write_file?path=%2Ftmp%2Fx&jwt=token-abc"),
+            b"payload",
+        )
+    ]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("failure_type", "expected_error"),
+    [
+        (httpx.ReadError, APIConnectionError),
+        (httpx.ReadTimeout, APITimeoutError),
+        (None, InternalServerError),
+    ],
+)
+async def test_async_direct_vm_failure_does_not_retry_unreplayable_write(
+    monkeypatch: pytest.MonkeyPatch,
+    failure_type: type[httpx.TransportError] | None,
+    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)
+    requests: list[tuple[httpx.URL, bytes]] = []
+
+    async def payload() -> AsyncIterator[bytes]:
+        yield b"payload"
+
+    class Transport(httpx.AsyncBaseTransport):
+        @override
+        async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+            body = b"".join([chunk async for chunk in cast(AsyncIterator[bytes], request.stream)])
+            requests.append((request.url, body))
+            if failure_type is not None:
+                raise failure_type("connection failed after sending the request", request=request)
+            return httpx.Response(500, json={"error": "boom"})
+
+    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(expected_error):
+            await client.browsers.fs.write_file("sess-1", cast(Any, payload()), path="/tmp/x")
+
+        assert client.browser_route_cache.get("sess-1") is not None
+
+    assert requests == [
+        (
+            httpx.URL("http://browser-session.test/browser/kernel/fs/write_file?path=%2Ftmp%2Fx&jwt=token-abc"),
+            b"payload",
+        )
+    ]
+
+
 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")