Skip to content

Commit c00090d

Browse files
nroscinonroscino
committed
chore: allow script evaluation on service worker
1 parent 13a18e2 commit c00090d

5 files changed

Lines changed: 203 additions & 72 deletions

File tree

‎src/tools/ToolDefinition.ts‎

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

77
import type {ParsedArguments} from '../cli.js';
8+
import {ExtensionServiceWorker} from '../McpContext.js';
89
import {zod} from '../third_party/index.js';
910
import type {
1011
Dialog,
@@ -189,6 +190,10 @@ export type Context = Readonly<{
189190
uninstallExtension(id: string): Promise<void>;
190191
listExtensions(): InstalledExtension[];
191192
getExtension(id: string): InstalledExtension | undefined;
193+
getExtensionServiceWorkers(): ExtensionServiceWorker[];
194+
getExtensionServiceWorkerId(
195+
extensionServiceWorker: ExtensionServiceWorker,
196+
): string | undefined;
192197
}>;
193198

194199
export function defineTool<Schema extends zod.ZodRawShape>(

‎src/tools/script.ts‎

Lines changed: 131 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -4,23 +4,28 @@
44
* SPDX-License-Identifier: Apache-2.0
55
*/
66

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

1011
import {ToolCategory} from './categories.js';
12+
import type {Context} from './ToolDefinition.js';
1113
import {definePageTool} 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 = definePageTool(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,121 @@ 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 context.getElementByUid(el.uid, request.page);
53-
frames.add(handle.frame);
54-
args.push(handle);
55-
}
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.",
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?.categoryExtensions
52+
? {
53+
serviceWorkerId: zod
54+
.string()
55+
.optional()
56+
.describe(
57+
`An optional service worker id to evaluate the script in.`,
58+
),
59+
}
60+
: {}),
61+
},
62+
handler: async (request, response, context) => {
63+
const args: Array<JSHandle<unknown>> = [];
64+
try {
65+
const frames = new Set<Frame>();
66+
for (const el of request.params.args ?? []) {
67+
const handle = await context.getElementByUid(el.uid, request.page);
68+
frames.add(handle.frame);
69+
args.push(handle);
70+
}
71+
72+
const evaluatable = await getEvaluatable(
73+
context,
74+
frames,
75+
cliArgs?.categoryExtensions,
76+
request.params.serviceWorkerId as string | undefined,
6177
);
62-
} else {
63-
pageOrFrame = [...frames.values()][0] ?? request.page;
64-
}
65-
const fn = await pageOrFrame.evaluateHandle(
66-
`(${request.params.function})`,
67-
);
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,
78+
79+
const fn = await evaluatable.evaluateHandle(
80+
`(${request.params.function})`,
7681
);
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()));
84-
}
85-
},
82+
args.unshift(fn);
83+
84+
await context.waitForEventsAfterAction(async () => {
85+
const result = await evaluatable.evaluate(
86+
async (fn, ...args) => {
87+
// @ts-expect-error no types.
88+
return JSON.stringify(await fn(...args));
89+
},
90+
...args,
91+
);
92+
response.appendResponseLine('Script ran on page and returned:');
93+
response.appendResponseLine('```json');
94+
response.appendResponseLine(`${result}`);
95+
response.appendResponseLine('```');
96+
});
97+
} finally {
98+
void Promise.allSettled(args.map(arg => arg.dispose()));
99+
}
100+
},
101+
};
86102
});
103+
104+
const getEvaluatable = async (
105+
context: Context,
106+
frames: Set<Frame>,
107+
enableExtensions?: boolean,
108+
serviceWorkerId?: string,
109+
): Promise<Evaluatable> => {
110+
if (enableExtensions && serviceWorkerId) {
111+
return getWebWorker(context, serviceWorkerId);
112+
}
113+
return getPageOrFrame(context, frames);
114+
};
115+
116+
const getPageOrFrame = async (
117+
context: Context,
118+
frames: Set<Frame>,
119+
): Promise<Page | Frame> => {
120+
let pageOrFrame: Page | Frame;
121+
// We can't evaluate the element handle across frames
122+
if (frames.size > 1) {
123+
throw new Error(
124+
"Elements from different frames can't be evaluated together.",
125+
);
126+
} else {
127+
pageOrFrame = [...frames.values()][0] ?? context.getSelectedPage();
128+
}
129+
130+
return pageOrFrame;
131+
};
132+
133+
const getWebWorker = async (
134+
context: Context,
135+
serviceWorkerId: string,
136+
): Promise<WebWorker> => {
137+
const serviceWorkers = context.getExtensionServiceWorkers();
138+
139+
const serviceWorker = serviceWorkers.find(
140+
(sw: ExtensionServiceWorker) =>
141+
context.getExtensionServiceWorkerId(sw) === serviceWorkerId,
142+
);
143+
144+
if (serviceWorker && serviceWorker.target) {
145+
const worker = await serviceWorker.target.worker();
146+
147+
if (!worker) {
148+
throw new Error('Service worker target not found.');
149+
}
150+
151+
return worker;
152+
} else {
153+
throw new Error('Service worker not found.');
154+
}
155+
};

