Skip to content

Commit 91d5d24

Browse files
yury-syury-s
authored
chore: handle list roots in the server, with timeout (#898)
1 parent 92554ab commit 91d5d24

File tree

4 files changed

+72
-78
lines changed

4 files changed

+72
-78
lines changed

‎src/browserServerBackend.ts‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -45,11 +45,9 @@ export class BrowserServerBackend implements ServerBackend {
4545
this._tools = filteredTools(config);
4646
}
4747

48-
async initialize(server: mcpServer.Server): Promise<void> {
49-
const capabilities = server.getClientCapabilities();
48+
async initialize(clientVersion: mcpServer.ClientVersion, roots: mcpServer.Root[]): Promise<void> {
5049
let rootPath: string | undefined;
51-
if (capabilities?.roots) {
52-
const { roots } = await server.listRoots();
50+
if (roots.length > 0) {
5351
const firstRootUri = roots[0]?.uri;
5452
const url = firstRootUri ? new URL(firstRootUri) : undefined;
5553
rootPath = url ? fileURLToPath(url) : undefined;
@@ -60,7 +58,7 @@ export class BrowserServerBackend implements ServerBackend {
6058
config: this._config,
6159
browserContextFactory: this._browserContextFactory,
6260
sessionLog: this._sessionLog,
63-
clientInfo: { ...server.getClientVersion(), rootPath },
61+
clientInfo: { ...clientVersion, rootPath },
6462
});
6563
}
6664

‎src/mcp/proxyBackend.ts‎

Lines changed: 4 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,9 @@ import { logUnhandledError } from '../utils/log.js';
2323
import { packageJSON } from '../utils/package.js';
2424

2525

26-
import type { Server } from '@modelcontextprotocol/sdk/server/index.js';
27-
import type { ServerBackend } from './server.js';
26+
import type { ServerBackend, ClientVersion, Root } from './server.js';
2827
import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
29-
import type { Root, Tool, CallToolResult, CallToolRequest } from '@modelcontextprotocol/sdk/types.js';
28+
import type { Tool, CallToolResult, CallToolRequest } from '@modelcontextprotocol/sdk/types.js';
3029

3130
export type MCPProvider = {
3231
name: string;
@@ -48,14 +47,8 @@ export class ProxyBackend implements ServerBackend {
4847
this._contextSwitchTool = this._defineContextSwitchTool();
4948
}
5049

51-
async initialize(server: Server): Promise<void> {
52-
const version = server.getClientVersion();
53-
const capabilities = server.getClientCapabilities();
54-
if (capabilities?.roots && version && clientsWithRoots.includes(version.name)) {
55-
const { roots } = await server.listRoots();
56-
this._roots = roots;
57-
}
58-
50+
async initialize(clientVersion: ClientVersion, roots: Root[]): Promise<void> {
51+
this._roots = roots;
5952
await this._setCurrentClient(this._mcpProviders[0]);
6053
}
6154

@@ -136,5 +129,3 @@ export class ProxyBackend implements ServerBackend {
136129
this._currentClient = client;
137130
}
138131
}
139-
140-
const clientsWithRoots = ['Visual Studio Code', 'Visual Studio Code - Insiders'];

‎src/mcp/server.ts‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,17 +20,18 @@ import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprot
2020
import { ManualPromise } from '../utils/manualPromise.js';
2121
import { logUnhandledError } from '../utils/log.js';
2222

23-
import type { Tool, CallToolResult, CallToolRequest } from '@modelcontextprotocol/sdk/types.js';
23+
import type { Tool, CallToolResult, CallToolRequest, Root } from '@modelcontextprotocol/sdk/types.js';
2424
import type { Transport } from '@modelcontextprotocol/sdk/shared/transport.js';
2525
export type { Server } from '@modelcontextprotocol/sdk/server/index.js';
26-
export type { Tool, CallToolResult, CallToolRequest } from '@modelcontextprotocol/sdk/types.js';
26+
export type { Tool, CallToolResult, CallToolRequest, Root } from '@modelcontextprotocol/sdk/types.js';
2727

2828
const serverDebug = debug('pw:mcp:server');
2929

30+
export type ClientVersion = { name: string, version: string };
3031
export interface ServerBackend {
3132
name: string;
3233
version: string;
33-
initialize?(server: Server): Promise<void>;
34+
initialize?(clientVersion: ClientVersion, roots: Root[]): Promise<void>;
3435
listTools(): Promise<Tool[]>;
3536
callTool(name: string, args: CallToolRequest['params']['arguments']): Promise<CallToolResult>;
3637
serverClosed?(): void;
@@ -78,8 +79,20 @@ export function createServer(backend: ServerBackend, runHeartbeat: boolean): Ser
7879
};
7980
}
8081
});
81-
addServerListener(server, 'initialized', () => {
82-
backend.initialize?.(server).then(() => initializedPromise.resolve()).catch(logUnhandledError);
82+
addServerListener(server, 'initialized', async () => {
83+
try {
84+
const capabilities = server.getClientCapabilities();
85+
let clientRoots: Root[] = [];
86+
if (capabilities?.roots) {
87+
const { roots } = await server.listRoots(undefined, { timeout: 2_000 }).catch(() => ({ roots: [] }));
88+
clientRoots = roots;
89+
}
90+
const clientVersion = server.getClientVersion() ?? { name: 'unknown', version: 'unknown' };
91+
await backend.initialize?.(clientVersion, clientRoots);
92+
initializedPromise.resolve();
93+
} catch (e) {
94+
logUnhandledError(e);
95+
}
8396
});
8497
addServerListener(server, 'close', () => backend.serverClosed?.());
8598
return server;

‎tests/roots.spec.ts‎

Lines changed: 47 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -23,65 +23,57 @@ import { createHash } from '../src/utils/guid.js';
2323

2424
const p = process.platform === 'win32' ? 'c:\\non\\existent\\folder' : '/non/existent/folder';
2525

26-
for (const mode of ['default', 'proxy']) {
27-
const extraArgs = mode === 'proxy' ? ['--connect-tool'] : [];
28-
29-
test.describe(`${mode} mode`, () => {
30-
test('should use separate user data by root path', async ({ startClient, server }, testInfo) => {
31-
const { client } = await startClient({
32-
args: extraArgs,
33-
clientName: 'Visual Studio Code', // Simulate VS Code client, roots only work with it
34-
roots: [
35-
{
36-
name: 'test',
37-
uri: 'file://' + p.replace(/\\/g, '/'),
38-
}
39-
],
40-
});
41-
42-
await client.callTool({
43-
name: 'browser_navigate',
44-
arguments: { url: server.HELLO_WORLD },
45-
});
26+
test('should use separate user data by root path', async ({ startClient, server }, testInfo) => {
27+
const { client } = await startClient({
28+
clientName: 'Visual Studio Code',
29+
roots: [
30+
{
31+
name: 'test',
32+
uri: 'file://' + p.replace(/\\/g, '/'),
33+
}
34+
],
35+
});
4636

47-
const hash = createHash(p);
48-
const [file] = await fs.promises.readdir(testInfo.outputPath('ms-playwright'));
49-
expect(file).toContain(hash);
50-
});
37+
await client.callTool({
38+
name: 'browser_navigate',
39+
arguments: { url: server.HELLO_WORLD },
40+
});
5141

42+
const hash = createHash(p);
43+
const [file] = await fs.promises.readdir(testInfo.outputPath('ms-playwright'));
44+
expect(file).toContain(hash);
45+
});
5246

53-
test('check that trace is saved in workspace', async ({ startClient, server }, testInfo) => {
54-
const rootPath = testInfo.outputPath('workspace');
55-
const { client } = await startClient({
56-
args: ['--save-trace', ...extraArgs],
57-
clientName: 'Visual Studio Code - Insiders', // Simulate VS Code client, roots only work with it
58-
roots: [
59-
{
60-
name: 'workspace',
61-
uri: pathToFileURL(rootPath).toString(),
62-
},
63-
],
64-
});
47+
test('check that trace is saved in workspace', async ({ startClient, server }, testInfo) => {
48+
const rootPath = testInfo.outputPath('workspace');
49+
const { client } = await startClient({
50+
args: ['--save-trace'],
51+
clientName: 'My client',
52+
roots: [
53+
{
54+
name: 'workspace',
55+
uri: pathToFileURL(rootPath).toString(),
56+
},
57+
],
58+
});
6559

66-
expect(await client.callTool({
67-
name: 'browser_navigate',
68-
arguments: { url: server.HELLO_WORLD },
69-
})).toHaveResponse({
70-
code: expect.stringContaining(`page.goto('http://localhost`),
71-
});
60+
expect(await client.callTool({
61+
name: 'browser_navigate',
62+
arguments: { url: server.HELLO_WORLD },
63+
})).toHaveResponse({
64+
code: expect.stringContaining(`page.goto('http://localhost`),
65+
});
7266

73-
const [file] = await fs.promises.readdir(path.join(rootPath, '.playwright-mcp'));
74-
expect(file).toContain('traces');
75-
});
67+
const [file] = await fs.promises.readdir(path.join(rootPath, '.playwright-mcp'));
68+
expect(file).toContain('traces');
69+
});
7670

77-
test('should list all tools when listRoots is slow', async ({ startClient, server }, testInfo) => {
78-
const { client } = await startClient({
79-
clientName: 'Visual Studio Code', // Simulate VS Code client, roots only work with it
80-
roots: [],
81-
rootsResponseDelay: 1000,
82-
});
83-
const tools = await client.listTools();
84-
expect(tools.tools.length).toBeGreaterThan(20);
85-
});
71+
test('should list all tools when listRoots is slow', async ({ startClient, server }, testInfo) => {
72+
const { client } = await startClient({
73+
clientName: 'Another custom client',
74+
roots: [],
75+
rootsResponseDelay: 1000,
8676
});
87-
}
77+
const tools = await client.listTools();
78+
expect(tools.tools.length).toBeGreaterThan(20);
79+
});

0 commit comments

Comments
 (0)