chore: implement local browser rendering /v1/session endpoint (#10220)

* chore: implement local browser rendering /v1/session endpoint

* fix: ensure websocket is closed

* test: ensure puppeteer is pre-installed for miniflare tests

Author: ruifigueira

.changeset/petite-paws-allow.md

@@ -0,0 +1,5 @@
+---
+"miniflare": minor
+---
+
+Add /v1/session endpoint for Browser Rendering local mode

fixtures/browser-rendering/src/playwright.ts

@@ -32,6 +32,24 @@ export default {
 					const pText = await page.locator("p").first().textContent();
 					return new Response(pText);
 				}
+
+				case "disconnect": {
+					const { sessionId } = await playwright.acquire(env.MYBROWSER);
+					const browser = await playwright.connect(env.MYBROWSER, sessionId);
+					// closing a browser obtained with playwright.connect actually disconnects
+					// (it doesn's close the porcess)
+					await browser.close();
+					const sessionInfo = await playwright
+						.sessions(env.MYBROWSER)
+						.then((sessions) =>
+							sessions.find((s) => s.sessionId === sessionId)
+						);
+					return new Response(
+						sessionInfo.connectionId
+							? "Browser not disconnected"
+							: "Browser disconnected"
+					);
+				}
 			}
 
 			let img = await env.BROWSER_KV_DEMO.get(url, { type: "arrayBuffer" });

fixtures/browser-rendering/src/puppeteer.ts

@@ -32,6 +32,22 @@ export default {
 					const pText = await page.$eval("p", (el) => el.textContent.trim());
 					return new Response(pText);
 				}
+
+				case "disconnect": {
+					const browser = await puppeteer.launch(env.MYBROWSER);
+					const sessionId = browser.sessionId();
+					await browser.disconnect();
+					const sessionInfo = await puppeteer
+						.sessions(env.MYBROWSER)
+						.then((sessions) =>
+							sessions.find((s) => s.sessionId === sessionId)
+						);
+					return new Response(
+						sessionInfo.connectionId
+							? "Browser not disconnected"
+							: "Browser disconnected"
+					);
+				}
 			}
 
 			let img = await env.BROWSER_KV_DEMO.get(url, { type: "arrayBuffer" });

fixtures/browser-rendering/test/index.spec.ts

@@ -62,6 +62,14 @@ describe("Local Browser", () => {
 					`New paragraph text set by ${lib === "playwright" ? "Playwright" : "Puppeteer"}!`
 				);
 			});
+
+			it("Disconnect a browser, and check its session connection status", async () => {
+				await expect(
+					fetchText(
+						`http://${ip}:${port}/?lib=${lib}&url=https://example.com&action=disconnect`
+					)
+				).resolves.toEqual(`Browser disconnected`);
+			});
 		});
 	}
 });

packages/miniflare/package.json

@@ -96,6 +96,7 @@
 		"npx-import": "^1.1.4",
 		"postal-mime": "^2.4.3",
 		"pretty-bytes": "^6.0.0",
+		"puppeteer": "22.8.2",
 		"rimraf": "catalog:default",
 		"source-map": "^0.6.1",
 		"ts-dedent": "^2.2.0",

packages/miniflare/src/index.ts

@@ -851,7 +851,7 @@ export class Miniflare {
 	#workerOpts: PluginWorkerOptions[];
 	#log: Log;
 
-	// key is the browser wsEndpoint, value is the browser process
+	// key is the browser session ID, value is the browser process
 	#browserProcesses: Map<string, ChildProcess> = new Map();
 
 	readonly #runtime?: Runtime;
@@ -1265,19 +1265,23 @@ export class Miniflare {
 					// workaround for CI environments, to avoid sandboxing issues
 					args: process.env.CI ? ["--no-sandbox"] : [],
 				});
+				const sessionId = crypto.randomUUID();
+				const browserProcess = browser.process();
+				browserProcess.on("exit", () => {
+					this.#browserProcesses.delete(sessionId);
+				});
 				const wsEndpoint = browser.wsEndpoint();
