Memoize Playwright page target IDs

.github/workflows/server-test.yaml

@@ -36,6 +36,11 @@ jobs:
           cache: true
           cache-dependency-path: server/go.sum
 
+      - name: Set up Node.js
+        uses: actions/setup-node@v4
+        with:
+          node-version: 22
+
       # categorygen's checks (unclassified route, category that isn't control or
       # platform, classified route with no handler) only run when the generator
       # does, and its only other caller is `make oapi-generate`, which needs the
@@ -50,6 +55,10 @@ jobs:
         run: make test-unit
         working-directory: server
 
+      - name: Run Playwright daemon unit tests
+        run: make test-runtime
+        working-directory: server
+
   test-server-e2e:
     runs-on: ubuntu-latest
     needs: [build-headful, build-headless]

images/chromium-headful/Dockerfile

@@ -372,7 +372,7 @@ COPY --from=server-builder /out/kernel-images-supervisord-shim /usr/local/bin/ke
 COPY --from=server-builder /out/wrapper /wrapper
 
 # Copy and compile the Playwright daemon
-COPY server/runtime/playwright-daemon.ts /tmp/playwright-daemon.ts
+COPY server/runtime/playwright-daemon.ts server/runtime/page-target-id-cache.ts /tmp/
 RUN esbuild /tmp/playwright-daemon.ts \
     --bundle \
     --platform=node \
@@ -382,7 +382,7 @@ RUN esbuild /tmp/playwright-daemon.ts \
     --external:playwright-core \
     --external:patchright \
     --external:esbuild \
-    && rm /tmp/playwright-daemon.ts
+    && rm /tmp/playwright-daemon.ts /tmp/page-target-id-cache.ts
 
 RUN useradd -m -s /bin/bash kernel
 

images/chromium-headless/image/Dockerfile

@@ -268,7 +268,7 @@ COPY --from=server-builder /out/chromium-launcher /usr/local/bin/chromium-launch
 COPY --from=server-builder /out/kernel-images-supervisord-shim /usr/local/bin/kernel-images-supervisord-shim
 
 # Copy and compile the Playwright daemon
-COPY server/runtime/playwright-daemon.ts /tmp/playwright-daemon.ts
+COPY server/runtime/playwright-daemon.ts server/runtime/page-target-id-cache.ts /tmp/
 RUN esbuild /tmp/playwright-daemon.ts \
     --bundle \
     --platform=node \
@@ -278,6 +278,6 @@ RUN esbuild /tmp/playwright-daemon.ts \
     --external:playwright-core \
     --external:patchright \
     --external:esbuild \
-    && rm /tmp/playwright-daemon.ts
+    && rm /tmp/playwright-daemon.ts /tmp/page-target-id-cache.ts
 
 ENTRYPOINT [ "/wrapper" ]

server/Makefile

@@ -1,5 +1,5 @@
 SHELL := /bin/bash
-.PHONY: oapi-generate build dev test test-unit test-e2e clean
+.PHONY: oapi-generate build dev test test-unit test-runtime test-e2e clean
 
 BIN_DIR ?= $(CURDIR)/bin
 RECORDING_DIR ?= $(CURDIR)/recordings
@@ -30,15 +30,18 @@ build: | $(BIN_DIR)
 dev: build $(RECORDING_DIR)
 	OUTPUT_DIR=$(RECORDING_DIR) DISPLAY_NUM=$(DISPLAY_NUM) ./bin/api
 
-# `test` runs unit + e2e. The two are split so callers (e.g. the Hypeman CI job)
-# can run just the e2e suite, and so e2e logs stream as they run instead of
-# waiting for all unit tests to complete.
-test: test-unit test-e2e
+# `test` runs Go unit, runtime unit, and e2e tests. The suites are split so
+# callers (e.g. the Hypeman CI job) can run just the e2e suite, and so e2e logs
+# stream as they run instead of waiting for all unit tests to complete.
+test: test-unit test-runtime test-e2e
 
 test-unit:
 	go vet ./...
 	go test -v -race $$(go list ./... | grep -v /e2e$$)
 
