Skip to content

Commit 76ff677

Browse files
nroscinonroscino
committed
chore: allow script evaluation on service worker
1 parent 9628dab commit 76ff677

7 files changed

Lines changed: 295 additions & 88 deletions

File tree

‎src/McpContext.ts‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ import type {
4545
GeolocationOptions,
4646
TextSnapshot,
4747
TextSnapshotNode,
48+
ExtensionServiceWorker,
4849
} from './types.js';
4950
import {
5051
ExtensionRegistry,
@@ -59,12 +60,6 @@ export type {
5960
TextSnapshotNode,
6061
} from './types.js';
6162

62-
export interface ExtensionServiceWorker {
63-
url: string;
64-
target: Target;
65-
id: string;
66-
}
67-
6863
interface McpContextOptions {
6964
// Whether the DevTools windows are exposed as pages for debugging of DevTools.
7065
experimentalDevToolsDebugging: boolean;

‎src/tools/ToolDefinition.ts‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
*/
66

77
import type {ParsedArguments} from '../cli.js';
8+
import type {McpPage} from '../McpPage.js';
89
import {zod} from '../third_party/index.js';
910
import type {
1011
Dialog,
@@ -14,7 +15,11 @@ import type {
1415
Viewport,
1516
} from '../third_party/index.js';
1617
import type {InsightName, TraceResult} from '../trace-processing/parse.js';
17-
import type {TextSnapshotNode, GeolocationOptions} from '../types.js';
18+
import type {
19+
TextSnapshotNode,
20+
GeolocationOptions,
21+
ExtensionServiceWorker,
22+
} from '../types.js';
1823
import type {InstalledExtension} from '../utils/ExtensionRegistry.js';
1924
import type {PaginationOptions} from '../utils/types.js';
2025

@@ -136,6 +141,7 @@ export type Context = Readonly<{
136141
recordedTraces(): TraceResult[];
137142
storeTraceRecording(result: TraceResult): void;
138143
getPageById(pageId: number): ContextPage;
144+
resolvePageById(pageId?: number): ContextPage;
139145
newPage(
140146
background?: boolean,
141147
isolatedContextName?: string,
@@ -188,6 +194,11 @@ export type Context = Readonly<{
188194
uninstallExtension(id: string): Promise<void>;
189195
listExtensions(): InstalledExtension[];
190196
getExtension(id: string): InstalledExtension | undefined;
197+
getSelectedMcpPage(): McpPage;
198+
getExtensionServiceWorkers(): ExtensionServiceWorker[];
199+
getExtensionServiceWorkerId(
200+
extensionServiceWorker: ExtensionServiceWorker,
201+
): string | undefined;
191202
}>;
192203

193204
export type ContextPage = Readonly<{

‎src/tools/script.ts‎

Lines changed: 153 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -5,22 +5,27 @@
55
*/
66

77
import {zod} from '../third_party/index.js';
8-
import type {Frame, JSHandle, Page} from '../third_party/index.js';
8+
import type {Frame, JSHandle, Page, WebWorker} from '../third_party/index.js';
9+
import type {ExtensionServiceWorker} from '../types.js';
910

1011
import {ToolCategory} from './categories.js';
11-
import {definePageTool} from './ToolDefinition.js';
12+
import type {Context, Response} from './ToolDefinition.js';
13+
import {defineTool, pageIdSchema} from './ToolDefinition.js';
1214

13-
export const evaluateScript = definePageTool({
14-
name: 'evaluate_script',
15-
description: `Evaluate a JavaScript function inside the currently selected page. Returns the response as JSON,
15+
export type Evaluatable = Page | Frame | WebWorker;
16+
17+
export const evaluateScript = defineTool(cliArgs => {
18+
return {
19+
name: 'evaluate_script',
20+
description: `Evaluate a JavaScript function inside the currently selected page. Returns the response as JSON,
1621
so returned values have to be JSON-serializable.`,
17-
annotations: {
18-
category: ToolCategory.DEBUGGING,
19-
readOnlyHint: false,
20-
},
21-
schema: {
22-
function: zod.string().describe(
23-
`A JavaScript function declaration to be executed by the tool in the currently selected page.
22+
annotations: {
23+
category: ToolCategory.DEBUGGING,
24+
readOnlyHint: false,
25+
},
26+
schema: {
27+
function: zod.string().describe(
28+
`A JavaScript function declaration to be executed by the tool in the currently selected page.
2429
Example without arguments: \`() => {
2530
return document.title
2631
}\` or \`async () => {
@@ -30,57 +35,143 @@ Example with arguments: \`(el) => {
3035
return el.innerText;
3136
}\`
3237
`,
33-
),
34-
args: zod
35-
.array(
36-
zod.object({
37-
uid: zod
38-
.string()
39-
.describe(
40-
'The uid of an element on the page from the page content snapshot',
41-
),
42-
}),
43-
)
44-
.optional()
45-
.describe(`An optional list of arguments to pass to the function.`),
46-
},
47-
handler: async (request, response, context) => {
48-
const args: Array<JSHandle<unknown>> = [];
49-
try {
50-
const frames = new Set<Frame>();
51-
for (const el of request.params.args ?? []) {
52-
const handle = await request.page.getElementByUid(el.uid);
53-
frames.add(handle.frame);
54-
args.push(handle);
38+
),
39+
args: zod
40+
.array(
41+
zod.object({
42+
uid: zod
43+
.string()
44+
.describe(
45+
'The uid of an element on the page from the page content snapshot',
46+
),
47+
}),
48+
)
49+
.optional()
50+
.describe(`An optional list of arguments to pass to the function.`),
51+
...(cliArgs?.experimentalPageIdRouting ? pageIdSchema : {}),
52+
...(cliArgs?.categoryExtensions
53+
? {
54+
serviceWorkerId: zod
55+
.string()
56+
.optional()
57+
.describe(
58+
`An optional service worker id to evaluate the script in.`,
59+
),
60+
}
61+
: {}),
62+
},
63+
handler: async (request, response, context) => {
64+
const {
65+
serviceWorkerId,
66+
args: uidArgs,
67+
function: fnString,
68+
pageId,
69+
} = request.params;
70+
71+
if (cliArgs?.categoryExtensions && serviceWorkerId) {
72+
if (uidArgs && uidArgs.length > 0) {
73+
throw new Error(
74+
'args (element uids) cannot be used when evaluating in a service worker.',
75+
);
76+
}
77+
if (pageId) {
78+
throw new Error('specify either a pageId or a serviceWorkerId.');
79+
}
80+
81+
const worker = await getWebWorker(context, serviceWorkerId);
82+
await performEvaluation(worker, fnString, [], response, context);
83+
return;
5584
}
56-
let pageOrFrame: Page | Frame;
57-
// We can't evaluate the element handle across frames
58-
if (frames.size > 1) {
59-
throw new Error(
60-
"Elements from different frames can't be evaluated together.",
61-
);
62-
} else {
63-
pageOrFrame = [...frames.values()][0] ?? request.page.pptrPage;
85+
86+
const mcpPage = cliArgs?.experimentalPageIdRouting
87+
? context.resolvePageById(request.params.pageId)
88+
: context.getSelectedMcpPage();
89+
const page: Page = mcpPage.pptrPage;
90+
91+
const args: Array<JSHandle<unknown>> = [];
92+
try {
93+
const frames = new Set<Frame>();
94+
for (const el of uidArgs ?? []) {
95+
const handle = await context.getElementByUid(el.uid, page);
96+
frames.add(handle.frame);
97+
args.push(handle);
98+
}
99+
100+
const evaluatable = await getPageOrFrame(page, frames);
101+
102+
await performEvaluation(evaluatable, fnString, args, response, context);
103+
} finally {
104+
void Promise.allSettled(args.map(arg => arg.dispose()));
64105
}
65-
const fn = await pageOrFrame.evaluateHandle(
66-
`(${request.params.function})`,
106+
},
107+
};
108+
});
109+
110+
const performEvaluation = async (
111+
evaluatable: Evaluatable,
112+
fnString: string,
113+
args: Array<JSHandle<unknown>>,
114+
response: Response,
115+
context: Context,
116+
) => {
117+
const fn = await evaluatable.evaluateHandle(`(${fnString})`);
118+
try {
119+
await context.waitForEventsAfterAction(async () => {
120+
const result = await evaluatable.evaluate(
121+
async (fn, ...args) => {
122+
// @ts-expect-error no types for function fn
123+
return JSON.stringify(await fn(...args));
124+
},
125+
fn,
126+
...args,
67127
);
68-
args.unshift(fn);
69-
await context.waitForEventsAfterAction(async () => {
70-
const result = await pageOrFrame.evaluate(
71-
async (fn, ...args) => {
72-
// @ts-expect-error no types.
73-
return JSON.stringify(await fn(...args));
74-
},
75-
...args,
76-
);
77-
response.appendResponseLine('Script ran on page and returned:');
78-
response.appendResponseLine('```json');
79-
response.appendResponseLine(`${result}`);
80-
response.appendResponseLine('```');
81-
});
82-
} finally {
83-
void Promise.allSettled(args.map(arg => arg.dispose()));
128+
response.appendResponseLine('Script ran on page and returned:');
129+
response.appendResponseLine('```json');
130+
response.appendResponseLine(`${result}`);
131+
response.appendResponseLine('```');
132+
});
133+
} finally {
134+
void fn.dispose();
135+
}
136+
};
137+
138+
const getPageOrFrame = async (
139+
page: Page,
140+
frames: Set<Frame>,
141+
): Promise<Page | Frame> => {
142+
let pageOrFrame: Page | Frame;
143+
// We can't evaluate the element handle across frames
144+
if (frames.size > 1) {
145+
throw new Error(
146+
"Elements from different frames can't be evaluated together.",
147+
);
148+
} else {
149+
pageOrFrame = [...frames.values()][0] ?? page;
150+
}
151+
152+
return pageOrFrame;
153+
};
154+
155+
const getWebWorker = async (
156+
context: Context,
157+
serviceWorkerId: string,
158+
): Promise<WebWorker> => {
159+
const serviceWorkers = context.getExtensionServiceWorkers();
160+
161+
const serviceWorker = serviceWorkers.find(
162+
(sw: ExtensionServiceWorker) =>
163+
context.getExtensionServiceWorkerId(sw) === serviceWorkerId,
164+
);
165+
166+
if (serviceWorker && serviceWorker.target) {
167+
const worker = await serviceWorker.target.worker();
168+
169+
if (!worker) {
170+
throw new Error('Service worker target not found.');
84171
}
85-
},
86-
});
172+
173+
return worker;
174+
} else {
175+
throw new Error('Service worker not found.');
176+
}
177+
};

‎src/tools/tools.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ export const createTools = (args: ParsedArguments) => {
4444
const tools = [];
4545
for (const tool of rawTools) {
4646
if (typeof tool === 'function') {
47-
tools.push(tool(args));
47+
tools.push(tool(args) as unknown as ToolDefinition);
4848
} else {
4949
tools.push(tool as ToolDefinition);
5050
}

‎src/types.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,13 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66

7-
import type {SerializedAXNode, Viewport} from './third_party/index.js';
7+
import type {SerializedAXNode, Viewport, Target} from './third_party/index.js';
8+
9+
export interface ExtensionServiceWorker {
10+
url: string;
11+
target: Target;
12+
id: string;
13+
}
814

915
export interface TextSnapshotNode extends SerializedAXNode {
1016
id: string;

‎tests/tools/extensions.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ const EXTENSION_PATH = path.join(
2424
'../../../tests/tools/fixtures/extension',
2525
);
2626

27-
function extractId(response: McpResponse) {
27+
export function extractId(response: McpResponse) {
2828
const responseLine = response.responseLines[0];
2929
assert.ok(responseLine, 'Response should not be empty');
3030
const match = responseLine.match(/Extension installed\. Id: (.+)/);

0 commit comments

Comments
 (0)