fix(desktop): hand keyboard focus back after an automation click

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

@@ -180,20 +180,24 @@ const {
   browserWindowConstructor,
   clipboardItemConstructor,
   createFromPath,
+  fromFrame,
   fromId,
   getFocusedWebContents,
+  getFocusedWindow,
   mkdir,
   showItemInFolder,
   webviewSend,
   writeFile,
   writeClipboard,
 } = vi.hoisted(() => ({
   browserWindowConstructor: vi.fn(),
+  getFocusedWindow: vi.fn<() => Electron.BrowserWindow | null>(() => null),
   clipboardItemConstructor: vi.fn(),
   createFromPath: vi.fn((): { readonly isEmpty: () => boolean; readonly toPNG: () => Buffer } => ({
     isEmpty: () => false,
     toPNG: () => Buffer.from("png"),
   })),
+  fromFrame: vi.fn<typeof Electron.webContents.fromFrame>(() => undefined),
   fromId: vi.fn<(_id?: number) => Electron.WebContents | null>((_id?: number) => null),
   getFocusedWebContents: vi.fn(() => null),
   mkdir: vi.fn((_path: string) => undefined),
@@ -204,7 +208,7 @@ const {
 }));
 
 vi.mock("electron", () => ({
-  BrowserWindow: browserWindowConstructor,
+  BrowserWindow: Object.assign(browserWindowConstructor, { getFocusedWindow }),
   ClipboardItem: class {
     constructor(data: Record<string, unknown>) {
       clipboardItemConstructor(data);
@@ -223,6 +227,7 @@ vi.mock("electron", () => ({
     fromPartition: vi.fn(),
   },
   webContents: {
+    fromFrame,
     fromId,
     getFocusedWebContents,
   },
@@ -540,9 +545,12 @@ const makeTestPictureInPictureWindow = (loadURL: () => Promise<void> = async ()
 describe("PreviewManager", () => {
   beforeEach(() => {
     browserWindowConstructor.mockReset();
+    fromFrame.mockReset();
     fromId.mockClear();
     getFocusedWebContents.mockReset();
     getFocusedWebContents.mockReturnValue(null);
+    getFocusedWindow.mockReset();
+    getFocusedWindow.mockReturnValue({} as Electron.BrowserWindow);
     mkdir.mockClear();
     writeFile.mockClear();
     showItemInFolder.mockClear();
@@ -3922,6 +3930,155 @@ describe("PreviewManager", () => {
     ),
   );
 
+  for (const [scenario, name] of [
+    ["success", "hands keyboard focus back to the previous renderer after an automation click"],
+    ["failure", "restores keyboard focus and removes the hook after a failing automation click"],
+    ["third-renderer", "leaves a third renderer focused and removes the automation click hook"],
+    ["app-blurred", "does not restore keyboard focus after the app loses focus during a click"],
+    ["guest-focused", "does not preserve host focus when typing inside the preview"],
+    ["hook-rejected", "completes an automation click when the remember-focus hook rejects"],
+  ] as const) {
+    effectIt.effect(name, () =>
+      withManager((manager) =>
+        Effect.gen(function* () {
+          const composer = { isConnected: true, focus: vi.fn() };
+          const previewElement = {
+            getAttribute: (name: string) => (name === "data-preview-tab" ? "tab_1" : null),
+          };
+          const hostDocument = {
+            activeElement: composer as typeof composer | typeof previewElement,
+          };
+          const hostContext = NodeVM.createContext({ document: hostDocument });
+          const executeJavaScript = vi.fn(async (script: string) =>
+            NodeVM.runInContext(script, hostContext),
+          );
+          if (scenario === "hook-rejected") {
+            executeJavaScript.mockRejectedValueOnce(new Error("renderer rejected the focus hook"));
+          }
+          const hostFrame = { frameTreeNodeId: 7 };
+          const guestFrame = { frameTreeNodeId: 42 };
+          const thirdFrame = { frameTreeNodeId: 9 };
+          const host = Object.assign(makeTestHostWebContents(), {
+            focusedFrame: scenario === "guest-focused" ? guestFrame : hostFrame,
+            executeJavaScript,
+          });
+          const thirdRenderer = Object.assign(makeTestHostWebContents(), { id: 9 });
+          let humanInput: ((_event: unknown, signal: unknown) => void) | undefined;
+          const sendCommand = vi.fn(async (method: string, params?: Record<string, unknown>) => {
+            if (method === "Runtime.evaluate") {
+              return { result: { value: { width: 800, height: 600 } } };
+            }
+            if (method === "Input.dispatchMouseEvent" && params?.type === "mousePressed") {
+              host.focusedFrame = guestFrame;
+              hostDocument.activeElement = previewElement;
+              if (scenario === "third-renderer") host.focusedFrame = thirdFrame;
+              if (scenario === "app-blurred") {
+                getFocusedWebContents.mockReturnValue(null);
+                getFocusedWindow.mockReturnValue(null);
+              }
+              humanInput?.({}, { kind: "pointer", x: params.x, y: params.y, button: 0 });
+            }
+            return undefined;
+          });
+          const capturePage = vi.fn(async () => ({
+            toJPEG: () => Buffer.from("unused-focus-frame"),
+            getSize: () => ({ width: 800, height: 600 }),
+          }));
+          const guest = Object.assign(makeTestPreviewWebContents(capturePage, 42, host), {
+            isDevToolsOpened: () => false,
+            ipc: {
+              on: vi.fn((channel: string, listener: typeof humanInput) => {
+                if (channel === "preview:human-input") humanInput = listener;
+              }),
+              off: vi.fn(),
+            },
+            debugger: {
+              isAttached: () => false,
+              attach: vi.fn(),
+              sendCommand,
+              on: vi.fn(),
+              off: vi.fn(),
+            },
+          });
+          fromId.mockReturnValue(guest);
+          // Electron reports the guest even while the host composer has DOM focus.
+          getFocusedWebContents.mockReturnValue(guest as never);
+          fromFrame.mockImplementation((frame) => {
+            if (frame === hostFrame) return host as unknown as Electron.WebContents;
+            if (frame === guestFrame) return guest;
+            if (frame === thirdFrame) return thirdRenderer as unknown as Electron.WebContents;
+            return undefined;
+          });
+
+          yield* manager.createTab("tab_1");
+          yield* manager.registerWebview("tab_1", 42);
+          executeJavaScript.mockClear();
+          sendCommand.mockClear();
+          vi.mocked(guest.executeJavaScript).mockClear();
+          const click = yield* manager
+            .automationClick("tab_1", { x: scenario === "failure" ? 5000 : 120, y: 80 })
+            .pipe(Effect.exit, Effect.forkChild({ startImmediately: true }));
+          yield* TestClock.adjust(200);
+          const exit = yield* Fiber.join(click);
+          expect(exit).toMatchObject({ _tag: scenario === "failure" ? "Failure" : "Success" });
+          if (Exit.isFailure(exit)) {
+            expect(Option.getOrThrow(Cause.findErrorOption(exit.cause))).toMatchObject({
+              _tag: "PreviewAutomationCoordinatesOutsideViewportError",
+            });
+          }
+          expect(sendCommand).toHaveBeenCalled();
+          expect(
+            sendCommand.mock.calls.filter(([method]) => method === "Input.dispatchMouseEvent"),
+          ).toEqual(
+            scenario === "failure"
+              ? []
+              : [
+                  [
+                    "Input.dispatchMouseEvent",
+                    { type: "mousePressed", x: 120, y: 80, button: "left", clickCount: 1 },
+                  ],
+                  [
+                    "Input.dispatchMouseEvent",
+                    { type: "mouseReleased", x: 120, y: 80, button: "left", clickCount: 1 },
+                  ],
+                ],
+          );
+          expect(guest.executeJavaScript).not.toHaveBeenCalled();
+          expect(thirdRenderer.executeJavaScript).not.toHaveBeenCalled();
+          expect(Object.keys(hostContext)).toEqual(["document"]);
+
+          if (scenario === "guest-focused") {
+            expect(executeJavaScript).not.toHaveBeenCalled();
+            expect(composer.focus).not.toHaveBeenCalled();
+            return;
+          }
+          const shouldRestore = scenario !== "third-renderer" && scenario !== "app-blurred";
+          const scripts = executeJavaScript.mock.calls.map(([script]) => script);
+          const reference = scripts[0]?.match(/globalThis\[[^\]]+\]/)?.[0];
+          expect(reference).toBeDefined();
+          expect(scripts[0]).toContain("document.activeElement");
+          expect(scripts.slice(1)).toEqual(
+            shouldRestore ? [`${reference}?.()`, `delete ${reference}`] : [`delete ${reference}`],
+          );
+          expect(executeJavaScript.mock.invocationCallOrder[0]).toBeLessThan(
+            sendCommand.mock.invocationCallOrder[0]!,
+          );
+          expect(executeJavaScript.mock.invocationCallOrder[1]).toBeGreaterThan(
+            sendCommand.mock.invocationCallOrder.at(-1)!,
+          );
+          if (shouldRestore && scenario !== "hook-rejected") {
+            expect(composer.focus).toHaveBeenCalledExactlyOnceWith({ preventScroll: true });
+            expect(executeJavaScript.mock.invocationCallOrder[2]).toBeGreaterThan(
+              composer.focus.mock.invocationCallOrder[0]!,
+            );
+          } else {
+            expect(composer.focus).not.toHaveBeenCalled();
+          }
+        }),
+      ),
+    );
+  }
+
   effectIt.effect("types in background webviews and enables native key input", () =>
     withManager((manager) =>
       Effect.gen(function* () {

apps/desktop/src/preview/Manager.ts

@@ -3731,7 +3731,95 @@ const makeNativeOperations = Effect.fn("PreviewManager.makeOperations")(function
     );
   });
 
+  // A webview and its embedder can both report isFocused() on macOS. Use the
+  // embedder's focused frame to distinguish the guest from the host renderer.
+  const focusedRenderer = () => {
+    const focused = webContents.getFocusedWebContents();
+    const frame = (focused?.hostWebContents ?? focused)?.focusedFrame;
+    return (frame && webContents.fromFrame(frame)) || focused;
+  };
+
+  // Dispatching input moves keyboard focus into the guest renderer as a side
+  // effect. Hand it back to whatever had it before, so the user's next keystroke
+  // does not land in the previewed page, which may not even be visible.
+  const restoreFocusedWebContents = Effect.fn("PreviewManager.restoreFocusedWebContents")(
+    function* (
+      operation: string,
+      tabId: string,
+      wc: Electron.WebContents,
+      previouslyFocused: Electron.WebContents | null,
+      focusKey: string,
+    ) {
+      if (!previouslyFocused || previouslyFocused.id === wc.id || previouslyFocused.isDestroyed()) {
+        return;
+      }
+      // A newer selection the user made while the action ran wins over the restore.
+      const focusedNow = yield* attempt({ operation, tabId, webContentsId: wc.id }, () =>
+        focusedRenderer(),
+      ).pipe(Effect.orElseSucceed(() => null));
+      if (focusedNow && focusedNow.id !== wc.id && focusedNow.id !== previouslyFocused.id) {
+        return;
+      }
+      // The user left T3 for another app while the action ran; do not pull them back.
+      if (focusedNow === null && BrowserWindow.getFocusedWindow() === null) {
+        return;
+      }
+      yield* attemptPromise({ operation, tabId, webContentsId: previouslyFocused.id }, () =>
+        previouslyFocused.executeJavaScript(`globalThis[${focusKey}]?.()`),
+      ).pipe(Effect.ignore);
+    },
+  );
+
   const performAutomationClick = Effect.fn("PreviewManager.performAutomationClick")(function* (
+    tabId: string,
+    wc: Electron.WebContents,
+    input: PreviewAutomationClickInput,
+    send: SendCommand,
+  ) {
+    const previouslyFocused = yield* attempt(
+      { operation: "automationClick.getFocusedWebContents", tabId, webContentsId: wc.id },
+      () => focusedRenderer(),
+    );
+    const context = { operation: "automationClick.preserveFocus", tabId, webContentsId: wc.id };
+    const focusKey = yield* encodeJson(context, `__t3AutomationFocus_${NodeCrypto.randomUUID()}`);
+    const previewTab = yield* encodeJson(context, tabId);
+    // Native window focus does not restore the embedder's DOM focus after a
+    // webview click. Keep the element reference in its renderer.
+    if (previouslyFocused && previouslyFocused.id !== wc.id) {
+      yield* attemptPromise(context, () =>
+        previouslyFocused.executeJavaScript(`(() => {
+          const element = document.activeElement;
+          globalThis[${focusKey}] = () => {
+            const active = document.activeElement;
+            if (element?.isConnected && (active === element ||
+                active?.getAttribute("data-preview-tab") === ${previewTab})) {
+              element.focus({ preventScroll: true });
+            }
+          };
+        })()`),
+      ).pipe(Effect.ignore);
+    }
+    yield* dispatchAutomationClick(tabId, input, send).pipe(
+      Effect.ensuring(
+        restoreFocusedWebContents(
+          "automationClick.restoreFocusedWebContents",
+          tabId,
+          wc,
+          previouslyFocused,
+          focusKey,
+        ),
+      ),
+      Effect.ensuring(
+        previouslyFocused && previouslyFocused.id !== wc.id
+          ? attemptPromise(context, () =>
+              previouslyFocused.executeJavaScript(`delete globalThis[${focusKey}]`),
+            ).pipe(Effect.ignore)
+          : Effect.void,
+      ),
+    );
+  });
+
+  const dispatchAutomationClick = Effect.fn("PreviewManager.dispatchAutomationClick")(function* (
     tabId: string,
     input: PreviewAutomationClickInput,
     send: SendCommand,
@@ -3794,7 +3882,7 @@ const makeNativeOperations = Effect.fn("PreviewManager.makeOperations")(function
   ) {
     const wc = yield* requireWebContents(tabId);
     yield* withControlSession(tabId, wc, "click", (send) =>
-      performAutomationClick(tabId, input, send),
+      performAutomationClick(tabId, wc, input, send),
     );
   });