fix(desktop): a hidden preview tab no longer keeps the display awake after a screenshot

apps/desktop/src/preview/DesktopBrowserHost.test.ts

@@ -607,20 +607,21 @@ describe("DesktopBrowserHost", () => {
         params: { sessionId: "t3-preview-page-1" },
       });
       expect(yield* Queue.take(replies)).toBe(3);
-      expect(painting).toEqual([true, false]);
+      expect(painting).toEqual([true]);
 
       // The agent's screenshot waits its turn; painting lets the abandoned one finish first.
       yield* command(host, 4, "Page.captureScreenshot");
       yield* settle;
-      expect(painting).toEqual([true, false, true]);
+      expect(painting).toEqual([true]);
       debuggee.release();
       expect(yield* Queue.take(replies)).toBe(2);
       yield* settle;
       expect(debuggee.sent).toEqual([undefined, 4]);
+      expect(painting).toEqual([true]);
       debuggee.release();
       expect(yield* Queue.take(replies)).toBe(4);
       yield* settle;
-      expect(painting).toEqual([true, false, true, false]);
+      expect(painting).toEqual([true, false]);
     }).pipe(Effect.scoped),
   );
 
@@ -645,14 +646,19 @@ describe("DesktopBrowserHost", () => {
       expect(yield* Queue.take(replies)).toBe(2);
       yield* raw({ id: 3, method: "Page.captureScreenshot", sessionId: "t3-preview-page-1" });
       yield* settle;
-      expect(painting).toEqual([true, false, true]);
+      expect(painting).toEqual([true]);
       yield* raw({
         id: 4,
         method: "Target.detachFromTarget",
         params: { sessionId: "t3-preview-page-1" },
       });
       expect(yield* Queue.take(replies)).toBe(4);
-      expect(painting).toEqual([true, false, true, false]);
+
+      // The agent's answers late; the viewer's is dropped unsent, so nothing is left to paint.
+      debuggee.release();
+      yield* settle;
+      expect(debuggee.sent).toEqual([1]);
+      expect(painting).toEqual([true, false]);
       yield* TestClock.adjust("1 minute");
       expect(debuggee.focusEmulation).toEqual([true, false]);
     }).pipe(Effect.scoped),
@@ -682,7 +688,7 @@ describe("DesktopBrowserHost", () => {
     }).pipe(Effect.scoped),
   );
 
