fix(highlight): highlight in the world the selector resolves in

packages/playwright-core/src/server/highlightController.ts

@@ -14,8 +14,9 @@
  * limitations under the License.
  */
 
-import type { Frame } from './frames';
+import type { FrameExecutionContext } from './dom';
 import type { Page } from './page';
+import type * as types from './types';
 import type { ParsedSelector } from '@isomorphic/selectorParser';
 
 export type HighlightOptions = {
@@ -27,6 +28,9 @@ type HighlightEntry = HighlightOptions & {
   selector: string;
 };
 
+// Custom selector engines run in the main world, so highlights may live in either world.
+const worlds: types.World[] = ['utility', 'main'];
+
 export class HighlightController {
   private _page: Page;
   private _entries = new Map<string, HighlightEntry>();
@@ -59,9 +63,10 @@ export class HighlightController {
   async hideHighlights() {
     this._entries.clear();
     await Promise.all(this._page.frames().map(frame => frame.raceAgainstEvaluationStallingEvents(async () => {
-      const context = frame.existingContext('utility');
-      const injectedScript = await context?.injectedScript();
-      await injectedScript?.evaluate(injected => injected.hideHighlight());
+      await Promise.all(worlds.map(async world => {
+        const injectedScript = await frame.existingContext(world)?.injectedScript();
+        await injectedScript?.evaluate(injected => injected.hideHighlight());
+      }));
     }).catch(() => {})));
   }
 
@@ -78,25 +83,32 @@ export class HighlightController {
     if (this._page.isClosed())
       return;
 
-    const perFrame = new Map<Frame, { selector: ParsedSelector, cssStyle?: string }[]>();
+    const perContext = new Map<FrameExecutionContext, { selector: ParsedSelector, cssStyle?: string }[]>();
     for (const entry of this._entries.values()) {
       const results = await this._page.mainFrame().selectors.resolveFramesForSelector(entry.selector, { strict: false, anyFrame: entry.anyFrame }).catch(() => []);
       for (const { frame, info } of results) {
-        let list = perFrame.get(frame);
+        const context = frame.existingContext(info.world);
+        if (!context)
+          continue;
+        let list = perContext.get(context);
         if (!list) {
           list = [];
-          perFrame.set(frame, list);
+          perContext.set(context, list);
         }
         list.push({ selector: info.parsed, cssStyle: entry.style });
       }
     }
 
     await Promise.all(this._page.frames().map(async frame => {
-      const highlights = perFrame.get(frame) || [];
       await frame.raceAgainstEvaluationStallingEvents(async () => {
-        const context = frame.existingContext('utility');
-        const injectedScript = await context?.injectedScript();
-        await injectedScript?.evaluate((injected, highlights) => injected.setHighlights(highlights), highlights);
+        await Promise.all(worlds.map(async world => {
+          const context = frame.existingContext(world);
+          if (!context)
+            return;
+          const highlights = perContext.get(context) || [];
+          const injectedScript = await context.injectedScript();
+          await injectedScript.evaluate((injected, highlights) => injected.setHighlights(highlights), highlights);
+        }));
       }).catch(() => {});
     }));
 

tests/library/locator-highlight.spec.ts

@@ -118,3 +118,36 @@ test('Page.hideHighlight clears all locator highlights', async ({ browser, serve
 
   await context.close();
 });
+
+test('highlight should work with a custom selector engine that runs in the main world', async ({ playwright, browser }) => {
+  // Engines registered without "contentScript" run in the main world and can see page globals.
+  const createEngine = () => ({
+    query(root, selector) {
+      return window['__engineEnabled'] ? root.querySelector(selector) : null;
+    },
+    queryAll(root, selector) {
+      return window['__engineEnabled'] ? Array.from(root.querySelectorAll(selector)) : [];
+    },
+  });
+  await playwright.selectors.register('highlight-tag', createEngine);
+
+  const context = await browser.newContext();
+  const page = await context.newPage();
+  await page.setContent('<button>Button</button>');
+  await page.evaluate(() => window['__engineEnabled'] = true);
+
+  const button = page.locator('highlight-tag=button');
+  await button.highlight();
+  await expect(page.locator('x-pw-highlight')).toHaveCount(1);
+
+  await button.hideHighlight();
+  await expect(page.locator('x-pw-highlight')).toHaveCount(0);
+
+  await button.highlight();
+  await expect(page.locator('x-pw-highlight')).toHaveCount(1);
+
+  await page.hideHighlight();
+  await expect(page.locator('x-pw-highlight')).toHaveCount(0);
+
+  await context.close();
+});