Preserve direct request ownership

Author: tnsardesai

src/kernel/_client.py

@@ -213,7 +213,7 @@ def __init__(
         self.browser_route_cache = _browser_route_cache or BrowserRouteCache()
         self._browser_routing = browser_routing_config_from_env()
         install_direct_vm_auth_stripping(self._client)
-        install_stale_direct_vm_auth_eviction(self._client, cache=self.browser_route_cache)
+        install_stale_direct_vm_auth_eviction(self._client)
 
     @cached_property
     def deployments(self) -> DeploymentsResource:
@@ -612,7 +612,7 @@ def __init__(
         self.browser_route_cache = _browser_route_cache or BrowserRouteCache()
         self._browser_routing = browser_routing_config_from_env()
         install_async_direct_vm_auth_stripping(self._client)
-        install_async_stale_direct_vm_auth_eviction(self._client, cache=self.browser_route_cache)
+        install_async_stale_direct_vm_auth_eviction(self._client)
 
     @cached_property
     def deployments(self) -> AsyncDeploymentsResource:

src/kernel/lib/browser_routing/raw_http.py

@@ -7,7 +7,7 @@
 import httpx
 
 from .util import sanitize_curl_raw_params
-from .routing import BrowserRoute
+from .routing import BrowserRoute, mark_direct_vm_headers
 from ..._types import Body, Timeout, NotGiven, not_given
 from ..._models import FinalRequestOptions
 
@@ -33,7 +33,7 @@ def request_via_browser_route(
         method=method.upper(),
         url=route.base_url.rstrip("/") + "/curl/raw",
         params=query,
-        headers=headers or {},
+        headers=mark_direct_vm_headers(headers),
         content=_normalize_binary_content(content),
         json_data=json,
         timeout=_normalize_timeout(timeout),
@@ -93,7 +93,7 @@ async def async_request_via_browser_route(
         method=method.upper(),
         url=route.base_url.rstrip("/") + "/curl/raw",
         params=query,
-        headers=headers or {},
+        headers=mark_direct_vm_headers(headers),
         content=_normalize_binary_content(content),
         json_data=json,
         timeout=_normalize_timeout(timeout),

src/kernel/lib/browser_routing/routing.py

@@ -33,7 +33,7 @@ class BrowserRoutingConfig:
     subresources: tuple[str, ...] = field(default_factory=tuple)
 
 
-_EVICTION_HOOK_CACHE_ATTR = "_kernel_browser_route_cache"
+_DIRECT_VM_EVICTION_HOOK_ATTR = "_kernel_direct_vm_eviction_hook"
 _DIRECT_VM_AUTH_HOOK_ATTR = "_kernel_direct_vm_auth_hook"
 _DIRECT_VM_REQUEST_MARKER_HEADER = "x-kernel-direct-vm-request"
 _STALE_DIRECT_VM_AUTH_REQUEST_EXTENSION = "kernel_stale_direct_vm_auth"
@@ -207,7 +207,7 @@ 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:
+def install_stale_direct_vm_auth_eviction(client: httpx.Client) -> 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
@@ -222,34 +222,38 @@ def install_stale_direct_vm_auth_eviction(client: httpx.Client, *, cache: Browse
     failing body or raising.
     """
     hooks = client.event_hooks.setdefault("response", [])
-    if _has_eviction_hook(hooks, cache):
+    if any(getattr(hook, _DIRECT_VM_EVICTION_HOOK_ATTR, False) for hook in hooks):
         return
 
     def handle_response(response: httpx.Response) -> None:
-        request_cache = _direct_vm_route_cache(response.request) or cache
+        request_cache = _direct_vm_route_cache(response.request)
+        if request_cache is None:
+            return
         _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=request_cache)
 
-    setattr(handle_response, _EVICTION_HOOK_CACHE_ATTR, cache)
+    setattr(handle_response, _DIRECT_VM_EVICTION_HOOK_ATTR, True)
     hooks.insert(0, handle_response)
 
 
-def install_async_stale_direct_vm_auth_eviction(client: httpx.AsyncClient, *, cache: BrowserRouteCache) -> None:
+def install_async_stale_direct_vm_auth_eviction(client: httpx.AsyncClient) -> None:
     """Async counterpart of `install_stale_direct_vm_auth_eviction`."""
     hooks = client.event_hooks.setdefault("response", [])
-    if _has_eviction_hook(hooks, cache):
+    if any(getattr(hook, _DIRECT_VM_EVICTION_HOOK_ATTR, False) for hook in hooks):
         return
 
     async def handle_response(response: httpx.Response) -> None:
-        request_cache = _direct_vm_route_cache(response.request) or cache
+        request_cache = _direct_vm_route_cache(response.request)
+        if request_cache is None:
+            return
         _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=request_cache)
 
-    setattr(handle_response, _EVICTION_HOOK_CACHE_ATTR, cache)
+    setattr(handle_response, _DIRECT_VM_EVICTION_HOOK_ATTR, True)
     hooks.insert(0, handle_response)
 
 
@@ -281,12 +285,6 @@ async def strip_auth(request: httpx.Request) -> None:
     hooks.append(strip_auth)
 
 
-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 _reject_unreplayable_direct_vm_redirect(response: httpx.Response, *, cache: BrowserRouteCache) -> None:
     if not response.has_redirect_location or not direct_vm_request_body_is_known_unreplayable(response.request):
         return
@@ -469,6 +467,12 @@ def rewrite_direct_vm_options(
     return rewritten
 
 
+def mark_direct_vm_headers(headers: Mapping[str, str] | None) -> dict[str, str]:
+    marked = dict(headers or {})
+    marked[_DIRECT_VM_REQUEST_MARKER_HEADER] = "true"
+    return marked
+
+
 def prepare_direct_vm_request(request: httpx.Request, *, cache: BrowserRouteCache) -> None:
     if request.headers.pop(_DIRECT_VM_REQUEST_MARKER_HEADER, None) is None:
         return

tests/test_browser_routing.py

@@ -190,6 +190,8 @@ def test_browser_request_uses_curl_raw() -> None:
     request = cast(httpx.Request, cast(Any, route.calls[0]).request)
     assert "curl/raw" in str(request.url)
     assert request.url.params.get("jwt") == "token-abc"
+    assert request.headers.get("Authorization") is None
+    assert request.headers.get("x-kernel-direct-vm-request") is None
 
 
 @respx.mock
@@ -344,6 +346,8 @@ async def test_async_raw_browser_create_warms_route_cache() -> None:
     assert routed.content == b"ok"
     request = cast(httpx.Request, cast(Any, routed_request.calls[0]).request)
     assert request.url.params.get("jwt") == "token-abc"
+    assert request.headers.get("Authorization") is None
+    assert request.headers.get("x-kernel-direct-vm-request") is None
 
 
 @respx.mock
@@ -1727,9 +1731,15 @@ 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(
+@pytest.mark.parametrize(
+    ("status_code", "expected_error"),
+    [(307, APIConnectionError), (401, AuthenticationError)],
+)
+def test_direct_vm_failure_evicts_originating_cache_when_http_client_is_shared(
     monkeypatch: pytest.MonkeyPatch,
     origin_first: bool,
+    status_code: int,
+    expected_error: type[Exception],
 ) -> None:
     monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False)
     requests: list[tuple[httpx.URL, bytes]] = []
@@ -1740,7 +1750,7 @@ 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,
+                status_code,
                 headers={"location": "http://other-vm.test/browser/kernel/fs/write_file"},
             )
 
@@ -1760,8 +1770,9 @@ def make_client() -> Kernel:
         other_client, client = make_client(), make_client()
 
     try:
+        assert len(http_client.event_hooks["response"]) == 1
         _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 client.browser_route_cache.get("sess-1") is None
         assert other_client.browser_route_cache.get("sess-1") is None
@@ -1778,9 +1789,15 @@ def make_client() -> Kernel:
 
 @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(
+@pytest.mark.parametrize(
+    ("status_code", "expected_error"),
+    [(308, APIConnectionError), (403, PermissionDeniedError)],
+)
+async def test_async_direct_vm_failure_evicts_originating_cache_when_http_client_is_shared(
     monkeypatch: pytest.MonkeyPatch,
     origin_first: bool,
+    status_code: int,
+    expected_error: type[Exception],
 ) -> None:
     monkeypatch.delenv("KERNEL_BROWSER_ROUTING_SUBRESOURCES", raising=False)
     requests: list[tuple[httpx.URL, bytes]] = []
@@ -1791,7 +1808,7 @@ 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,
+                status_code,
                 headers={"location": "http://other-vm.test/browser/kernel/fs/write_file"},
             )
 
@@ -1811,10 +1828,11 @@ def make_client() -> AsyncKernel:
         other_client, client = make_client(), make_client()
 
     try:
+        assert len(http_client.event_hooks["response"]) == 1
         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, _UnreplayableAsyncBody(b"payload")),