refactor: move dialog handling to McpPage (#1059)

- renamed getSelectedPage to getSelectedPptrPage
- removes getSelectedPptrPage from tool interfaces
- moves dialog handling to McpPage
- makes responses to be optionally McpPage-scoped

Author: OrKoN

src/McpContext.ts

@@ -23,7 +23,6 @@ import type {
   BrowserContext,
   ConsoleMessage,
   Debugger,
-  Dialog,
   ElementHandle,
   HTTPRequest,
   Page,
@@ -205,7 +204,7 @@ export class McpContext implements Context {
   }
 
   #resolveTargetPage(): Page {
-    return this.#requestPage?.pptrPage ?? this.getSelectedPage();
+    return this.#requestPage?.pptrPage ?? this.getSelectedPptrPage();
   }
 
   resolveCdpRequestId(cdpRequestId: string): number | undefined {
@@ -327,7 +326,7 @@ export class McpContext implements Context {
   }
 
   async restoreEmulation(targetPage?: Page) {
-    const page = targetPage ?? this.getSelectedPage();
+    const page = targetPage ?? this.getSelectedPptrPage();
     const mcpPage = this.#getMcpPage(page);
     const currentSetting = mcpPage.emulationSettings;
     await this.emulate(currentSetting, targetPage);
@@ -344,7 +343,7 @@ export class McpContext implements Context {
     },
     targetPage?: Page,
   ): Promise<void> {
-    const page = targetPage ?? this.getSelectedPage();
+    const page = targetPage ?? this.getSelectedPptrPage();
     const mcpPage = this.#getMcpPage(page);
     const newSettings: EmulationSettings = {...mcpPage.emulationSettings};
     let timeoutsNeedUpdate = false;
@@ -492,23 +491,7 @@ export class McpContext implements Context {
     return this.#options.performanceCrux;
   }
 
-  getDialog(page?: Page): Dialog | undefined {
-    const targetPage =
-      page ?? this.#requestPage?.pptrPage ?? this.#selectedPage;
-    if (!targetPage) {
-      return undefined;
-    }
-    return this.#mcpPages.get(targetPage)?.dialog;
-  }
-
-  clearDialog(page?: Page): void {
-    const targetPage = page ?? this.#selectedPage;
-    if (targetPage) {
-      this.#mcpPages.get(targetPage)?.clearDialog();
-    }
-  }
-
-  getSelectedPage(): Page {
+  getSelectedPptrPage(): Page {
     const page = this.#selectedPage;
     if (!page) {
       throw new Error('No page selected');
@@ -522,7 +505,7 @@ export class McpContext implements Context {
   }
 
   getSelectedMcpPage(): McpPage {
-    const page = this.getSelectedPage();
+    const page = this.getSelectedPptrPage();
     return this.#getMcpPage(page);
   }
 
@@ -555,7 +538,7 @@ export class McpContext implements Context {
   }
 
   #getSelectedMcpPage(): McpPage {
-    return this.#getMcpPage(this.getSelectedPage());
+    return this.#getMcpPage(this.getSelectedPptrPage());
   }
 
   isPageSelected(page: Page): boolean {
@@ -595,7 +578,7 @@ export class McpContext implements Context {
   }
 
   #updateSelectedPageTimeouts() {
-    const page = this.getSelectedPage();
+    const page = this.getSelectedPptrPage();
     // For waiters 5sec timeout should be sufficient.
     // Increased in case we throttle the CPU
     const cpuMultiplier = this.getCpuThrottlingRate();
@@ -922,7 +905,7 @@ export class McpContext implements Context {
     devtoolsData: DevToolsData | undefined = undefined,
     targetPage?: Page,
   ): Promise<void> {
-    const page = targetPage ?? this.getSelectedPage();
+    const page = targetPage ?? this.getSelectedPptrPage();
     const mcpPage = this.#getMcpPage(page);
     const rootNode = await page.accessibility.snapshot({
       includeIframes: true,
@@ -1063,7 +1046,7 @@ export class McpContext implements Context {
     action: () => Promise<unknown>,
     options?: {timeout?: number},
   ): Promise<void> {
-    const page = this.getSelectedPage();
+    const page = this.getSelectedPptrPage();
     const cpuMultiplier = this.getCpuThrottlingRate();
     const networkMultiplier = getNetworkMultiplierFromString(
       this.getNetworkConditions(),
@@ -1085,7 +1068,7 @@ export class McpContext implements Context {
     timeout?: number,
     targetPage?: Page,
   ): Promise<Element> {
-    const page = targetPage ?? this.getSelectedPage();
+    const page = targetPage ?? this.getSelectedPptrPage();
     const frames = page.frames();
 
     let locator = this.#locatorClass.race(

src/McpPage.ts

@@ -59,6 +59,10 @@ export class McpPage implements ContextPage {
     return this.#dialog;
   }
 
+  getDialog(): Dialog | undefined {
+    return this.dialog;
+  }
+
   clearDialog(): void {
     this.#dialog = undefined;
   }

src/McpResponse.ts

@@ -10,6 +10,7 @@ import {IssueFormatter} from './formatters/IssueFormatter.js';
 import {NetworkFormatter} from './formatters/NetworkFormatter.js';
 import {SnapshotFormatter} from './formatters/SnapshotFormatter.js';
 import type {McpContext} from './McpContext.js';
+import type {McpPage} from './McpPage.js';
 import {UncaughtError} from './PageCollector.js';
 import {DevTools} from './third_party/index.js';
 import type {
@@ -70,11 +71,16 @@ export class McpResponse implements Response {
   #devToolsData?: DevToolsData;
   #tabId?: string;
   #args: ParsedArguments;
+  #page?: McpPage;
 
   constructor(args: ParsedArguments) {
     this.#args = args;
   }
 
+  setPage(page: McpPage): void {
+    this.#page = page;
+  }
+
   attachDevToolsData(data: DevToolsData): void {
     this.#devToolsData = data;
   }
@@ -517,7 +523,7 @@ export class McpResponse implements Response {
       structuredContent.colorScheme = colorScheme;
     }
 
-    const dialog = context.getDialog();
+    const dialog = this.#page?.getDialog();
     if (dialog) {
       const defaultValueIfNeeded =
         dialog.type() === 'prompt'

src/server.ts

@@ -153,7 +153,8 @@ export async function createMcpServer(
     const schema =
       'pageScoped' in tool &&
       tool.pageScoped &&
-      serverArgs.experimentalPageIdRouting
+      serverArgs.experimentalPageIdRouting &&
+      !serverArgs.slim
         ? {...tool.schema, ...pageIdSchema}
         : tool.schema;
 
@@ -179,10 +180,13 @@ export async function createMcpServer(
           try {
             if ('pageScoped' in tool && tool.pageScoped) {
               const page =
-                serverArgs.experimentalPageIdRouting && params.pageId
+                serverArgs.experimentalPageIdRouting &&
+                params.pageId &&
+                !serverArgs.slim
                   ? context.resolvePageById(params.pageId)
                   : context.getSelectedMcpPage();
               context.setRequestPage(page);
+              response.setPage(page);
               await tool.handler(
                 {
                   params,

src/tools/ToolDefinition.ts

@@ -136,10 +136,6 @@ export type Context = Readonly<{
   isCruxEnabled(): boolean;
   recordedTraces(): TraceResult[];
   storeTraceRecording(result: TraceResult): void;
-  // TODO: Remove once slim tools are converted to pageScoped: true.
-  getSelectedPage(): Page;
-  getDialog(page?: Page): Dialog | undefined;
-  clearDialog(page?: Page): void;
   getPageById(pageId: number): Page;
   newPage(
     background?: boolean,
@@ -198,6 +194,9 @@ export type ContextPage = Readonly<{
   readonly pptrPage: Page;
   getAXNodeByUid(uid: string): TextSnapshotNode | undefined;
   getElementByUid(uid: string): Promise<ElementHandle<Element>>;
+
+  getDialog(): Dialog | undefined;
+  clearDialog(): void;
 }>;
 
 export function defineTool<Schema extends zod.ZodRawShape>(

src/tools/pages.ts

@@ -187,7 +187,7 @@ export const navigatePage = definePageTool({
           void dialog.dismiss();
         }
         // We are not going to report the dialog like regular dialogs.
-        context.clearDialog(page.pptrPage);
+        page.clearDialog();
       }
     };
 
@@ -333,9 +333,9 @@ export const handleDialog = definePageTool({
       .optional()
       .describe('Optional prompt text to enter into the dialog.'),
   },
-  handler: async (request, response, context) => {
+  handler: async (request, response, _context) => {
     const page = request.page;
-    const dialog = context.getDialog(page.pptrPage);
+    const dialog = page.getDialog();
     if (!dialog) {
       throw new Error('No open dialog found');
     }
@@ -363,7 +363,7 @@ export const handleDialog = definePageTool({
       }
     }
 
-    context.clearDialog(page.pptrPage);
+    page.clearDialog();
     response.setIncludePages(true);
   },
 });

src/tools/slim/tools.ts

@@ -7,9 +7,9 @@
 import type {Dialog} from '../../third_party/index.js';
 import {zod} from '../../third_party/index.js';
 import {ToolCategory} from '../categories.js';
-import {defineTool} from '../ToolDefinition.js';
+import {definePageTool} from '../ToolDefinition.js';
 
-export const screenshot = defineTool({
+export const screenshot = definePageTool({
   name: 'screenshot',
   description: `Takes a screenshot`,
   annotations: {
@@ -19,8 +19,8 @@ export const screenshot = defineTool({
   },
   schema: {},
   handler: async (request, response, context) => {
-    const page = context.getSelectedPage();
-    const screenshot = await page.screenshot({
+    const page = request.page;
+    const screenshot = await page.pptrPage.screenshot({
       type: 'png',
       optimizeForSpeed: true,
     });
@@ -32,7 +32,7 @@ export const screenshot = defineTool({
   },
 });
 
-export const navigate = defineTool({
+export const navigate = definePageTool({
   name: 'navigate',
   description: `Loads a URL`,
   annotations: {
@@ -42,8 +42,9 @@ export const navigate = defineTool({
   schema: {
     url: zod.string().describe('URL to navigate to'),
   },
-  handler: async (request, response, context) => {
-    const page = context.getSelectedPage();
+  handler: async (request, response) => {
+    const page = request.page;
+
     const options = {
       timeout: 30_000,
     };
@@ -53,22 +54,22 @@ export const navigate = defineTool({
         response.appendResponseLine(`Accepted a beforeunload dialog.`);
         void dialog.accept();
         // We are not going to report the dialog like regular dialogs.
-        context.clearDialog();
+        page.clearDialog();
       }
     };
 
-    page.on('dialog', dialogHandler);
+    page.pptrPage.on('dialog', dialogHandler);
 
     try {
-      await page.goto(request.params.url, options);
-      response.appendResponseLine(`Navigated to ${page.url()}.`);
+      await page.pptrPage.goto(request.params.url, options);
+      response.appendResponseLine(`Navigated to ${page.pptrPage.url()}.`);
     } finally {
-      page.off('dialog', dialogHandler);
+      page.pptrPage.off('dialog', dialogHandler);
     }
   },
 });
 
-export const evaluate = defineTool({
+export const evaluate = definePageTool({
   name: 'evaluate',
   description: `Evaluates a JavaScript script`,
   annotations: {
@@ -78,10 +79,10 @@ export const evaluate = defineTool({
   schema: {
     script: zod.string().describe(`JS script to run on the page`),
   },
-  handler: async (request, response, context) => {
-    const page = context.getSelectedPage();
+  handler: async (request, response) => {
+    const page = request.page;
     try {
-      const result = await page.evaluate(request.params.script);
+      const result = await page.pptrPage.evaluate(request.params.script);
       response.appendResponseLine(JSON.stringify(result));
     } catch (err) {
       response.appendResponseLine(String(err.message));

tests/McpContext.test.ts

@@ -22,7 +22,7 @@ describe('McpContext', () => {
 
   it('list pages', async () => {
     await withMcpContext(async (_response, context) => {
-      const page = context.getSelectedPage();
+      const page = context.getSelectedPptrPage();
       await page.setContent(
         html`<button>Click me</button>
           <input
@@ -104,7 +104,7 @@ describe('McpContext', () => {
   it('resolves uid from a non-selected page snapshot', async () => {
     await withMcpContext(async (_response, context) => {
       // Page 1: set content and snapshot
-      const page1 = context.getSelectedPage();
+      const page1 = context.getSelectedPptrPage();
       await page1.setContent(html`<button>Page1 Button</button>`);
       await context.createTextSnapshot(false, undefined, page1);
 
@@ -230,7 +230,7 @@ describe('McpContext', () => {
     it('resolves for default context page alongside isolated contexts', async () => {
       await withMcpContext(async (_response, context) => {
         // Default context page (already exists from withMcpContext setup).
-        const defaultPage = context.getSelectedPage();
+        const defaultPage = context.getSelectedPptrPage();
         await defaultPage.setContent(html`<button>Default Button</button>`);
         await context.createTextSnapshot(false, undefined, defaultPage);
         const defaultUid = '1_1';