diff --git a/src/McpContext.ts b/src/McpContext.ts index 65c6b0671..76cbfdbd2 100644 --- a/src/McpContext.ts +++ b/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 { - 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 { - 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, options?: {timeout?: number}, ): Promise { - 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 { - const page = targetPage ?? this.getSelectedPage(); + const page = targetPage ?? this.getSelectedPptrPage(); const frames = page.frames(); let locator = this.#locatorClass.race( diff --git a/src/McpPage.ts b/src/McpPage.ts index fb23a5f48..73717464b 100644 --- a/src/McpPage.ts +++ b/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; } diff --git a/src/McpResponse.ts b/src/McpResponse.ts index 5cbea123e..9a0f19134 100644 --- a/src/McpResponse.ts +++ b/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' diff --git a/src/server.ts b/src/server.ts index ec0df3706..02c0cac20 100644 --- a/src/server.ts +++ b/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, diff --git a/src/tools/ToolDefinition.ts b/src/tools/ToolDefinition.ts index c6c1f2978..982cf9b75 100644 --- a/src/tools/ToolDefinition.ts +++ b/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>; + + getDialog(): Dialog | undefined; + clearDialog(): void; }>; export function defineTool( diff --git a/src/tools/pages.ts b/src/tools/pages.ts index fa724ee4b..9ec277712 100644 --- a/src/tools/pages.ts +++ b/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); }, }); diff --git a/src/tools/slim/tools.ts b/src/tools/slim/tools.ts index 0360eda0d..712dcff63 100644 --- a/src/tools/slim/tools.ts +++ b/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)); diff --git a/tests/McpContext.test.ts b/tests/McpContext.test.ts index f1499032b..387a5bfc1 100644 --- a/tests/McpContext.test.ts +++ b/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` { 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``); 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``); await context.createTextSnapshot(false, undefined, defaultPage); const defaultUid = '1_1'; diff --git a/tests/McpResponse.test.ts b/tests/McpResponse.test.ts index dff5a06da..141f4f0bb 100644 --- a/tests/McpResponse.test.ts +++ b/tests/McpResponse.test.ts @@ -64,7 +64,7 @@ describe('McpResponse', () => { it('does not include anything in response if snapshot is null', async t => { await withMcpContext(async (response, context) => { - const page = context.getSelectedPage(); + const page = context.getSelectedPptrPage(); page.accessibility.snapshot = async () => null; const {content, structuredContent} = await response.handle( 'test', @@ -80,7 +80,7 @@ describe('McpResponse', () => { it('returns correctly formatted snapshot for a simple tree', async t => { await withMcpContext(async (response, context) => { - const page = context.getSelectedPage(); + const page = context.getSelectedPptrPage(); await page.setContent( html` { it('returns values for textboxes', async t => { await withMcpContext(async (response, context) => { - const page = context.getSelectedPage(); + const page = context.getSelectedPptrPage(); await page.setContent( html`