-  it.effect("stops painting for a screenshot still queued when the agent's turn ends", () =>
+  it.effect("holds no focus for a screenshot still queued when the agent's turn ends", () =>
     Effect.gen(function* () {
       const host = yield* DesktopBrowserHost.make;
       const debuggee = makeDebuggee();
@@ -697,15 +703,18 @@ describe("DesktopBrowserHost", () => {
       yield* settle;
       yield* turn(false);
       yield* settle;
-      expect(painting).toEqual([true, false]);
+      expect(painting).toEqual([true]);
 
-      // The first answers late; the second still goes in turn, but holds nothing.
+      // The first answers late; the second still goes in turn, painting but holding no focus.
       debuggee.releaseFirst();
       yield* settle;
       expect(debuggee.sent).toEqual([1, 2]);
-      expect(painting).toEqual([true, false]);
       yield* TestClock.adjust("1 minute");
       expect(debuggee.focusEmulation).toEqual([true, false]);
+      expect(painting).toEqual([true]);
+      debuggee.release();
+      yield* settle;
+      expect(painting).toEqual([true, false]);
     }).pipe(Effect.scoped),
   );
 
@@ -952,7 +961,7 @@ describe("DesktopBrowserHost", () => {
     }).pipe(Effect.scoped),
   );
 
-  it.effect("stops painting for a screenshot Chromium never answers", () =>
+  it.effect("stops painting for a screenshot Chromium never answers once its document goes", () =>
     Effect.gen(function* () {
       const host = yield* DesktopBrowserHost.make;
       const debuggee = makeDebuggee();
@@ -961,32 +970,34 @@ describe("DesktopBrowserHost", () => {
       const turn = (active: boolean) =>
         host.handleCommandLine(encodeJson({ type: "turn", threadId: key.threadId, active }));
 
-      // The server gave up on it, and the agent's turn ends.
+      // The server gave up on it, and the agent's turn ends; Chromium still holds the capture.
       yield* turn(true);
       yield* command(host, 1, "Page.captureScreenshot");
       yield* settle;
       yield* turn(false);
       yield* settle;
-      expect(painting).toEqual([true, false]);
+      expect(painting).toEqual([true]);
 
       // The page loads a new document, which also frees the screenshot turn.
       debuggee.navigate();
+      expect(painting).toEqual([true, false]);
       yield* command(host, 2, "Page.captureScreenshot");
       yield* settle;
       debuggee.navigate();
       yield* settle;
       expect(painting).toEqual([true, false, true, false]);
 
-      // The server releases the relay.
+      // The server releases the relay; Chromium still holds the capture.
       yield* command(host, 3, "Page.captureScreenshot");
       yield* settle;
       yield* host.handleCommandLine(encodeJson({ type: "release", ...key }));
       yield* settle;
-      expect(painting).toEqual([true, false, true, false, true, false]);
+      expect(painting).toEqual([true, false, true, false, true]);
 
       // The guest detaches; a replacement starts with no screenshot pending.
       const replacement = makeDebuggee();
       host.attach(key, replacement.tab);
+      expect(painting).toEqual([true, false, true, false, true, false]);
       yield* command(host, 4, "Page.captureScreenshot");
       yield* settle;
       host.detach(key);

apps/desktop/src/preview/DesktopBrowserHost.ts

@@ -48,7 +48,7 @@ interface AttachedTab {
   readonly debuggee: DesktopBrowserTabDebugger;
   relay: CdpRelayConnection | null;
   screenshotTurn: object | null;
-  /** Each holds focus and painting while `held` and whoever queued it still waits on it. */
+  /** Each keeps the guest painting while whoever queued it still waits, and focused while `held`. */
   readonly waitingScreenshots: Array<{
     readonly start: () => void;
     readonly live: () => boolean;
@@ -61,10 +61,9 @@ interface AttachedTab {
   focusChange: Promise<void> | null;
   focusUses: number;
   /** Each pending command, by the extra session waiting on it; the main connection's have none. */
-  readonly focusPending: Map<
-    object,
-    { readonly owner: string | undefined; readonly paints: boolean }
-  >;
+  readonly focusPending: Map<object, string | undefined>;
+  /** Screenshots Chromium was sent and has not answered; each keeps its wake lock until then. */
+  readonly capturing: Set<object>;
   /** Whether the client keeps the guest painting, as published on `painting`. */
   painting: boolean;
   /**
@@ -148,7 +147,7 @@ export const make = Effect.gen(function* () {
     tab: AttachedTab,
     relay: CdpRelayConnection,
     ownerLive: () => boolean,
-    /** Sends only while `current` holds; holds focus and painting only while `held`. */
+    /** Sends only while `current` holds; holds focus only while `held`. */
     capture: (current: () => boolean, held: boolean) => Promise<unknown>,
   ) =>
     new Promise<unknown>((resolve, reject) => {
@@ -236,38 +235,42 @@ export const make = Effect.gen(function* () {
     const active =
       tabs.get(keyOf(tab.key)) === tab &&
       (tab.screencasts.size > 0 ||
-        waitingHolds(tab) ||
-        [...tab.focusPending.values()].some((hold) => hold.paints));
+        tab.capturing.size > 0 ||
+        tab.waitingScreenshots.some((waiting) => waiting.live()));
     if (tab.painting === active) return;
     tab.painting = active;
     runFork(PubSub.publish(painting, { key: tab.key, active }));
   };
 
   /** A command may wait on the page's next frame, so the guest stays focused until it settles. */
-  const holdFocus = <A>(
-    tab: AttachedTab,
-    method: string,
-    owner: string | undefined,
-    work: () => Promise<A>,
-  ) => {
+  const holdFocus = <A>(tab: AttachedTab, owner: string | undefined, work: () => Promise<A>) => {
     const pending = {};
-    tab.focusPending.set(pending, { owner, paints: method === "Page.captureScreenshot" });
+    tab.focusPending.set(pending, owner);
     tab.focusUses += 1;
     setFocusEmulation(tab, true);
     releaseFocusWhenIdle(tab);
-    syncPainting(tab);
     return work().finally(() => {
       if (!tab.focusPending.delete(pending)) return;
       tab.focusUses += 1;
       releaseFocusWhenIdle(tab);
+    });
+  };
+
+  /** Chromium answers a capture, and frees its wake lock, only once the guest paints a frame. */
+  const paintUntilAnswered = <A>(tab: AttachedTab, send: () => Promise<A>) => {
+    const sent = {};
+    tab.capturing.add(sent);
+    syncPainting(tab);
+    return send().finally(() => {
+      tab.capturing.delete(sent);
       syncPainting(tab);
     });
   };
 
   /**
    * Chromium can leave a command unanswered, so the screenshot guard's release points
-   * also end every hold: a released relay has no one waiting, and a new document answers
-   * nothing pending from the old one. A released relay's screencast has no viewer to stop it.
+   * also end every focus hold: a released relay has no one waiting, and a new document
+   * answers nothing pending from the old one. A released relay's screencast has no viewer.
    */
   const dropFocusHolds = (tab: AttachedTab) => {
     tab.focusPending.clear();
@@ -285,24 +288,24 @@ export const make = Effect.gen(function* () {
   /** A detached session, such as a closed viewer, waits on nothing it sent. */
   const dropSessionHolds = (tab: AttachedTab, owner: string) => {
     tab.screencasts.delete(owner);
-    for (const [pending, hold] of tab.focusPending) {
-      if (hold.owner === owner) tab.focusPending.delete(pending);
+    for (const [pending, holder] of tab.focusPending) {
+      if (holder === owner) tab.focusPending.delete(pending);
     }
     releaseFocusWhenIdle(tab);
     syncPainting(tab);
   };
 
   /**
    * A command the server stopped waiting on, such as a timed-out screenshot or evaluation,
-   * may never settle; the turn's end ends every hold. The agent also sends on its own extra
+   * may never settle; the turn's end ends every focus hold. The agent also sends on its own extra
    * session, so this cannot spare other sessions; a viewer's screencast still holds.
    */
   const endAgentTurn = (threadId: string) => {
     agentTurns.delete(threadId);
     for (const tab of tabs.values()) {
       if (tab.key.threadId !== threadId) continue;
       tab.agentUsed = false;
-      // Still sent in turn, so the queue keeps moving, but no longer kept painting.
+      // Still sent in turn, so the queue keeps moving, but no longer kept focused.
       for (const waiting of tab.waitingScreenshots) waiting.held = false;
       dropFocusHolds(tab);
     }
@@ -323,8 +326,10 @@ export const make = Effect.gen(function* () {
           if (sessionId === undefined && method === "Page.startScreencast") {
             tab.screencasts.add(owner);
             if (owner !== undefined) streamingOwners.add(owner);
+            syncPainting(tab);
           } else if (sessionId === undefined && method === "Page.stopScreencast") {
             tab.screencasts.delete(owner);
+            syncPainting(tab);
           }
           // A person's input through a viewer, or Playwright's own housekeeping, is not the agent's.
           if (
@@ -337,9 +342,11 @@ export const make = Effect.gen(function* () {
           const forward = (current?: () => boolean) => {
             const change = tab.focusChange;
             const run = () =>
-              tab.relay === relay && (current === undefined || current())
-                ? send(method, params, sessionId)
-                : Promise.reject(new Error("The page this screenshot was for has moved on."));
+              tab.relay !== relay || (current !== undefined && !current())
+                ? Promise.reject(new Error("The page this screenshot was for has moved on."))
+                : current === undefined
+                  ? send(method, params, sessionId)
+                  : paintUntilAnswered(tab, () => send(method, params, sessionId));
             return change === null ? run() : change.then(run);
           };
           return method === "Page.captureScreenshot"
@@ -348,17 +355,17 @@ export const make = Effect.gen(function* () {
                 relay,
                 () => owner === undefined || !detachedOwners.has(owner),
                 (current, held) =>
-                  held ? holdFocus(tab, method, owner, () => forward(current)) : forward(current),
+                  held ? holdFocus(tab, owner, () => forward(current)) : forward(current),
               )
             : method === "Page.startScreencast"
-              ? holdFocus(tab, method, owner, forward).catch((cause: unknown) => {
+              ? holdFocus(tab, owner, forward).catch((cause: unknown) => {
                   // A viewer can outlive a refused stream; nothing is streaming for it.
                   if (sessionId === undefined && tab.relay === relay) tab.screencasts.delete(owner);
                   releaseFocusWhenIdle(tab);
                   syncPainting(tab);
                   throw cause;
                 })
-              : holdFocus(tab, method, owner, forward);
+              : holdFocus(tab, owner, forward);
         },
         sessionDetached: (owner) => {
           detachedOwners.add(owner);
@@ -445,6 +452,7 @@ export const make = Effect.gen(function* () {
       focusChange: null,
       focusUses: 0,
       focusPending: new Map(),
+      capturing: new Set(),
       painting: false,
       agentUsed: false,
       focusReleasing: false,
@@ -459,7 +467,8 @@ export const make = Effect.gen(function* () {
       // A screenshot pending when the page loads a new document, a reload included, can
       // stay unanswered. A crashed page lands here once it reloads, well after its renderer died.
       onNewDocument: () => {
-        // Before the next screenshot starts, so its own hold survives.
+        // Before the next screenshot starts, so its own holds survive.
+        tab.capturing.clear();
         dropFocusHolds(tab);
         if (tab.screenshotTurn !== null) nextScreenshot(tab);
       },