make in-page tools work with pageId routing

Author: wolfib

src/McpContext.ts

@@ -31,7 +31,6 @@ import type {
 } from './third_party/index.js';
 import {Locator} from './third_party/index.js';
 import {PredefinedNetworkConditions} from './third_party/index.js';
-import type {ToolGroup, ToolDefinition} from './tools/inPage.js';
 import {listPages} from './tools/pages.js';
 import {CLOSE_PAGE_ERROR} from './tools/ToolDefinition.js';
 import type {Context, DevToolsData} from './tools/ToolDefinition.js';
@@ -102,7 +101,6 @@ export class McpContext implements Context {
   #screenRecorderData: {recorder: ScreenRecorder; filePath: string} | null =
     null;
 
-  #inPageTools?: ToolGroup<ToolDefinition>;
   #nextPageId = 1;
   #extensionPages = new WeakMap<Target, Page>();
 
@@ -466,14 +464,6 @@ export class McpContext implements Context {
     this.#updateSelectedPageTimeouts();
   }
 
-  setInPageTools(toolGroup?: ToolGroup<ToolDefinition>) {
-    this.#inPageTools = toolGroup;
-  }
-
-  getInPageTools(): ToolGroup<ToolDefinition> | undefined {
-    return this.#inPageTools;
-  }
-
   #updateSelectedPageTimeouts() {
     const page = this.#getSelectedMcpPage();
     // For waiters 5sec timeout should be sufficient.

src/McpPage.ts

@@ -10,6 +10,7 @@ import type {
   Page,
   Viewport,
 } from './third_party/index.js';
+import type {ToolGroup, ToolDefinition} from './tools/inPage.js';
 import {takeSnapshot} from './tools/snapshot.js';
 import type {ContextPage} from './tools/ToolDefinition.js';
 import type {
@@ -46,6 +47,8 @@ export class McpPage implements ContextPage {
   #dialog?: Dialog;
   #dialogHandler: (dialog: Dialog) => void;
 
+  inPageTools: ToolGroup<ToolDefinition> | undefined;
+
   constructor(page: Page, id: number) {
     this.pptrPage = page;
     this.id = id;
@@ -67,6 +70,10 @@ export class McpPage implements ContextPage {
     this.#dialog = undefined;
   }
 
+  getInPageTools(): ToolGroup<ToolDefinition> | undefined {
+    return this.inPageTools;
+  }
+
   get networkConditions(): string | null {
     return this.emulationSettings.networkConditions ?? null;
   }

src/McpResponse.ts

@@ -421,8 +421,9 @@ export class McpResponse implements Response {
 
     let inPageTools: ToolGroup<ToolDefinition> | undefined;
     if (this.#listInPageTools) {
-      inPageTools = await getToolGroup(context.getSelectedMcpPage());
-      context.setInPageTools(inPageTools);
+      const page = this.#page ?? context.getSelectedMcpPage();
+      inPageTools = await getToolGroup(page);
+      page.inPageTools = inPageTools;
     }
 
     let consoleMessages: Array<ConsoleFormatter | IssueFormatter> | undefined;

src/tools/ToolDefinition.ts

@@ -198,7 +198,6 @@ export type Context = Readonly<{
   triggerExtensionAction(id: string): Promise<void>;
   listExtensions(): InstalledExtension[];
   getExtension(id: string): InstalledExtension | undefined;
-  getInPageTools(): ToolGroup<InPageToolDefinition> | undefined;
   getSelectedMcpPage(): McpPage;
   getExtensionServiceWorkers(): ExtensionServiceWorker[];
   getExtensionServiceWorkerId(
@@ -213,6 +212,8 @@ export type ContextPage = Readonly<{
 
   getDialog(): Dialog | undefined;
   clearDialog(): void;
+
+  getInPageTools(): ToolGroup<InPageToolDefinition> | undefined;
 }>;
 
 export function defineTool<Schema extends zod.ZodRawShape>(

src/tools/inPage.ts

@@ -70,8 +70,7 @@ export const executeInPageTool = definePageTool({
       .optional()
       .describe('The JSON-stringified parameters to pass to the tool'),
   },
-  handler: async (request, response, context) => {
-    const page = context.getSelectedMcpPage();
+  handler: async (request, response) => {
     const toolName = request.params.toolName;
     let params: Record<string, unknown> = {};
     if (request.params.params) {
@@ -88,7 +87,7 @@ export const executeInPageTool = definePageTool({
       }
     }
 
-    const toolGroup = context.getInPageTools();
+    const toolGroup = request.page.getInPageTools();
     const tool = toolGroup?.tools.find(t => t.name === toolName);
     if (!tool) {
       throw new Error(`Tool ${toolName} not found`);
@@ -102,7 +101,7 @@ export const executeInPageTool = definePageTool({
       );
     }
 
-    const result = await page.pptrPage.evaluate(
+    const result = await request.page.pptrPage.evaluate(
       async (name, args) => {
         if (!window.__dtmcp?.executeTool) {
           throw new Error('No tools found on the page');

tests/tools/inPage.test.ts

@@ -20,6 +20,7 @@ describe('inPage', () => {
       await withMcpContext(
         async (response, context) => {
           const page = await context.newPage();
+          response.setPage(page);
 
           await page.pptrPage.evaluate(() => {
             window.__dtmcp = {
@@ -76,6 +77,7 @@ describe('inPage', () => {
       await withMcpContext(
         async (response, context) => {
           const page = await context.newPage();
+          response.setPage(page);
           await page.pptrPage.evaluate(() => {
             window.addEventListener('devtoolstooldiscovery', (e: Event) => {
               // @ts-expect-error Event has `respondWith`
@@ -105,6 +107,7 @@ describe('inPage', () => {
       await withMcpContext(
         async (response, context) => {
           const page = await context.newPage();
+          response.setPage(page);
           await page.pptrPage.evaluate(() => {
             window.addEventListener('devtoolstooldiscovery', () => {
               // do nothing
@@ -133,6 +136,7 @@ describe('inPage', () => {
       await withMcpContext(
         async (response, context) => {
           const page = await context.newPage();
+          response.setPage(page);
           await listInPageTools.handler({params: {}, page}, response, context);
 
           const result = await response.handle('list_in_page_tools', context);
@@ -155,6 +159,7 @@ describe('inPage', () => {
       evaluateFn: () => void,
     ) {
       const page = await context.newPage();
+      response.setPage(page);
       await page.pptrPage.evaluate(evaluateFn);
       await listInPageTools.handler({params: {}, page}, response, context);
       await response.handle('list_in_page_tools', context);