-				this.#browserProcesses.set(wsEndpoint, browser.process());
-				response = new Response(wsEndpoint);
+				const startTime = Date.now();
+				this.#browserProcesses.set(sessionId, browserProcess);
+				response = Response.json({ wsEndpoint, sessionId, startTime });
 			} else if (url.pathname === "/browser/status") {
-				const wsEndpoint = url.searchParams.get("wsEndpoint");
-				assert(wsEndpoint !== null, "Missing wsEndpoint query parameter");
-				const process = this.#browserProcesses.get(wsEndpoint);
-				const status = {
-					stopped: !process || process.exitCode !== null,
-				};
-				response = new Response(JSON.stringify(status), {
-					headers: { "Content-Type": "application/json" },
-				});
+				const sessionId = url.searchParams.get("sessionId");
+				assert(sessionId !== null, "Missing sessionId query parameter");
+				const process = this.#browserProcesses.get(sessionId);
+				response = new Response(null, { status: process ? 200 : 410 });
+			} else if (url.pathname === "/browser/sessionIds") {
+				const sessionIds = this.#browserProcesses.keys();
+				response = Response.json(Array.from(sessionIds));
 			} else if (url.pathname === "/core/store-temp-file") {
 				const prefix = url.searchParams.get("prefix");
 				const folder = prefix ? `files/${prefix}` : "files";
@@ -2001,10 +2005,10 @@ export class Miniflare {
 	}
 
 	async #assembleAndUpdateConfig() {
-		for (const [wsEndpoint, process] of this.#browserProcesses.entries()) {
+		for (const [sessionId, process] of this.#browserProcesses.entries()) {
 			// .close() isn't enough
 			process.kill("SIGKILL");
-			this.#browserProcesses.delete(wsEndpoint);
+			this.#browserProcesses.delete(sessionId);
 		}
 		// This function must be run with `#runtimeMutex` held
 		const initial = !this.#runtimeEntryURL;
@@ -2709,10 +2713,10 @@ export class Miniflare {
 		try {
 			await this.#waitForReady(/* disposing */ true);
 		} finally {
-			for (const [wsEndpoint, process] of this.#browserProcesses.entries()) {
+			for (const [sessionId, process] of this.#browserProcesses.entries()) {
 				// .close() isn't enough
 				process.kill("SIGKILL");
-				this.#browserProcesses.delete(wsEndpoint);
+				this.#browserProcesses.delete(sessionId);
 			}
 
 			// Remove exit hook, we're cleaning up what they would've cleaned up now

packages/miniflare/src/workers/browser-rendering/binding.worker.ts

@@ -12,23 +12,28 @@ function isClosed(ws: WebSocket | undefined): boolean {
 	return !ws || ws.readyState === WebSocket.CLOSED;
 }
 
+export type SessionInfo = {
+	wsEndpoint: string;
+	sessionId: string;
+	startTime: number;
+	connectionId?: string;
+	connectionStartTime?: number;
+};
+
 export class BrowserSession extends DurableObject<Env> {
-	endpoint?: string;
+	sessionInfo?: SessionInfo;
 	ws?: WebSocket;
 	server?: WebSocket;
 
 	async fetch(_request: Request) {
 		assert(
-			this.endpoint !== undefined,
-			"endpoint must be set before connecting"
+			this.sessionInfo !== undefined,
+			"sessionInfo must be set before connecting"
 		);
 
 		// sometimes the websocket doesn't get the close event, so we need to close them explicitly if needed
 		if (isClosed(this.ws) || isClosed(this.server)) {
-			this.ws?.close();
-			this.server?.close();
-			this.ws = undefined;
-			this.server = undefined;
+			this.closeWebSockets();
 		} else {
 			assert.fail("WebSocket already initialized");
 		}
@@ -38,7 +43,7 @@ export class BrowserSession extends DurableObject<Env> {
 
 		server.accept();
 
-		const wsEndpoint = this.endpoint.replace("ws://", "http://");
+		const wsEndpoint = this.sessionInfo.wsEndpoint.replace("ws://", "http://");
 
 		const response = await fetch(wsEndpoint, {
 			headers: {
@@ -85,37 +90,51 @@ export class BrowserSession extends DurableObject<Env> {
 		});
 		this.ws = ws;
 		this.server = server;
+		this.sessionInfo.connectionId = crypto.randomUUID();
+		this.sessionInfo.connectionStartTime = Date.now();
 
 		return new Response(null, {
 			status: 101,
 			webSocket: client,
 		});
 	}
-	async setEndpoint(endpoint: string) {
-		this.endpoint = endpoint;
+
+	async setSessionInfo(sessionInfo: SessionInfo) {
+		this.sessionInfo = sessionInfo;
+	}
+
+	async getSessionInfo(): Promise<SessionInfo | undefined> {
+		if (isClosed(this.ws) || isClosed(this.server)) {
+			this.closeWebSockets();
+		}
+		return this.sessionInfo;
 	}
 
 	async #checkStatus() {
-		if (this.endpoint) {
+		if (this.sessionInfo) {
 			const url = new URL("http://example.com/browser/status");
-			url.searchParams.set("wsEndpoint", this.endpoint);
+			url.searchParams.set("sessionId", this.sessionInfo.sessionId);
 			const resp = await this.env[CoreBindings.SERVICE_LOOPBACK].fetch(url);
-			const { stopped } = resp.ok
-				? ((await resp.json()) as { stopped: boolean })
-				: {};
 
-			if (stopped) {
+			if (!resp.ok) {
 				// Browser process has exited, we should close the WebSocket
 				// TODO should we send a error code?
-				this.ws?.close();
-				this.server?.close();
-				this.ws = undefined;
-				this.server = undefined;
-				this.ctx.storage.deleteAll();
+				this.closeWebSockets();
 				return;
 			}
 		}
 	}
+
+	closeWebSockets() {
+		this.ws?.close();
+		this.server?.close();
+		this.ws = undefined;
+		this.server = undefined;
+		if (this.sessionInfo) {
+			this.sessionInfo.connectionId = undefined;
+			this.sessionInfo.connectionStartTime = undefined;
+		}
+	}
 }
 
 export default {
@@ -126,18 +145,39 @@ export default {
 				const resp = await env[CoreBindings.SERVICE_LOOPBACK].fetch(
 					"http://example.com/browser/launch"
 				);
-				const wsEndpoint = await resp.text();
-				const sessionId = crypto.randomUUID();
-				const id = env.BrowserSession.idFromName(sessionId);
-				await env.BrowserSession.get(id).setEndpoint(wsEndpoint);
-				return Response.json({ sessionId });
+				const sessionInfo: SessionInfo = await resp.json();
+				const id = env.BrowserSession.idFromName(sessionInfo.sessionId);
+				await env.BrowserSession.get(id).setSessionInfo(sessionInfo);
+				return Response.json({ sessionId: sessionInfo.sessionId });
 			}
 			case "/v1/connectDevtools": {
 				const sessionId = url.searchParams.get("browser_session");
 				assert(sessionId !== null, "browser_session must be set");
 				const id = env.BrowserSession.idFromName(sessionId);
 				return env.BrowserSession.get(id).fetch(request);
 			}
+			case "/v1/sessions": {
+				const sessionIds = (await env[CoreBindings.SERVICE_LOOPBACK]
+					.fetch("http://example.com/browser/sessionIds")
+					.then((resp) => resp.json())) as string[];
+				const sessions = await Promise.all(
+					sessionIds.map(async (sessionId) => {
+						const id = env.BrowserSession.idFromName(sessionId);
+						return env.BrowserSession.get(id)
+							.getSessionInfo()
+							.then((sessionInfo) => {
+								if (!sessionInfo) return null;
+								return {
+									sessionId: sessionInfo.sessionId,
+									startTime: sessionInfo.startTime,
+									connectionId: sessionInfo.connectionId,
+									connectionStartTime: sessionInfo.connectionStartTime,
+								};
+							});
+					})
+				).then((results) => results.filter(Boolean));
+				return Response.json({ sessions });
+			}
 			default:
 				return new Response("Not implemented", { status: 405 });
 		}