Fall back to the control plane on a stale session JWT

src/lib/browser-routing.ts

@@ -1,4 +1,5 @@
 import type { Fetch, RequestInfo, RequestInit } from '../internal/builtin-types';
+import { CancelReadableStream } from '../internal/shims';
 import { joinURL } from './join-url';
 
 export type BrowserRoute = {
@@ -22,6 +23,15 @@ export class BrowserRouteCache {
     this.entries.delete(sessionId);
   }
 
+  deleteIfJwt(sessionId: string, jwt: string): boolean {
+    const route = this.entries.get(sessionId);
+    if (!route || route.jwt !== jwt) {
+      return false;
+    }
+    this.entries.delete(sessionId);
+    return true;
+  }
+
   clear(): void {
     this.entries.clear();
   }
@@ -262,7 +272,13 @@ async function routeRequest(
 
   const headers = new Headers(request.headers);
   headers.delete('authorization');
-  return innerFetch(target.toString(), buildRoutedInit(input, request, init, headers));
+  const routed = await innerFetch(target.toString(), buildRoutedInit(input, request, init, headers));
+  if ((routed.status === 401 || routed.status === 403) && target.searchParams.get('jwt')) {
+    cache.deleteIfJwt(sessionId, target.searchParams.get('jwt') ?? '');
+    await CancelReadableStream(routed.body);
+    return innerFetch(input, init);
+  }
+  return routed;
 }
 
 function buildRoutedInit(

tests/lib/browser-routing.test.ts

@@ -590,4 +590,78 @@ describe('browser routing', () => {
       ]);
     });
   });
+
+  test('falls back to the control plane on a stale direct-VM JWT', async () => {
+    await withBrowserRoutingEnv(undefined, async () => {
+      const calls: Array<{ url: string; headers: Headers }> = [];
+      const kernel = new Kernel({
+        apiKey: 'k',
+        baseURL: 'https://api.example/',
+        fetch: async (input, init?: RequestInit) => {
+          const url = normalizeURL(input);
+          const headers = input instanceof Request ? new Headers(input.headers) : new Headers(init?.headers);
+          calls.push({ url, headers });
+          if (url === 'https://api.example/browsers') {
+            return Response.json({
+              session_id: 'sess-1',
+              base_url: 'http://browser-session.test/browser/kernel',
+              cdp_ws_url: 'wss://browser-session.test/browser/cdp?jwt=token-abc',
+            });
+          }
+          if (url.includes('browser-session.test')) {
+            return new Response('Invalid JWT', { status: 401, headers: { 'content-type': 'text/plain' } });
+          }
+          return new Response(new Uint8Array([1, 2, 3]), {
+            status: 200,
+            headers: { 'content-type': 'image/png' },
+          });
+        },
+      });
+
+      await kernel.browsers.create();
+      await kernel.browsers.computer.captureScreenshot('sess-1');
+
+      expect(calls[1]?.url).toBe(
+        'http://browser-session.test/browser/kernel/computer/screenshot?jwt=token-abc',
+      );
+      expect(calls[2]?.url).toBe('https://api.example/browsers/sess-1/computer/screenshot');
+      expect(calls[2]?.headers.get('authorization')).toBe('Bearer k');
+      expect(kernel.browserRouteCache.get('sess-1')).toBeUndefined();
+    });
+  });
+
+  test('does not evict a refreshed route after a stale JWT 401', async () => {
+    await withBrowserRoutingEnv(undefined, async () => {
+      const kernel = new Kernel({
+        apiKey: 'k',
+        baseURL: 'https://api.example/',
+        fetch: async (input) => {
+          const url = normalizeURL(input);
+          if (url === 'https://api.example/browsers') {
+            return Response.json({
+              session_id: 'sess-1',
+              base_url: 'http://browser-session.test/browser/kernel',
+              cdp_ws_url: 'wss://browser-session.test/browser/cdp?jwt=token-abc',
+            });
+          }
+          if (url.includes('browser-session.test')) {
+            kernel.browserRouteCache.set({
+              sessionId: 'sess-1',
+              baseURL: 'http://browser-session.test/browser/kernel',
+              jwt: 'jwt-FRESH',
+            });
+            return new Response('Invalid JWT', { status: 401, headers: { 'content-type': 'text/plain' } });
+          }
+          return new Response(new Uint8Array([1, 2, 3]), {
+            status: 200,
+            headers: { 'content-type': 'image/png' },
+          });
+        },
+      });
+
+      await kernel.browsers.create();
+      await kernel.browsers.computer.captureScreenshot('sess-1');
+      expect(kernel.browserRouteCache.get('sess-1')).toMatchObject({ jwt: 'jwt-FRESH' });
+    });
+  });
 });