diff --git a/src/filesystem/__tests__/roots-gate.test.ts b/src/filesystem/__tests__/roots-gate.test.ts new file mode 100644 index 0000000000..2bb9dc00bd --- /dev/null +++ b/src/filesystem/__tests__/roots-gate.test.ts @@ -0,0 +1,41 @@ +import { describe, it, expect } from 'vitest'; +import { createRootsGate } from '../roots-gate.js'; + +describe('createRootsGate', () => { + it('resolves immediately when gate is pre-resolved', async () => { + const gate = createRootsGate(); + gate.resolve(); + await expect(gate.waitForReady()).resolves.toBeUndefined(); + }); + + it('resolves immediately when called after the gate has already resolved', async () => { + const gate = createRootsGate(); + gate.resolve(); + await expect(gate.waitForReady()).resolves.toBeUndefined(); + await expect(gate.waitForReady()).resolves.toBeUndefined(); + }); + + it('resolves after a delay when resolve is called later', async () => { + const gate = createRootsGate(); + setTimeout(() => gate.resolve(), 50); + await expect(gate.waitForReady()).resolves.toBeUndefined(); + }); + + it('rejects with timeout when gate is never resolved', async () => { + const gate = createRootsGate(100); + await expect(gate.waitForReady()).rejects.toThrow( + 'Roots initialization timed out after 100ms' + ); + }); + + it('resolves all concurrent waiters together', async () => { + const gate = createRootsGate(); + const results = Promise.all([ + gate.waitForReady(), + gate.waitForReady(), + gate.waitForReady(), + ]); + gate.resolve(); + await expect(results).resolves.toEqual([undefined, undefined, undefined]); + }); +}); diff --git a/src/filesystem/index.ts b/src/filesystem/index.ts index 7b67e63e58..8e36fd42d2 100644 --- a/src/filesystem/index.ts +++ b/src/filesystem/index.ts @@ -14,6 +14,7 @@ import { z } from "zod"; import { minimatch } from "minimatch"; import { normalizePath, expandHome } from './path-utils.js'; import { getValidRootDirectories } from './roots-utils.js'; +import { createRootsGate } from './roots-gate.js'; import { // Function imports formatSize, @@ -92,6 +93,13 @@ allowedDirectories = accessibleDirectories; // Initialize the global allowedDirectories in lib.ts setAllowedDirectories(allowedDirectories); +// Gate to block tool handlers until roots/directories are initialized. +// Without this, tool calls can race ahead of oninitialized and see empty allowedDirectories. +const rootsGate = createRootsGate(); +if (allowedDirectories.length > 0) { + rootsGate.resolve(); +} + // Schema definitions const ReadTextFileArgsSchema = z.object({ path: z.string(), @@ -189,6 +197,7 @@ async function readFileAsBase64Stream(filePath: string): Promise { // read_file (deprecated) and read_text_file const readTextFileHandler = async (args: z.infer) => { + await rootsGate.waitForReady(); const validPath = await validatePath(args.path); if (args.head && args.tail) { @@ -265,6 +274,7 @@ server.registerTool( annotations: { readOnlyHint: true } }, async (args: z.infer) => { + await rootsGate.waitForReady(); const validPath = await validatePath(args.path); const extension = path.extname(validPath).toLowerCase(); const mimeTypes: Record = { @@ -316,6 +326,7 @@ server.registerTool( annotations: { readOnlyHint: true } }, async (args: z.infer) => { + await rootsGate.waitForReady(); const results = await Promise.all( args.paths.map(async (filePath: string) => { try { @@ -352,6 +363,7 @@ server.registerTool( annotations: { readOnlyHint: false, idempotentHint: true, destructiveHint: true } }, async (args: z.infer) => { + await rootsGate.waitForReady(); const validPath = await validatePath(args.path); await writeFileContent(validPath, args.content); const text = `Successfully wrote to ${args.path}`; @@ -382,6 +394,7 @@ server.registerTool( annotations: { readOnlyHint: false, idempotentHint: false, destructiveHint: true } }, async (args: z.infer) => { + await rootsGate.waitForReady(); const validPath = await validatePath(args.path); const result = await applyFileEdits(validPath, args.edits, args.dryRun); return { @@ -407,6 +420,7 @@ server.registerTool( annotations: { readOnlyHint: false, idempotentHint: true, destructiveHint: false } }, async (args: z.infer) => { + await rootsGate.waitForReady(); const validPath = await validatePath(args.path); await fs.mkdir(validPath, { recursive: true }); const text = `Successfully created directory ${args.path}`; @@ -433,6 +447,7 @@ server.registerTool( annotations: { readOnlyHint: true } }, async (args: z.infer) => { + await rootsGate.waitForReady(); const validPath = await validatePath(args.path); const entries = await fs.readdir(validPath, { withFileTypes: true }); const formatted = entries @@ -462,6 +477,7 @@ server.registerTool( annotations: { readOnlyHint: true } }, async (args: z.infer) => { + await rootsGate.waitForReady(); const validPath = await validatePath(args.path); const entries = await fs.readdir(validPath, { withFileTypes: true }); @@ -541,6 +557,7 @@ server.registerTool( annotations: { readOnlyHint: true } }, async (args: z.infer) => { + await rootsGate.waitForReady(); interface TreeEntry { name: string; type: 'file' | 'directory'; @@ -611,6 +628,7 @@ server.registerTool( annotations: { readOnlyHint: false, idempotentHint: false, destructiveHint: true } }, async (args: z.infer) => { + await rootsGate.waitForReady(); const validSourcePath = await validatePath(args.source); const validDestPath = await validatePath(args.destination); await fs.rename(validSourcePath, validDestPath); @@ -642,6 +660,7 @@ server.registerTool( annotations: { readOnlyHint: true } }, async (args: z.infer) => { + await rootsGate.waitForReady(); const validPath = await validatePath(args.path); const results = await searchFilesWithValidation(validPath, args.pattern, allowedDirectories, { excludePatterns: args.excludePatterns }); const text = results.length > 0 ? results.join("\n") : "No matches found"; @@ -668,6 +687,7 @@ server.registerTool( annotations: { readOnlyHint: true } }, async (args: z.infer) => { + await rootsGate.waitForReady(); const validPath = await validatePath(args.path); const info = await getFileStats(validPath); const text = Object.entries(info) @@ -694,6 +714,7 @@ server.registerTool( annotations: { readOnlyHint: true } }, async () => { + await rootsGate.waitForReady(); const text = `Allowed directories:\n${allowedDirectories.join('\n')}`; return { content: [{ type: "text" as const, text }], @@ -729,25 +750,29 @@ server.server.setNotificationHandler(RootsListChangedNotificationSchema, async ( // Handles post-initialization setup, specifically checking for and fetching MCP roots. server.server.oninitialized = async () => { - const clientCapabilities = server.server.getClientCapabilities(); - - if (clientCapabilities?.roots) { - try { - const response = await server.server.listRoots(); - if (response && 'roots' in response) { - await updateAllowedDirectoriesFromRoots(response.roots); + try { + const clientCapabilities = server.server.getClientCapabilities(); + + if (clientCapabilities?.roots) { + try { + const response = await server.server.listRoots(); + if (response && 'roots' in response) { + await updateAllowedDirectoriesFromRoots(response.roots); + } else { + console.error("Client returned no roots set, keeping current settings"); + } + } catch (error) { + console.error("Failed to request initial roots from client:", error instanceof Error ? error.message : String(error)); + } + } else { + if (allowedDirectories.length > 0) { + console.error("Client does not support MCP Roots, using allowed directories set from server args:", allowedDirectories); } else { - console.error("Client returned no roots set, keeping current settings"); + throw new Error(`Server cannot operate: No allowed directories available. Server was started without command-line directories and client either does not support MCP roots protocol or provided empty roots. Please either: 1) Start server with directory arguments, or 2) Use a client that supports MCP roots protocol and provides valid root directories.`); } - } catch (error) { - console.error("Failed to request initial roots from client:", error instanceof Error ? error.message : String(error)); - } - } else { - if (allowedDirectories.length > 0) { - console.error("Client does not support MCP Roots, using allowed directories set from server args:", allowedDirectories); - }else{ - throw new Error(`Server cannot operate: No allowed directories available. Server was started without command-line directories and client either does not support MCP roots protocol or provided empty roots. Please either: 1) Start server with directory arguments, or 2) Use a client that supports MCP roots protocol and provides valid root directories.`); } + } finally { + rootsGate.resolve(); } }; diff --git a/src/filesystem/roots-gate.ts b/src/filesystem/roots-gate.ts new file mode 100644 index 0000000000..aed0de89d3 --- /dev/null +++ b/src/filesystem/roots-gate.ts @@ -0,0 +1,41 @@ +/** + * Promise-based readiness gate that blocks tool handlers until roots/directories + * are initialized. Prevents race conditions where tool calls arrive before + * oninitialized has finished loading roots. + */ + +export interface RootsGate { + resolve: () => void; + waitForReady: () => Promise; +} + +export function createRootsGate(timeoutMs: number = 10000): RootsGate { + let resolve!: () => void; + const promise = new Promise((res) => { + resolve = res; + }); + + function waitForReady(): Promise { + return new Promise((resolveWait, rejectWait) => { + const timer = setTimeout(() => { + rejectWait(new Error(`Roots initialization timed out after ${timeoutMs}ms`)); + }, timeoutMs); + + promise.then( + () => { + clearTimeout(timer); + resolveWait(); + }, + (error) => { + clearTimeout(timer); + rejectWait(error); + } + ); + }); + } + + return { + resolve, + waitForReady, + }; +}