fix(browser): __piloRefMap fallback when data-pilo-ref is stripped

packages/core/src/browser/playwrightBrowser.ts

@@ -773,9 +773,7 @@ export class PlaywrightBrowser implements AriaBrowser {
     const locator = this.page.locator(`[data-pilo-ref="${ref}"]`);
     const count = await locator.count();
 
-    if (count === 0) {
-      throw new InvalidRefException(ref);
-    }
+    if (count === 1) return locator;
 
     if (count > 1) {
       // This shouldn't happen with data-pilo-ref, but let's be defensive
@@ -785,7 +783,41 @@ export class PlaywrightBrowser implements AriaBrowser {
       );
     }
 
-    return locator;
+    // count === 0: the attribute is gone (commonly stripped by React reconciliation
+    // after our snapshot set it). Fall back to the JS-side __piloRefMap which holds
+    // direct Element references that survive attribute strips. We re-attach the
+    // attribute so the subsequent locator query (and any future lookup in this
+    // iteration) finds the element via the existing selector path.
+    //
+    // The ownerDocument check restricts recovery to the main frame's document.
+    // Same-origin iframe elements may be in __piloRefMap (the snapshot walks them
+    // inline) but cannot be located by page.locator(...) — that searches the main
+    // frame only. Recovering them would push the failure to click time as an
+    // opaque locator timeout instead of an immediate InvalidRefException.
+    const reattached = await this.page.evaluate((r) => {
+      const map = (globalThis as { __piloRefMap?: Map<string, Element> }).__piloRefMap;
+      const el = map?.get(r);
+      if (!el || !el.isConnected || el.ownerDocument !== document) return false;
+      el.setAttribute("data-pilo-ref", r);
+      return true;
+    }, ref);
+
+    if (!reattached) throw new InvalidRefException(ref);
+
+    // Re-validate the post-recovery locator with the same count() check the
+    // happy path uses. Guards against a race where another reconciliation
+    // pass strips the attribute again, or where setAttribute somehow lands
+    // on multiple elements — both turn into the same loud failure as today.
+    const recoveredLocator = this.page.locator(`[data-pilo-ref="${ref}"]`);
+    const recoveredCount = await recoveredLocator.count();
+    if (recoveredCount === 1) return recoveredLocator;
+    if (recoveredCount > 1) {
+      throw new InvalidRefException(
+        ref,
+        `Multiple elements found with reference '${ref}'. This may indicate a page structure issue.`,
+      );
+    }
+    throw new InvalidRefException(ref);
   }
 
   async performAction(ref: string, action: PageAction, value?: string): Promise<void> {

packages/core/test/playwrightBrowser.test.ts

@@ -603,12 +603,14 @@ describe("PlaywrightBrowser", () => {
 
     describe("validateElementRef", () => {
       it("should throw InvalidRefException when element doesn't exist", async () => {
-        // Mock the page and locator
+        // Mock the page and locator. evaluate returns false so the __piloRefMap
+        // fallback also misses (the simulated DOM has no element for this ref).
         const mockLocator = {
           count: vi.fn().mockResolvedValue(0),
         };
         const mockPage = {
           locator: vi.fn().mockReturnValue(mockLocator),
+          evaluate: vi.fn().mockResolvedValue(false),
         };
         (browser as any).page = mockPage;
 
@@ -650,6 +652,89 @@ describe("PlaywrightBrowser", () => {
         expect(result).toBe(mockLocator);
         expect(mockPage.locator).toHaveBeenCalledWith('[data-pilo-ref="valid"]');
       });
+
+      it("recovers via __piloRefMap when attribute is stripped but element still connected", async () => {
+        let attributeSet = false;
+        const mockLocator0 = { count: vi.fn().mockResolvedValue(0) };
+        const mockLocator1 = { count: vi.fn().mockResolvedValue(1) };
+        const mockPage = {
+          locator: vi.fn().mockImplementation(() => (attributeSet ? mockLocator1 : mockLocator0)),
+          evaluate: vi.fn().mockImplementation(async () => {
+            attributeSet = true;
+            return true;
+          }),
+        };
+        (browser as any).page = mockPage;
+
+        const result = await (browser as any).validateElementRef("stripped");
+        expect(result).toBe(mockLocator1);
+        expect(mockPage.locator).toHaveBeenCalledTimes(2);
+        expect(mockPage.evaluate).toHaveBeenCalledTimes(1);
+      });
+
+      it("throws InvalidRefException when __piloRefMap has no entry for the ref", async () => {
+        const mockLocator = { count: vi.fn().mockResolvedValue(0) };
+        const mockPage = {
+          locator: vi.fn().mockReturnValue(mockLocator),
+          evaluate: vi.fn().mockResolvedValue(false),
+        };
+        (browser as any).page = mockPage;
+
+        await expect((browser as any).validateElementRef("hallucinated")).rejects.toThrow(
+          InvalidRefException,
+        );
+        expect(mockPage.evaluate).toHaveBeenCalledTimes(1);
+      });
+
+      it("throws InvalidRefException when __piloRefMap entry is no longer connected to DOM", async () => {
+        const mockLocator = { count: vi.fn().mockResolvedValue(0) };
+        const mockPage = {
+          locator: vi.fn().mockReturnValue(mockLocator),
+          // isConnected check fails inside evaluate, so the page-side callback returns false
+          evaluate: vi.fn().mockResolvedValue(false),
+        };
+        (browser as any).page = mockPage;
+
+        await expect((browser as any).validateElementRef("removed")).rejects.toThrow(
+          InvalidRefException,
+        );
+      });
+
+      it("throws InvalidRefException when post-recovery locator still matches zero (race)", async () => {
+        // Simulates a reconciliation race: evaluate reports success, but by the
+        // time we re-query the locator, the attribute has been stripped again.
+        const mockLocator = { count: vi.fn().mockResolvedValue(0) };
+        const mockPage = {
+          locator: vi.fn().mockReturnValue(mockLocator),
+          evaluate: vi.fn().mockResolvedValue(true),
+        };
+        (browser as any).page = mockPage;
+
+        await expect((browser as any).validateElementRef("racy")).rejects.toThrow(
+          InvalidRefException,
+        );
+        // First locator returns 0 → triggers fallback. Second locator (post-evaluate)
+        // also returns 0 → throws.
+        expect(mockPage.locator).toHaveBeenCalledTimes(2);
+      });
+
+      it("throws InvalidRefException when post-recovery locator matches multiple elements", async () => {
+        let attributeSet = false;
+        const mockLocator0 = { count: vi.fn().mockResolvedValue(0) };
+        const mockLocator2 = { count: vi.fn().mockResolvedValue(2) };
+        const mockPage = {
+          locator: vi.fn().mockImplementation(() => (attributeSet ? mockLocator2 : mockLocator0)),
+          evaluate: vi.fn().mockImplementation(async () => {
+            attributeSet = true;
+            return true;
+          }),
+        };
+        (browser as any).page = mockPage;
+
+        await expect((browser as any).validateElementRef("dup")).rejects.toThrow(
+          "Multiple elements found with reference 'dup'",
+        );
+      });
     });
 
     describe("actionRequiresElement", () => {
@@ -803,6 +888,8 @@ describe("PlaywrightBrowser", () => {
         };
         const mockPage = {
           locator: vi.fn().mockReturnValue(mockLocator),
+          // evaluate returns false so the __piloRefMap fallback also misses
+          evaluate: vi.fn().mockResolvedValue(false),
         };
         (browser as any).page = mockPage;
 

packages/extension/src/background/ExtensionBrowser.ts

@@ -354,7 +354,22 @@ export class ExtensionBrowser implements AriaBrowser {
           const { ref: refParam, action: actionParam, value: valueParam } = JSON.parse(paramsJson);
 
           // Look up the element using the data-pilo-ref attribute that's now set by ariaSnapshot
-          const element = document.querySelector(`[data-pilo-ref="${refParam}"]`);
+          let element = document.querySelector(`[data-pilo-ref="${refParam}"]`);
+
+          if (!element) {
+            // Fallback: __piloRefMap holds direct Element references that survive
+            // attribute strips (commonly caused by React reconciliation). If the
+            // mapped element is still connected to the DOM of this frame, re-attach
+            // the attribute and use it. The ownerDocument check excludes same-origin
+            // iframe elements that the snapshot may have walked inline — those
+            // can't be acted on from this script's main-document context.
+            const map = (globalThis as { __piloRefMap?: Map<string, Element> }).__piloRefMap;
+            const mapped = map?.get(refParam);
+            if (mapped && mapped.isConnected && mapped.ownerDocument === document) {
+              mapped.setAttribute("data-pilo-ref", refParam);
+              element = mapped;
+            }
+          }
 
           if (!element) {
             return {