diff --git a/examples/pdf-server/main.ts b/examples/pdf-server/main.ts index bb7b57f42..a3e6def61 100644 --- a/examples/pdf-server/main.ts +++ b/examples/pdf-server/main.ts @@ -21,6 +21,7 @@ import { fileUrlToPath, allowedLocalFiles, cliLocalFiles, + cliLocalDirs, DEFAULT_PDF, allowedLocalDirs, writeFlags, @@ -162,6 +163,7 @@ async function main() { console.error(`[pdf-server] Registered local file: ${filePath}`); } else if (s.isDirectory()) { allowedLocalDirs.add(filePath); + cliLocalDirs.add(filePath); console.error(`[pdf-server] Registered local directory: ${filePath}`); } } else { diff --git a/examples/pdf-server/server.test.ts b/examples/pdf-server/server.test.ts index 0b2a5eb4d..c842c029b 100644 --- a/examples/pdf-server/server.test.ts +++ b/examples/pdf-server/server.test.ts @@ -20,6 +20,7 @@ import { startFileWatch, stopFileWatch, cliLocalFiles, + cliLocalDirs, isWritablePath, writeFlags, viewSourcePaths, @@ -786,6 +787,106 @@ describe("createServer useClientRoots option", () => { expect(server.server.oninitialized).toBeFunction(); server.close(); }); + + describe("refreshRoots against a live client", () => { + let savedDirs: Set; + let savedCliDirs: Set; + const tmpDirs: string[] = []; + + const mkTmpDir = () => { + const dir = fs.mkdtempSync(path.join(os.tmpdir(), "pdf-roots-")); + tmpDirs.push(dir); + return dir; + }; + + beforeEach(() => { + savedDirs = new Set(allowedLocalDirs); + savedCliDirs = new Set(cliLocalDirs); + allowedLocalDirs.clear(); + cliLocalDirs.clear(); + }); + + afterEach(() => { + allowedLocalDirs.clear(); + cliLocalDirs.clear(); + for (const d of savedDirs) allowedLocalDirs.add(d); + for (const d of savedCliDirs) cliLocalDirs.add(d); + while (tmpDirs.length) { + fs.rmSync(tmpDirs.pop()!, { recursive: true, force: true }); + } + }); + + // Poll until cond() holds; refreshRoots runs fire-and-forget from + // oninitialized / the list_changed notification handler, so the only + // observable signal is the allowedLocalDirs set itself. + async function waitFor(cond: () => boolean, ms = 3000): Promise { + const deadline = Date.now() + ms; + while (!cond()) { + if (Date.now() > deadline) throw new Error("waitFor timed out"); + await new Promise((r) => setTimeout(r, 10)); + } + } + + async function connectWithRoots( + getRoots: () => Array<{ uri: string }>, + ): Promise<{ server: ReturnType; client: Client }> { + const server = createServer({ useClientRoots: true }); + const client = new Client( + { name: "t", version: "1" }, + { capabilities: { roots: { listChanged: true } } }, + ); + client.setRequestHandler("roots/list", async () => ({ + roots: getRoots(), + })); + const [ct, st] = InMemoryTransport.createLinkedPair(); + await Promise.all([server.connect(st), client.connect(ct)]); + return { server, client }; + } + + it("keeps CLI-arg directories when the client advertises roots but returns an empty list", async () => { + const cliDir = mkTmpDir(); + const staleDir = mkTmpDir(); + // Mirrors main(): a CLI-arg directory lands in both sets. + cliLocalDirs.add(cliDir); + allowedLocalDirs.add(cliDir); + // A directory granted by an earlier roots refresh — the next refresh + // is expected to revoke it, which also proves refreshRoots ran. + allowedLocalDirs.add(staleDir); + + const { server, client } = await connectWithRoots(() => []); + try { + await waitFor(() => !allowedLocalDirs.has(staleDir)); + expect(allowedLocalDirs.has(cliDir)).toBe(true); + } finally { + await client.close(); + await server.close(); + } + }); + + it("still revokes a removed client root while keeping CLI-arg directories", async () => { + const cliDir = mkTmpDir(); + const rootDir = mkTmpDir(); + cliLocalDirs.add(cliDir); + allowedLocalDirs.add(cliDir); + + let roots: Array<{ uri: string }> = [{ uri: pathToFileUrl(rootDir) }]; + const { server, client } = await connectWithRoots(() => roots); + try { + // Initial refresh on init grants the client root. + await waitFor(() => allowedLocalDirs.has(rootDir)); + + // Client drops the root: the refresh must revoke it but leave the + // CLI-arg directory alone. + roots = []; + await client.sendRootsListChanged(); + await waitFor(() => !allowedLocalDirs.has(rootDir)); + expect(allowedLocalDirs.has(cliDir)).toBe(true); + } finally { + await client.close(); + await server.close(); + } + }); + }); }); describe("isWritablePath", () => { diff --git a/examples/pdf-server/server.ts b/examples/pdf-server/server.ts index 629ace85f..2f0fec376 100644 --- a/examples/pdf-server/server.ts +++ b/examples/pdf-server/server.ts @@ -100,6 +100,15 @@ export const allowedLocalDirs = new Set(); */ export const cliLocalFiles = new Set(); +/** + * Subset of allowedLocalDirs that came from CLI args (not MCP roots). + * refreshRoots() clears allowedLocalDirs on every refresh so the client can + * revoke previously-granted roots; entries in this set are re-added after + * each clear so a client that advertises the roots capability but returns + * an empty list cannot wipe directories the user passed on the command line. + */ +export const cliLocalDirs = new Set(); + /** * Write-permission flags. Object wrapper (not a bare `let`) so main.ts can * mutate via the exported binding without re-import gymnastics — same @@ -958,6 +967,10 @@ async function refreshRoots(server: Server): Promise { try { const { roots } = await server.listRoots(); allowedLocalDirs.clear(); + // CLI-arg directories are not roots — the client can't revoke them. + for (const dir of cliLocalDirs) { + allowedLocalDirs.add(dir); + } for (const root of roots) { if (isFileUrl(root.uri)) { const dir = fileUrlToPath(root.uri);