Evict the originating browser route cache

Author: tnsardesai

src/kernel/_client.py

@@ -378,7 +378,7 @@ def _prepare_options(self, options: Any) -> Any:
 
     @override
     def _prepare_request(self, request: httpx.Request) -> None:
-        prepare_direct_vm_request(request)
+        prepare_direct_vm_request(request, cache=self.browser_route_cache)
 
     @override
     def _should_retry_on_connection_error(self, request: httpx.Request) -> bool:
@@ -777,7 +777,7 @@ async def _prepare_options(self, options: Any) -> Any:
 
     @override
     async def _prepare_request(self, request: httpx.Request) -> None:
-        prepare_direct_vm_request(request)
+        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

@@ -38,6 +38,7 @@ class BrowserRoutingConfig:
 _DIRECT_VM_REQUEST_MARKER_HEADER = "x-kernel-direct-vm-request"
 _STALE_DIRECT_VM_AUTH_REQUEST_EXTENSION = "kernel_stale_direct_vm_auth"
 _DIRECT_VM_BODY_REPLAYABLE_REQUEST_EXTENSION = "kernel_direct_vm_body_replayable"
+_DIRECT_VM_ROUTE_CACHE_REQUEST_EXTENSION = "kernel_direct_vm_route_cache"
 
 
 _BROWSER_ROUTE_CACHEABLE_PATH = re.compile(r"^/(?:v\d+/)?browsers(?:/[^/]+)?/?$")
@@ -225,10 +226,11 @@ def install_stale_direct_vm_auth_eviction(client: httpx.Client, *, cache: Browse
         return
 
     def handle_response(response: httpx.Response) -> None:
-        _reject_unreplayable_direct_vm_redirect(response, cache=cache)
+        request_cache = _direct_vm_route_cache(response.request) or cache
+        _reject_unreplayable_direct_vm_redirect(response, cache=request_cache)
         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)
+            maybe_evict_browser_route_from_response(response, cache=request_cache)
 
     setattr(handle_response, _EVICTION_HOOK_CACHE_ATTR, cache)
     hooks.insert(0, handle_response)
@@ -241,10 +243,11 @@ def install_async_stale_direct_vm_auth_eviction(client: httpx.AsyncClient, *, ca
         return
 
     async def handle_response(response: httpx.Response) -> None:
-        _reject_unreplayable_direct_vm_redirect(response, cache=cache)
+        request_cache = _direct_vm_route_cache(response.request) or cache
+        _reject_unreplayable_direct_vm_redirect(response, cache=request_cache)
         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)
+            maybe_evict_browser_route_from_response(response, cache=request_cache)
 
     setattr(handle_response, _EVICTION_HOOK_CACHE_ATTR, cache)
     hooks.insert(0, handle_response)
@@ -466,7 +469,7 @@ def rewrite_direct_vm_options(
     return rewritten
 
 
-def prepare_direct_vm_request(request: httpx.Request) -> None:
+def prepare_direct_vm_request(request: httpx.Request, *, cache: BrowserRouteCache) -> None:
     if request.headers.pop(_DIRECT_VM_REQUEST_MARKER_HEADER, None) is None:
         return
 
@@ -475,9 +478,15 @@ def prepare_direct_vm_request(request: httpx.Request) -> None:
     request.extensions[_DIRECT_VM_BODY_REPLAYABLE_REQUEST_EXTENSION] = _classify_direct_vm_request_body_replayability(
         request
     )
+    request.extensions[_DIRECT_VM_ROUTE_CACHE_REQUEST_EXTENSION] = cache
     request.headers.pop("Authorization", None)
 
 
+def _direct_vm_route_cache(request: httpx.Request) -> BrowserRouteCache | None:
+    cache = request.extensions.get(_DIRECT_VM_ROUTE_CACHE_REQUEST_EXTENSION)
+    return cache if isinstance(cache, BrowserRouteCache) else None
+
+
 def _is_direct_vm_request(request: httpx.Request) -> bool:
     return isinstance(request.extensions.get(_DIRECT_VM_BODY_REPLAYABLE_REQUEST_EXTENSION), bool)
 

tests/test_browser_routing.py

@@ -1726,6 +1726,113 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
     ]
 
 
+@pytest.mark.parametrize("origin_first", [True, False])
+def test_direct_vm_redirect_evicts_originating_cache_when_http_client_is_shared(
+    monkeypatch: pytest.MonkeyPatch,
+    origin_first: bool,
+) -> None:
+    monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False)
+    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))
+            return httpx.Response(
+                307,
+                headers={"location": "http://other-vm.test/browser/kernel/fs/write_file"},
+            )
+
+    http_client = httpx.Client(transport=Transport(), follow_redirects=True)
+
+    def make_client() -> Kernel:
+        return Kernel(
+            base_url=base_url,
+            api_key=api_key,
+            http_client=http_client,
+            _strict_response_validation=True,
+        )
+
+    if origin_first:
+        client, other_client = make_client(), make_client()
+    else:
+        other_client, client = make_client(), make_client()
+
+    try:
+        _cache_browser(client)
+        with pytest.raises(APIConnectionError):
+            client.browsers.fs.write_file("sess-1", _UnseekableFile(b"payload"), path="/tmp/x")
+        assert client.browser_route_cache.get("sess-1") is None
+        assert other_client.browser_route_cache.get("sess-1") is None
+    finally:
+        http_client.close()
+
+    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("origin_first", [True, False])
+async def test_async_direct_vm_redirect_evicts_originating_cache_when_http_client_is_shared(
+    monkeypatch: pytest.MonkeyPatch,
+    origin_first: bool,
+) -> None:
+    monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False)
+    requests: list[tuple[httpx.URL, bytes]] = []
+
+    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))
+            return httpx.Response(
+                308,
+                headers={"location": "http://other-vm.test/browser/kernel/fs/write_file"},
+            )
+
+    http_client = httpx.AsyncClient(transport=Transport(), follow_redirects=True)
+
+    def make_client() -> AsyncKernel:
+        return AsyncKernel(
+            base_url=base_url,
+            api_key=api_key,
+            http_client=http_client,
+            _strict_response_validation=True,
+        )
+
+    if origin_first:
+        client, other_client = make_client(), make_client()
+    else:
+        other_client, client = make_client(), make_client()
+
+    try:
+        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, _UnreplayableAsyncBody(b"payload")),
+                path="/tmp/x",
+            )
+        assert client.browser_route_cache.get("sess-1") is None
+        assert other_client.browser_route_cache.get("sess-1") is None
+    finally:
+        await http_client.aclose()
+
+    assert requests == [
+        (
+            httpx.URL("http://browser-session.test/browser/kernel/fs/write_file?path=%2Ftmp%2Fx&jwt=token-abc"),
+            b"payload",
+        )
+    ]
+
+
 def test_direct_vm_auth_is_removed_after_httpx_and_request_hooks(monkeypatch: pytest.MonkeyPatch) -> None:
     monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False)
     requests: list[httpx.Request] = []