‎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
}

‎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: (.+)/);

‎tests/tools/script.test.ts‎

Lines changed: 65 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,19 +5,28 @@
55
*/
66

77
import assert from 'node:assert';
8+
import path from 'node:path';
89
import {describe, it} from 'node:test';
910

11+
import type {ParsedArguments} from '../../src/cli.js';
12+
import {installExtension} from '../../src/tools/extensions.js';
1013
import {evaluateScript} from '../../src/tools/script.js';
1114
import {serverHooks} from '../server.js';
1215
import {html, withMcpContext} from '../utils.js';
16+
import {extractId} from './extensions.test.js';
17+
18+
const EXTENSION_PATH = path.join(
19+
import.meta.dirname,
20+
'../../../tests/tools/fixtures/extension-sw',
21+
);
1322

1423
describe('script', () => {
1524
const server = serverHooks();
1625

1726
describe('browser_evaluate_script', () => {
1827
it('evaluates', async () => {
1928
await withMcpContext(async (response, context) => {
20-
await evaluateScript.handler(
29+
await evaluateScript().handler(
2130
{
2231
params: {function: String(() => 2 * 5)},
2332
page: context.getSelectedPage(),
@@ -31,7 +40,7 @@ describe('script', () => {
3140
});
3241
it('runs in selected page', async () => {
3342
await withMcpContext(async (response, context) => {
34-
await evaluateScript.handler(
43+
await evaluateScript().handler(
3544
{
3645
params: {function: String(() => document.title)},
3746
page: context.getSelectedPage(),
@@ -51,7 +60,7 @@ describe('script', () => {
5160
`);
5261

5362
response.resetResponseLineForTesting();
54-
await evaluateScript.handler(
63+
await evaluateScript().handler(
5564
{
5665
params: {function: String(() => document.title)},
5766
page: context.getSelectedPage(),
@@ -71,7 +80,7 @@ describe('script', () => {
7180

7281
await page.setContent(html`<script src="./scripts.js"></script> `);
7382

74-
await evaluateScript.handler(
83+
await evaluateScript().handler(
7584
{
7685
params: {
7786
function: String(() => {
@@ -100,7 +109,7 @@ describe('script', () => {
100109

101110
await page.setContent(html`<script src="./scripts.js"></script> `);
102111

103-
await evaluateScript.handler(
112+
await evaluateScript().handler(
104113
{
105114
params: {
106115
function: String(async () => {
@@ -126,7 +135,7 @@ describe('script', () => {
126135

127136
await context.createTextSnapshot();
128137

129-
await evaluateScript.handler(
138+
await evaluateScript().handler(
130139
{
131140
params: {
132141
function: String(async (el: Element) => {
@@ -152,7 +161,7 @@ describe('script', () => {
152161

153162
await context.createTextSnapshot();
154163

155-
await evaluateScript.handler(
164+
await evaluateScript().handler(
156165
{
157166
params: {
158167
function: String((container: Element, child: Element) => {
@@ -181,7 +190,7 @@ describe('script', () => {
181190
const page = context.getSelectedPage();
182191
await page.goto(server.getRoute('/main'));
183192
await context.createTextSnapshot();
184-
await evaluateScript.handler(
193+
await evaluateScript().handler(
185194
{
186195
params: {
187196
function: String((element: Element) => {
@@ -198,5 +207,53 @@ describe('script', () => {
198207
assert.strictEqual(JSON.parse(lineEvaluation), 'I am iframe button');
199208
});
200209
});
210+
it('evaluates inside extension service worker', async () => {
211+
await withMcpContext(
212+
async (response, context) => {
213+
await installExtension.handler(
214+
{params: {path: EXTENSION_PATH}},
215+
response,
216+
context,
217+
);
218+
219+
const extensionId = extractId(response);
220+
const swTarget = await context.browser.waitForTarget(
221+
t => t.type() === 'service_worker' && t.url().includes(extensionId),
222+
);
223+
224+
await context.createExtensionServiceWorkersSnapshot();
225+
const swList = context.getExtensionServiceWorkers();
226+
const sw = swList.find(s => s.target === swTarget);
227+
228+
if (!sw) {
229+
assert.fail('Service worker not found in context list');
230+
}
231+
232+
const swId = context.getExtensionServiceWorkerId(sw);
233+
234+
response.resetResponseLineForTesting();
235+
await evaluateScript({
236+
categoryExtensions: true,
237+
} as ParsedArguments).handler(
238+
{
239+
params: {
240+
function: String(() => {
241+
return 'chrome' in globalThis ? 'has-chrome' : 'no-chrome';
242+
}),
243+
serviceWorkerId: swId,
244+
},
245+
page: context.getSelectedPage(),
246+
},
247+
response,
248+
context,
249+
);
250+
251+
const lineEvaluation = response.responseLines.at(2)!;
252+
assert.strictEqual(JSON.parse(lineEvaluation), 'has-chrome');
253+
},
254+
{},
255+
{categoryExtensions: true} as ParsedArguments,
256+
);
257+
});
201258
});
202259
});

0 commit comments

Comments
 (0)