From a65702754222cb5373bffb92f0a47c11c4aecb0e Mon Sep 17 00:00:00 2001
From: Nidhi Nair <nidhi.nair93@gmail.com>
Date: Fri, 30 Jan 2026 17:00:46 +0530
Subject: [PATCH] fix: Use frames as reference for all activity

---
 src/actions.ts      |  5 ++-
 src/browser.test.ts | 75 ++++++++++++++++++++++++++++++++++++++++++++-
 src/browser.ts      | 21 +++++++------
 src/snapshot.ts     | 13 +++++---
 4 files changed, 96 insertions(+), 18 deletions(-)

diff --git a/src/actions.ts b/src/actions.ts
index 6162b48d..cccd6b03 100644
--- a/src/actions.ts
+++ b/src/actions.ts
@@ -619,10 +619,9 @@ async function handleEvaluate(
   command: EvaluateCommand,
   browser: BrowserManager
 ): Promise<Response<EvaluateData>> {
-  const page = browser.getPage();
+  const frame = browser.getFrame();
 
-  // Evaluate the script directly as a string expression
-  const result = await page.evaluate(command.script);
+  const result = await frame.evaluate(command.script);
 
   return successResponse(command.id, { result });
 }
diff --git a/src/browser.test.ts b/src/browser.test.ts
index a145d28a..5b024292 100644
--- a/src/browser.test.ts
+++ b/src/browser.test.ts
@@ -1,4 +1,4 @@
-import { describe, it, expect, beforeAll, afterAll, vi } from 'vitest';
+import { describe, it, expect, beforeAll, afterAll, vi, beforeEach } from 'vitest';
 import { BrowserManager } from './browser.js';
 import { chromium } from 'playwright-core';
 
@@ -661,4 +661,77 @@ describe('BrowserManager', () => {
       ).resolves.not.toThrow();
     });
   });
+
+  describe('frame switch', () => {
+    const IFRAME_HTML = `
+      <div id="main-content">Main document</div>
+      <iframe id="test-iframe" srcdoc="<h1>Iframe heading</h1><button id='frame-btn'>Frame button</button>"></iframe>
+    `;
+
+    beforeEach(async () => {
+      const page = browser.getPage();
+      await page.setContent(IFRAME_HTML, { waitUntil: 'load' });
+      browser.switchToMainFrame();
+    });
+
+    it('should return main frame when no frame selected', () => {
+      browser.switchToMainFrame();
+      const frame = browser.getFrame();
+      const page = browser.getPage();
+      expect(frame).toBe(page.mainFrame());
+    });
+
+    it('should switch to iframe by selector and snapshot shows iframe content', async () => {
+      await browser.switchToFrame({ selector: '#test-iframe' });
+      const { tree } = await browser.getSnapshot();
+      expect(tree).toContain('Iframe heading');
+      expect(tree).toContain('Frame button');
+      expect(tree).not.toContain('Main document');
+    });
+
+    it('should run eval in active frame context', async () => {
+      await browser.switchToFrame({ selector: '#test-iframe' });
+      const frame = browser.getFrame();
+      const result = await frame.evaluate(() => document.querySelector('h1')?.textContent);
+      expect(result).toBe('Iframe heading');
+    });
+
+    it('should resolve locators within active frame', async () => {
+      await browser.switchToFrame({ selector: '#test-iframe' });
+      const locator = browser.getLocator('#frame-btn');
+      const text = await locator.textContent();
+      expect(text).toBe('Frame button');
+    });
+
+    it('should support nested iframe switch (selector within current frame)', async () => {
+      const nestedHtml = `
+        <div>Outer iframe</div>
+        <iframe id="inner" srcdoc="<p>Inner content</p>"></iframe>
+      `;
+      const page = browser.getPage();
+      await page.setContent(
+        `<iframe id="outer" srcdoc="${nestedHtml.replace(/"/g, '&quot;')}"></iframe>`,
+        { waitUntil: 'load' }
+      );
+
+      await browser.switchToFrame({ selector: '#outer' });
+      const { tree } = await browser.getSnapshot();
+      expect(tree).toContain('Outer iframe');
+
+      await browser.switchToFrame({ selector: '#inner' });
+      const { tree: innerTree } = await browser.getSnapshot();
+      expect(innerTree).toContain('Inner content');
+    });
+
+    it('should return to main document after switchToMainFrame', async () => {
+      await browser.switchToFrame({ selector: '#test-iframe' });
+      const { tree: iframeTree } = await browser.getSnapshot();
+      expect(iframeTree).toContain('Iframe heading');
+
+      browser.switchToMainFrame();
+      const { tree: mainTree } = await browser.getSnapshot();
+      expect(mainTree).toContain('Main document');
+      expect(mainTree).not.toContain('Iframe heading');
+    });
+  });
 });
