Only evict a stale JWT if it is still the cached one

A later 401 must not delete a route that was refreshed in flight.

Author: tnsardesai

src/kernel/lib/browser_routing/routing.py

@@ -69,6 +69,14 @@ def set(self, route: BrowserRoute) -> None:
     def delete(self, session_id: str) -> None:
         self._routes.pop(_normalize_session_id(session_id), None)
 
+    def delete_if_jwt(self, session_id: str, jwt: str) -> bool:
+        key = _normalize_session_id(session_id)
+        route = self._routes.get(key)
+        if route is None or route.jwt != jwt.strip():
+            return False
+        del self._routes[key]
+        return True
+
     def values(self) -> list[BrowserRoute]:
         return list(self._routes.values())
 
@@ -119,9 +127,10 @@ def maybe_evict_browser_route_from_response(response: httpx.Response, *, cache:
     if not is_stale_direct_vm_auth_response(response):
         return
 
+    jwt = str(response.request.url.params.get("jwt") or "").strip()
     session_id = _session_id_from_direct_vm_response(response, cache=cache)
-    if session_id:
-        cache.delete(session_id)
+    if session_id and jwt:
+        cache.delete_if_jwt(session_id, jwt)
 
 
 def populate_browser_route_cache_from_value(value: object, *, cache: BrowserRouteCache) -> None:

tests/test_browser_routing.py

@@ -548,6 +548,37 @@ def _skip_retry_sleep(_self: object, **_kwargs: object) -> None:
     assert api_req.headers.get("Authorization") == f"Bearer {api_key}"
 
 
+def test_stale_direct_vm_jwt_does_not_evict_refreshed_route() -> None:
+    from kernel.lib.browser_routing.routing import maybe_evict_browser_route_from_response
+
+    cache = BrowserRouteCache()
+    cache.set(
+        BrowserRoute(
+            session_id="sess-1",
+            base_url="http://browser-session.test/browser/kernel",
+            jwt="token-abc",
+        )
+    )
+    cache.set(
+        BrowserRoute(
+            session_id="sess-1",
+            base_url="http://browser-session.test/browser/kernel",
+            jwt="jwt-FRESH",
+        )
+    )
+    request = httpx.Request(
+        "POST",
+        "http://browser-session.test/browser/kernel/computer/screenshot?jwt=token-abc",
+    )
+    maybe_evict_browser_route_from_response(
+        httpx.Response(401, text="Invalid JWT", request=request),
+        cache=cache,
+    )
+    route = cache.get("sess-1")
+    assert route is not None
+    assert route.jwt == "jwt-FRESH"
+
+
 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