diff --git a/src/everything/__tests__/roots.test.ts b/src/everything/__tests__/roots.test.ts new file mode 100644 index 0000000000..f81467dc7c --- /dev/null +++ b/src/everything/__tests__/roots.test.ts @@ -0,0 +1,87 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { syncRoots, roots } from '../server/roots.js'; +import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; + +const sessionA = 'session-a'; + +const makeMockServer = (overrides: { + listRoots: () => Promise; + capabilities?: Record; +}) => { + const setNotificationHandler = vi.fn(); + const sendLoggingMessage = vi.fn().mockResolvedValue(undefined); + const listRoots = vi.fn(overrides.listRoots); + const capabilities = overrides.capabilities ?? { roots: {} }; + + const mockServer = { + sendLoggingMessage, + server: { + getClientCapabilities: vi.fn(() => capabilities), + listRoots, + setNotificationHandler, + }, + } as unknown as McpServer; + + return { mockServer, listRoots, setNotificationHandler, sendLoggingMessage }; +}; + +describe('syncRoots', () => { + beforeEach(() => { + roots.clear(); + }); + + it('caches roots and does not re-fetch on subsequent calls', async () => { + const { mockServer, listRoots, setNotificationHandler } = makeMockServer({ + listRoots: () => Promise.resolve({ roots: [{ uri: 'file:///a', name: 'a' }] }), + }); + + const first = await syncRoots(mockServer, sessionA); + const second = await syncRoots(mockServer, sessionA); + + expect(first).toEqual([{ uri: 'file:///a', name: 'a' }]); + expect(second).toEqual([{ uri: 'file:///a', name: 'a' }]); + expect(listRoots).toHaveBeenCalledTimes(1); + expect(setNotificationHandler).toHaveBeenCalledTimes(1); + }); + + it('does not re-fetch on subsequent calls when listRoots throws', async () => { + const { mockServer, listRoots, setNotificationHandler } = makeMockServer({ + listRoots: () => Promise.reject(new Error('boom')), + }); + + await syncRoots(mockServer, sessionA); + await syncRoots(mockServer, sessionA); + await syncRoots(mockServer, sessionA); + + // The doc claims this function is idempotent and fetches once per session. + // After a failed fetch, subsequent calls must not retry the listRoots + // request, otherwise every tool call hits the client again. + expect(listRoots).toHaveBeenCalledTimes(1); + expect(setNotificationHandler).toHaveBeenCalledTimes(1); + }); + + it('does not re-fetch on subsequent calls when listRoots returns falsy or unshaped', async () => { + const { mockServer, listRoots, setNotificationHandler } = makeMockServer({ + listRoots: () => Promise.resolve(undefined), + }); + + await syncRoots(mockServer, sessionA); + await syncRoots(mockServer, sessionA); + + expect(listRoots).toHaveBeenCalledTimes(1); + expect(setNotificationHandler).toHaveBeenCalledTimes(1); + }); + + it('returns undefined and does not call listRoots when client lacks roots capability', async () => { + const { mockServer, listRoots, setNotificationHandler } = makeMockServer({ + listRoots: () => Promise.resolve({ roots: [] }), + capabilities: {}, + }); + + const result = await syncRoots(mockServer, sessionA); + + expect(result).toBeUndefined(); + expect(listRoots).not.toHaveBeenCalled(); + expect(setNotificationHandler).not.toHaveBeenCalled(); + }); +}); diff --git a/src/everything/server/roots.ts b/src/everything/server/roots.ts index 34b12b21ce..89de797b9f 100644 --- a/src/everything/server/roots.ts +++ b/src/everything/server/roots.ts @@ -11,7 +11,7 @@ export const roots: Map = new Map< >(); /** - * Get the latest the client roots list for the session. + * Get the latest client roots list for the session. * * - Request and cache the roots list for the session if it has not been fetched before. * - Return the cached roots list for the session if it exists. @@ -20,71 +20,79 @@ export const roots: Map = new Map< * notification handler. This ensures that updates are automatically fetched and handled * in real-time. * - * This function is idempotent. It should only request roots from the client once per session, - * returning the cached version thereafter. + * This function is idempotent. It requests roots from the client at most once per session, + * even if the initial request fails or the client returns no roots; callers get the + * cached result (possibly an empty list) on subsequent calls. Later `roots/list_changed` + * notifications are the only way the cache is refreshed. * * @param {McpServer} server - An instance of the MCP server used to communicate with the client. * @param {string} [sessionId] - An optional session id used to associate the roots list with a specific client session. - * - * @throws {Error} In case of a failure to request the roots from the client, an error log message is sent. */ export const syncRoots = async (server: McpServer, sessionId?: string) => { const clientCapabilities = server.server.getClientCapabilities() || {}; const clientSupportsRoots: boolean = clientCapabilities?.roots !== undefined; - // Fetch the roots list for this client - if (clientSupportsRoots) { - // Function to request the updated roots list from the client - const requestRoots = async () => { - try { - // Request the updated roots list from the client - const response = await server.server.listRoots(); - if (response && "roots" in response) { - // Store the roots list for this client - roots.set(sessionId, response.roots); + if (!clientSupportsRoots) { + return; + } - // Notify the client of roots received - await server.sendLoggingMessage( - { - level: "info", - logger: "everything-server", - data: `Roots updated: ${response?.roots?.length} root(s) received from client`, - }, - sessionId - ); - } else { - await server.sendLoggingMessage( - { - level: "info", - logger: "everything-server", - data: "Client returned no roots set", - }, - sessionId - ); - } - } catch (error) { - console.error( - `Failed to request roots from client ${sessionId}: ${ - error instanceof Error ? error.message : String(error) - }` + // Function to request the updated roots list from the client + const requestRoots = async () => { + try { + const response = await server.server.listRoots(); + if (response && "roots" in response) { + roots.set(sessionId, response.roots); + await server.sendLoggingMessage( + { + level: "info", + logger: "everything-server", + data: `Roots updated: ${response?.roots?.length} root(s) received from client`, + }, + sessionId + ); + } else { + // Client returned no usable roots. Cache an empty list so the session + // is marked as synced and we don't re-fetch on every tool call. + roots.set(sessionId, []); + await server.sendLoggingMessage( + { + level: "info", + logger: "everything-server", + data: "Client returned no roots set", + }, + sessionId ); } - }; - - // If the roots have not been synced for this client, - // set notification handler and request initial roots - if (!roots.has(sessionId)) { - // Set the list changed notification handler - server.server.setNotificationHandler( - RootsListChangedNotificationSchema, - requestRoots + } catch (error) { + // Keep the session marked as synced even on failure. Otherwise every + // subsequent tool call would re-invoke listRoots and re-register the + // list_changed notification handler. + if (!roots.has(sessionId)) { + roots.set(sessionId, []); + } + console.error( + `Failed to request roots from client ${sessionId}: ${ + error instanceof Error ? error.message : String(error) + }` ); - - // Request the initial roots list immediately - await requestRoots(); } + }; + + // If the roots have not been synced for this client, + // set notification handler and request initial roots. + if (!roots.has(sessionId)) { + // Mark the session as in-flight before awaiting the request so that a + // re-entrant call (e.g. a tool invoked while the initial fetch is still + // pending) sees a cached entry and doesn't install a second handler. + roots.set(sessionId, []); - // Return the roots list for this client - return roots.get(sessionId); + server.server.setNotificationHandler( + RootsListChangedNotificationSchema, + requestRoots + ); + + await requestRoots(); } + + return roots.get(sessionId); };