diff --git a/src/browser.ts b/src/browser.ts
index 9ce57567..9e564fdb 100644
--- a/src/browser.ts
+++ b/src/browser.ts
@@ -119,8 +119,8 @@ export class BrowserManager {
     compact?: boolean;
     selector?: string;
   }): Promise<EnhancedSnapshot> {
-    const page = this.getPage();
-    const snapshot = await getEnhancedSnapshot(page, options);
+    const frame = this.getFrame();
+    const snapshot = await getEnhancedSnapshot(frame, options);
     this.refMap = snapshot.refs;
     this.lastSnapshot = snapshot.tree;
     return snapshot;
@@ -144,14 +144,14 @@ export class BrowserManager {
     const refData = this.refMap[ref];
     if (!refData) return null;
 
-    const page = this.getPage();
+    const frame = this.getFrame();
 
     // Build locator with exact: true to avoid substring matches
     let locator: Locator;
     if (refData.name) {
-      locator = page.getByRole(refData.role as any, { name: refData.name, exact: true });
+      locator = frame.getByRole(refData.role as any, { name: refData.name, exact: true });
     } else {
-      locator = page.getByRole(refData.role as any);
+      locator = frame.getByRole(refData.role as any);
     }
 
     // If an nth index is stored (for disambiguation), use it
@@ -178,8 +178,8 @@ export class BrowserManager {
     if (locator) return locator;
 
     // Otherwise treat as regular selector
-    const page = this.getPage();
-    return page.locator(selectorOrRef);
+    const frame = this.getFrame();
+    return frame.locator(selectorOrRef);
   }
 
   /**
@@ -203,13 +203,16 @@ export class BrowserManager {
   }
 
   /**
-   * Switch to a frame by selector, name, or URL
+   * Switch to a frame by selector, name, or URL.
+   * Selector is resolved within the current frame (supports nested iframes).
+   * Name and URL search the entire frame tree.
    */
   async switchToFrame(options: { selector?: string; name?: string; url?: string }): Promise<void> {
     const page = this.getPage();
+    const currentFrame = this.getFrame();
 
     if (options.selector) {
-      const frameElement = await page.$(options.selector);
+      const frameElement = await currentFrame.$(options.selector);
       if (!frameElement) {
         throw new Error(`Frame not found: ${options.selector}`);
       }
diff --git a/src/snapshot.ts b/src/snapshot.ts
index de81aa90..11943942 100644
--- a/src/snapshot.ts
+++ b/src/snapshot.ts
@@ -17,7 +17,7 @@
  *   agent-browser click @e2             # Click element by ref
  */
 
-import type { Page, Locator } from 'playwright-core';
+import type { Page, Frame, Locator } from 'playwright-core';
 
 export interface RefMap {
   [ref: string]: {
@@ -137,17 +137,20 @@ function buildSelector(role: string, name?: string): string {
 }
 
 /**
- * Get enhanced snapshot with refs and optional filtering
+ * Get enhanced snapshot with refs and optional filtering.
+ * Accepts Page or Frame so snapshots can be scoped to the active frame (e.g. iframe).
  */
 export async function getEnhancedSnapshot(
-  page: Page,
+  pageOrFrame: Page | Frame,
   options: SnapshotOptions = {}
 ): Promise<EnhancedSnapshot> {
   resetRefs();
   const refs: RefMap = {};
 
-  // Get ARIA snapshot from Playwright
-  const locator = options.selector ? page.locator(options.selector) : page.locator(':root');
+  // Get ARIA snapshot from Playwright (works for both Page and Frame)
+  const locator = options.selector
+    ? pageOrFrame.locator(options.selector)
+    : pageOrFrame.locator(':root');
   const ariaTree = await locator.ariaSnapshot();
 
   if (!ariaTree) {