fix(web): restore renderer focus after preview presses

apps/web/src/components/preview/PreviewAutomationHosts.tsx

@@ -45,6 +45,7 @@ import { browserDefaultOpenViewport, resolveBrowserDefaults } from "~/browser/br
 import { runBrowserViewportMutation } from "~/browser/browserViewportActions";
 import { previewRuntimeTabId } from "~/browser/previewRuntimeTabId";
 import { isElectron } from "~/env";
+import { withPreviewAutomationFocus } from "~/lib/previewAutomationFocus";
 import { useEnvironments } from "~/state/environments";
 import { previewEnvironment } from "~/state/preview";
 import { useAtomQueryRunner } from "~/state/use-atom-query-runner";
@@ -652,11 +653,13 @@ function PreviewAutomationHost(props: { readonly environmentId: EnvironmentId })
             );
           }
           case "press": {
-            const ready = await requireReadyTab();
-            return await ready.bridge.automation.press(
-              ready.runtimeTabId,
-              request.input as Parameters<typeof ready.bridge.automation.press>[1],
-            );
+            return await withPreviewAutomationFocus(async () => {
+              const ready = await requireReadyTab();
+              return await ready.bridge.automation.press(
+                ready.runtimeTabId,
+                request.input as Parameters<typeof ready.bridge.automation.press>[1],
+              );
+            });
           }
           case "scroll": {
             const ready = await requireReadyTab();

apps/web/src/lib/previewAutomationFocus.test.ts

@@ -0,0 +1,226 @@
+import { afterEach, describe, expect, it, vi } from "vite-plus/test";
+
+import { withPreviewAutomationFocus } from "./previewAutomationFocus";
+
+class MockHTMLElement {
+  isConnected = true;
+  readonly focus = vi.fn((_options?: FocusOptions) => {
+    setActiveElement(this);
+  });
+}
+
+const setActiveElement = (activeElement: MockHTMLElement | null): void => {
+  (globalThis.document as unknown as { activeElement: MockHTMLElement | null }).activeElement =
+    activeElement;
+};
+
+afterEach(() => {
+  vi.unstubAllGlobals();
+});
+
+const setupDocument = (activeElement: MockHTMLElement | null, focused = true) => {
+  const body = new MockHTMLElement();
+  const documentElement = new MockHTMLElement();
+  let documentFocused = focused;
+  const documentListeners = new Map<string, Set<(event: Event) => void>>();
+  const windowListeners = new Map<string, Set<() => void>>();
+  vi.stubGlobal("HTMLElement", MockHTMLElement);
+  vi.stubGlobal("document", {
+    activeElement,
+    body,
+    documentElement,
+    hasFocus: () => documentFocused,
+    addEventListener: (type: string, listener: (event: Event) => void) => {
+      const listeners = documentListeners.get(type) ?? new Set();
+      listeners.add(listener);
+      documentListeners.set(type, listeners);
+    },
+    removeEventListener: (type: string, listener: (event: Event) => void) => {
+      documentListeners.get(type)?.delete(listener);
+    },
+  });
+  vi.stubGlobal("window", {
+    addEventListener: (type: string, listener: () => void) => {
+      const listeners = windowListeners.get(type) ?? new Set();
+      listeners.add(listener);
+      windowListeners.set(type, listeners);
+    },
+    removeEventListener: (type: string, listener: () => void) => {
+      windowListeners.get(type)?.delete(listener);
+    },
+  });
+  return {
+    body,
+    setDocumentFocused: (value: boolean) => {
+      documentFocused = value;
+    },
+    dispatchDocument: (type: string, target: MockHTMLElement, isTrusted = true) => {
+      for (const listener of documentListeners.get(type) ?? []) {
+        listener({ target, isTrusted } as unknown as Event);
+      }
+    },
+    dispatchWindow: (type: string) => {
+      for (const listener of windowListeners.get(type) ?? []) listener();
+    },
+  };
+};
+
+describe("withPreviewAutomationFocus", () => {
+  it("restores focus when automation leaves a connected host control focused", async () => {
+    const composer = new MockHTMLElement();
+    const hostButton = new MockHTMLElement();
+    const { dispatchDocument, dispatchWindow } = setupDocument(composer);
+
+    const result = await withPreviewAutomationFocus(async () => {
+      // Native guest focus briefly transfers the renderer window away and back.
+      dispatchWindow("blur");
+      dispatchWindow("focus");
+      setActiveElement(hostButton);
+      dispatchDocument("focusin", hostButton, false);
+      return "pressed";
+    });
+
+    expect(result).toBe("pressed");
+    expect(composer.focus).toHaveBeenCalledWith({ preventScroll: true });
+    expect(globalThis.document.activeElement).toBe(composer);
+  });
+
+  it("preserves newer DOM focus while the bridge operation is pending", async () => {
+    const composer = new MockHTMLElement();
+    const newerControl = new MockHTMLElement();
+    const { body, dispatchDocument } = setupDocument(composer);
+    let finish!: () => void;
+    let started!: () => void;
+    const operationStarted = new Promise<void>((resolve) => {
+      started = resolve;
+    });
+
+    const pending = withPreviewAutomationFocus(async () => {
+      setActiveElement(body);
+      started();
+      await new Promise<void>((resolve) => {
+        finish = resolve;
+      });
+    });
+
+    await operationStarted;
+    setActiveElement(newerControl);
+    // This models a newer programmatic or user focus event in the host.
+    dispatchDocument("focusin", newerControl);
+    finish();
+    await pending;
+
+    expect(composer.focus).not.toHaveBeenCalled();
+    expect(globalThis.document.activeElement).toBe(newerControl);
+  });
+
+  it("does not restore a detached prior element", async () => {
+    const detachedComposer = new MockHTMLElement();
+    const { body } = setupDocument(detachedComposer);
+    await withPreviewAutomationFocus(async () => {
+      detachedComposer.isConnected = false;
+      setActiveElement(body);
+    });
+    expect(detachedComposer.focus).not.toHaveBeenCalled();
+  });
+
+  it("does not restore when the document is unfocused at invocation", async () => {
+    const unfocusedComposer = new MockHTMLElement();
+    const unfocused = setupDocument(unfocusedComposer, false);
+    setActiveElement(unfocusedComposer);
+    await withPreviewAutomationFocus(async () => {
+      setActiveElement(unfocused.body);
+    });
+    expect(unfocusedComposer.focus).not.toHaveBeenCalled();
+  });
+
+  it("does not restore when the document loses focus during the operation", async () => {
+    const composer = new MockHTMLElement();
+    const { body, setDocumentFocused } = setupDocument(composer);
+    await withPreviewAutomationFocus(async () => {
+      setActiveElement(body);
+      setDocumentFocused(false);
+    });
+    expect(composer.focus).not.toHaveBeenCalled();
+  });
+
+  it.each(["pointerdown", "keydown"] as const)(
+    "preserves user focus after a native transfer and %s",
+    async (userEvent) => {
+      const composer = new MockHTMLElement();
+      const hostButton = new MockHTMLElement();
+      const { dispatchDocument, dispatchWindow } = setupDocument(composer);
+
+      await withPreviewAutomationFocus(async () => {
+        dispatchWindow("blur");
+        dispatchWindow("focus");
+        dispatchDocument(userEvent, hostButton);
+        setActiveElement(hostButton);
+        dispatchDocument("focusin", hostButton);
+      });
+
+      expect(composer.focus).not.toHaveBeenCalled();
+      expect(globalThis.document.activeElement).toBe(hostButton);
+    },
+  );
+
+  it("does not mask the operation rejection when restoration fails", async () => {
+    const composer = new MockHTMLElement();
+    const { body } = setupDocument(composer);
+    const error = new Error("press failed");
+    composer.focus.mockImplementation(() => {
+      throw new Error("focus failed");
+    });
+
+    await expect(
+      withPreviewAutomationFocus(async () => {
+        setActiveElement(body);
+        throw error;
+      }),
+    ).rejects.toBe(error);
+  });
+
+  it("does not mask the operation result when restoration fails", async () => {
+    const composer = new MockHTMLElement();
+    const { body } = setupDocument(composer);
+    composer.focus.mockImplementation(() => {
+      throw new Error("focus failed");
+    });
+
+    await expect(
+      withPreviewAutomationFocus(async () => {
+        setActiveElement(body);
+        return "pressed";
+      }),
+    ).resolves.toBe("pressed");
+  });
+
+  it("does not let an older overlapping operation reclaim focus", async () => {
+    const composer = new MockHTMLElement();
+    const { body } = setupDocument(composer);
+    let finish!: () => void;
+    let firstStarted!: () => void;
+    const firstIsStarted = new Promise<void>((resolve) => {
+      firstStarted = resolve;
+    });
+
+    const first = withPreviewAutomationFocus(async () => {
+      setActiveElement(body);
+      firstStarted();
+      await new Promise<void>((resolve) => {
+        finish = resolve;
+      });
+    });
+    await firstIsStarted;
+
+    const second = withPreviewAutomationFocus(async () => {
+      setActiveElement(body);
+    });
+
+    await second;
+    expect(composer.focus).not.toHaveBeenCalled();
+    finish();
+    await first;
+    expect(composer.focus).not.toHaveBeenCalled();
+  });
+});

apps/web/src/lib/previewAutomationFocus.ts

@@ -0,0 +1,108 @@
+let latestOperationId = 0;
+
+const getMeaningfulActiveElement = (): HTMLElement | null => {
+  if (typeof document === "undefined" || typeof HTMLElement === "undefined") return null;
+
+  const activeElement = document.activeElement;
+  if (
+    !(activeElement instanceof HTMLElement) ||
+    !activeElement.isConnected ||
+    activeElement === document.body ||
+    activeElement === document.documentElement
+  ) {
+    return null;
+  }
+  return activeElement;
+};
+
+const isDocumentFocused = (): boolean => {
+  if (typeof document === "undefined") return false;
+  return typeof document.hasFocus !== "function" || document.hasFocus();
+};
+
+/**
+ * Keeps preview automation from changing focus in the shared renderer.
+ */
+export async function withPreviewAutomationFocus<T>(operation: () => Promise<T>): Promise<T> {
+  const operationId = ++latestOperationId;
+  const previouslyFocused = getMeaningfulActiveElement();
+  const wasDocumentFocused = isDocumentFocused();
+  let userFocusObserved = false;
+  let pendingUserFocus = false;
+  let pendingVersion = 0;
+  let nativeFocusElement: HTMLElement | null = null;
+  let windowBlurred = false;
+  let windowRefocused = false;
+
+  const markPendingUserFocus = (event: Event): void => {
+    if (!event.isTrusted) return;
+    pendingUserFocus = true;
+    const version = ++pendingVersion;
+    queueMicrotask(() => {
+      if (pendingVersion === version) pendingUserFocus = false;
+    });
+  };
+  const onFocusIn = (event: Event): void => {
+    const nativeFocusTransfer = windowBlurred && windowRefocused;
+    const target = event.target instanceof HTMLElement ? event.target : null;
+    if (nativeFocusTransfer) nativeFocusElement = target;
+    if (event.isTrusted && (pendingUserFocus || (!nativeFocusTransfer && target?.isConnected))) {
+      userFocusObserved = true;
+    }
+    pendingUserFocus = false;
+    pendingVersion += 1;
+    windowBlurred = false;
+    windowRefocused = false;
+  };
+  const onWindowBlur = (): void => {
+    windowBlurred = true;
+    windowRefocused = false;
+  };
+  const onWindowFocus = (): void => {
+    if (windowBlurred) windowRefocused = true;
+  };
+
+  if (typeof document !== "undefined") {
+    document.addEventListener("pointerdown", markPendingUserFocus, true);
+    document.addEventListener("keydown", markPendingUserFocus, true);
+    document.addEventListener("focusin", onFocusIn, true);
+  }
+  if (typeof window !== "undefined") {
+    window.addEventListener("blur", onWindowBlur);
+    window.addEventListener("focus", onWindowFocus);
+  }
+
+  try {
+    return await operation();
+  } finally {
+    if (typeof document !== "undefined") {
+      document.removeEventListener("pointerdown", markPendingUserFocus, true);
+      document.removeEventListener("keydown", markPendingUserFocus, true);
+      document.removeEventListener("focusin", onFocusIn, true);
+    }
+    if (typeof window !== "undefined") {
+      window.removeEventListener("blur", onWindowBlur);
+      window.removeEventListener("focus", onWindowFocus);
+    }
+
+    const activeElement = getMeaningfulActiveElement();
+    const activeFocusIsExpected =
+      !activeElement || activeElement === previouslyFocused || activeElement === nativeFocusElement;
+    if (
+      operationId === latestOperationId &&
+      !userFocusObserved &&
+      wasDocumentFocused &&
+      !windowBlurred &&
+      isDocumentFocused() &&
+      previouslyFocused?.isConnected &&
+      activeFocusIsExpected &&
+      activeElement !== previouslyFocused
+    ) {
+      try {
+        previouslyFocused.focus({ preventScroll: true });
+      } catch {
+        // Focus restoration is best effort; never mask the automation result.
+      }
+    }
+  }
+}