+test-runtime:
+	node --test runtime/*.test.ts
+
 test-e2e:
 	@echo ""
 	@echo "=== Running e2e tests (this may take a few minutes) ==="

server/runtime/page-target-id-cache.test.ts

@@ -0,0 +1,157 @@
+import assert from 'node:assert/strict';
+import test from 'node:test';
+
+import { PageTargetIdCache } from './page-target-id-cache.ts';
+
+test('reuses a successfully discovered target ID', async () => {
+  const page = {};
+  let discoveries = 0;
+  const cache = new PageTargetIdCache<object>(async () => {
+    discoveries++;
+    return 'target-1';
+  });
+
+  assert.equal(await cache.get(page), 'target-1');
+  assert.equal(await cache.get(page), 'target-1');
+  assert.equal(discoveries, 1);
+});
+
+test('refreshes a cached target ID', async () => {
+  const page = {};
+  let targetId = 'stale-target';
+  let discoveries = 0;
+  const cache = new PageTargetIdCache<object>(async () => {
+    discoveries++;
+    return targetId;
+  });
+
+  assert.equal(await cache.get(page), 'stale-target');
+  targetId = 'fresh-target';
+  assert.equal(await cache.get(page), 'stale-target');
+  assert.equal(await cache.get(page, { refresh: true }), 'fresh-target');
+  assert.equal(await cache.get(page), 'fresh-target');
+  assert.equal(discoveries, 2);
+});
+
+test('does not cache failed discovery', async () => {
+  const page = {};
+  let discoveries = 0;
+  const cache = new PageTargetIdCache<object>(async () => {
+    discoveries++;
+    if (discoveries === 1) throw new Error('page closed');
+    return 'target-1';
+  });
+
+  await assert.rejects(cache.get(page), /page closed/);
+  assert.equal(await cache.get(page), 'target-1');
+  assert.equal(discoveries, 2);
+});
+
+test('shares an in-flight discovery between concurrent callers', async () => {
+  const page = {};
+  const pending = Promise.withResolvers<string>();
+  let discoveries = 0;
+  const cache = new PageTargetIdCache<object>(() => {
+    discoveries++;
+    return pending.promise;
+  });
+
+  const first = cache.get(page);
+  const second = cache.get(page);
+  assert.equal(discoveries, 1);
+
+  pending.resolve('target-1');
+  assert.deepEqual(await Promise.all([first, second]), ['target-1', 'target-1']);
+  assert.equal(discoveries, 1);
+});
+
+test('failed refresh evicts the stale target ID', async () => {
+  const page = {};
+  const results = ['stale-target', new Error('target replaced'), 'fresh-target'];
+  let discoveries = 0;
+  const cache = new PageTargetIdCache<object>(async () => {
+    const result = results[discoveries++];
+    if (result instanceof Error) throw result;
+    return result;
+  });
+
+  assert.equal(await cache.get(page), 'stale-target');
+  await assert.rejects(cache.get(page, { refresh: true }), /target replaced/);
+  assert.equal(await cache.get(page), 'fresh-target');
+  assert.equal(discoveries, 3);
+});
+
+test('builds an index while skipping pages that cannot be inspected', async () => {
+  const firstPage = {};
+  const closedPage = {};
+  const secondPage = {};
+  const targetIds = new Map<object, string>([
+    [firstPage, 'target-1'],
+    [secondPage, 'target-2'],
+  ]);
+  const cache = new PageTargetIdCache<object>(async page => {
+    const targetId = targetIds.get(page);
+    if (targetId === undefined) throw new Error('page closed');
+    return targetId;
+  });
+
+  const pageByTargetId = await cache.buildPageByTargetId([firstPage, closedPage, secondPage]);
+
+  assert.deepEqual([...pageByTargetId.entries()], [
+    ['target-1', firstPage],
+    ['target-2', secondPage],
+  ]);
+});
+
+test('refreshes cached IDs while rebuilding the index', async () => {
+  const page = {};
+  let targetId = 'stale-target';
+  let discoveries = 0;
+  const cache = new PageTargetIdCache<object>(async () => {
+    discoveries++;
+    return targetId;
+  });
+
+  assert.equal((await cache.buildPageByTargetId([page])).get('stale-target'), page);
+  targetId = 'fresh-target';
+  const refreshed = await cache.buildPageByTargetId([page], { refresh: true });
+
+  assert.equal(refreshed.has('stale-target'), false);
+  assert.equal(refreshed.get('fresh-target'), page);
+  assert.equal(discoveries, 2);
+});
+
+test('reset discards every cached target ID', async () => {
+  const firstPage = {};
+  const secondPage = {};
+  let discoveries = 0;
+  const cache = new PageTargetIdCache<object>(async () => `target-${++discoveries}`);
+
+  assert.equal(await cache.get(firstPage), 'target-1');
+  assert.equal(await cache.get(secondPage), 'target-2');
+  cache.reset();
+  assert.equal(await cache.get(firstPage), 'target-3');
+  assert.equal(await cache.get(secondPage), 'target-4');
+});
+
+test('reset prevents an in-flight discovery from repopulating the cache', async () => {
+  const page = {};
+  const pending = Promise.withResolvers<string>();
+  let discoveries = 0;
+  const cache = new PageTargetIdCache<object>(async () => {
+    discoveries++;
+    if (discoveries === 1) return pending.promise;
+    return 'fresh-target';
+  });
+
+  const staleDiscovery = cache.get(page);
+  cache.reset();
+  const freshDiscovery = cache.get(page);
+  assert.equal(discoveries, 2);
+
+  pending.resolve('stale-target');
+  assert.equal(await staleDiscovery, 'stale-target');
+  assert.equal(await freshDiscovery, 'fresh-target');
+  assert.equal(await cache.get(page), 'fresh-target');
+  assert.equal(discoveries, 2);
+});

server/runtime/page-target-id-cache.ts

@@ -0,0 +1,67 @@
+interface CacheOptions {
+  refresh?: boolean;
+}
+
+// A CDP page target ID is stable for the lifetime of its Playwright Page.
+// Weak keys let closed pages and their IDs be collected without explicit
+// eviction; refresh and reset cover target replacement and browser reconnects.
+export class PageTargetIdCache<Page extends object> {
+  #targetIdMemo = new WeakMap<Page, string>();
+  #targetIdInFlight = new WeakMap<Page, Promise<string>>();
+  #generation = 0;
+  readonly #discoverTargetId: (page: Page) => Promise<string>;
+
+  constructor(discoverTargetId: (page: Page) => Promise<string>) {
+    this.#discoverTargetId = discoverTargetId;
+  }
+
+  async get(page: Page, options: CacheOptions = {}): Promise<string> {
+    if (!options.refresh) {
+      const cached = this.#targetIdMemo.get(page);
+      if (cached !== undefined) return cached;
+    } else {
+      this.#targetIdMemo.delete(page);
+    }
+
+    const inFlight = this.#targetIdInFlight.get(page);
+    if (inFlight !== undefined) return inFlight;
+
+    const generation = this.#generation;
+    const discovery = this.#discoverTargetId(page)
+      .then(targetId => {
+        if (generation === this.#generation) {
+          this.#targetIdMemo.set(page, targetId);
+        }
+        return targetId;
+      })
+      .finally(() => {
+        if (this.#targetIdInFlight.get(page) === discovery) {
+          this.#targetIdInFlight.delete(page);
+        }
+      });
+
+    this.#targetIdInFlight.set(page, discovery);
+    return discovery;
+  }
+
+  async buildPageByTargetId(pages: readonly Page[], options: CacheOptions = {}): Promise<Map<string, Page>> {
+    const pageByTargetId = new Map<string, Page>();
+
+    for (const page of pages) {
+      try {
+        pageByTargetId.set(await this.get(page, options), page);
+      } catch {
+        // A crashed or closing page can fail target discovery. Exclude it from
+        // this snapshot without preventing other live pages from resolving.
+      }
+    }
+
+    return pageByTargetId;
+  }
+
+  reset(): void {
+    this.#targetIdMemo = new WeakMap<Page, string>();
+    this.#targetIdInFlight = new WeakMap<Page, Promise<string>>();
+    this.#generation++;
+  }
+}