Retry stale direct-VM auth without needing the cache

A metro 401/403 with a jwt query param is enough to fall back to the
control plane, so concurrent requests still retry after the first
eviction.

Author: tnsardesai

src/kernel/_client.py

@@ -356,7 +356,7 @@ def _prepare_request(self, request: httpx.Request) -> None:
 
     @override
     def _should_retry(self, response: httpx.Response) -> bool:
-        if should_retry_stale_direct_vm_auth(response, cache=self.browser_route_cache):
+        if should_retry_stale_direct_vm_auth(response):
             maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache)
             return True
         return super()._should_retry(response)
@@ -732,7 +732,7 @@ async def _prepare_request(self, request: httpx.Request) -> None:
 
     @override
     def _should_retry(self, response: httpx.Response) -> bool:
-        if should_retry_stale_direct_vm_auth(response, cache=self.browser_route_cache):
+        if should_retry_stale_direct_vm_auth(response):
             maybe_evict_browser_route_from_response(response, cache=self.browser_route_cache)
             return True
         return super()._should_retry(response)

src/kernel/lib/browser_routing/routing.py

@@ -116,7 +116,7 @@ def maybe_evict_browser_route_from_response(response: httpx.Response, *, cache:
             cache.delete(session_id)
         return
 
-    if response.status_code not in {401, 403}:
+    if not is_stale_direct_vm_auth_response(response):
         return
 
     session_id = _session_id_from_direct_vm_response(response, cache=cache)
@@ -173,10 +173,14 @@ def _session_id_from_direct_vm_response(response: httpx.Response, *, cache: Brow
     return None
 
 
-def should_retry_stale_direct_vm_auth(response: httpx.Response, *, cache: BrowserRouteCache) -> bool:
+def is_stale_direct_vm_auth_response(response: httpx.Response) -> bool:
     if response.status_code not in {401, 403}:
         return False
-    return _session_id_from_direct_vm_response(response, cache=cache) is not None
+    return bool(response.request.url.params.get("jwt"))
+
+
+def should_retry_stale_direct_vm_auth(response: httpx.Response) -> bool:
+    return is_stale_direct_vm_auth_response(response)
 
 
 def _session_id_from_browser_delete_path(path: str) -> str | None:

tests/test_browser_routing.py

@@ -546,3 +546,16 @@ def _skip_retry_sleep(_self: object, **_kwargs: object) -> None:
     assert api.called
     api_req = cast(httpx.Request, cast(Any, api.calls[0]).request)
     assert api_req.headers.get("Authorization") == f"Bearer {api_key}"
+
+
+def test_stale_direct_vm_auth_retry_does_not_require_cached_route() -> None:
+    from kernel.lib.browser_routing.routing import should_retry_stale_direct_vm_auth
+
+    request = httpx.Request(
+        "POST",
+        "http://browser-session.test/browser/kernel/computer/screenshot?jwt=token-abc",
+    )
+    response = httpx.Response(401, text="Invalid JWT", request=request)
+    empty = BrowserRouteCache()
+    assert should_retry_stale_direct_vm_auth(response) is True
+    assert empty.get("sess-1") is None