diff --git a/src/lib/browser-routing.ts b/src/lib/browser-routing.ts index d0272a7..673812e 100644 --- a/src/lib/browser-routing.ts +++ b/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( diff --git a/tests/lib/browser-routing.test.ts b/tests/lib/browser-routing.test.ts index 9332023..d6fb7ee 100644 --- a/tests/lib/browser-routing.test.ts +++ b/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' }); + }); + }); });