diff --git a/.changeset/tidy-tools-fetch.md b/.changeset/tidy-tools-fetch.md new file mode 100644 index 00000000000..4650a0a55b5 --- /dev/null +++ b/.changeset/tidy-tools-fetch.md @@ -0,0 +1,5 @@ +--- +'@primer/mcp': minor +--- + +Components: Add a batch tool for retrieving documentation for multiple components in one call. diff --git a/package-lock.json b/package-lock.json index bcb66fcfd6f..9c04b265417 100644 --- a/package-lock.json +++ b/package-lock.json @@ -28221,6 +28221,7 @@ }, "devDependencies": { "@modelcontextprotocol/inspector": "^0.16.6", + "@primer/vitest-config": "^0.0.0", "@types/turndown": "^5.0.5", "rimraf": "^6.0.1", "rolldown": "^1.1.2", diff --git a/packages/mcp/package.json b/packages/mcp/package.json index 21d3a194e1e..1416dbaa027 100644 --- a/packages/mcp/package.json +++ b/packages/mcp/package.json @@ -29,6 +29,7 @@ "scripts": { "clean": "rimraf dist", "build": "rolldown -c", + "test": "vitest --run", "type-check": "tsc --noEmit", "watch": "rolldown -c -w" }, @@ -43,6 +44,7 @@ }, "devDependencies": { "@modelcontextprotocol/inspector": "^0.16.6", + "@primer/vitest-config": "^0.0.0", "@types/turndown": "^5.0.5", "rimraf": "^6.0.1", "rolldown": "^1.1.2", diff --git a/packages/mcp/src/server.test.ts b/packages/mcp/src/server.test.ts new file mode 100644 index 00000000000..219a2b8de58 --- /dev/null +++ b/packages/mcp/src/server.test.ts @@ -0,0 +1,125 @@ +import {Client} from '@modelcontextprotocol/sdk/client/index.js' +import {InMemoryTransport} from '@modelcontextprotocol/sdk/inMemory.js' +import {afterAll, afterEach, beforeAll, beforeEach, describe, expect, it, vi} from 'vitest' +import {server} from './server' + +describe('get_component_batch', () => { + const client = new Client({name: 'mcp-server-test', version: '1.0.0'}) + + beforeAll(async () => { + const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair() + await server.connect(serverTransport) + await client.connect(clientTransport) + }) + + beforeEach(() => { + vi.stubGlobal('fetch', vi.fn()) + }) + + afterEach(() => { + vi.useRealTimers() + vi.unstubAllGlobals() + }) + + afterAll(async () => { + await client.close() + await server.close() + }) + + const callBatch = (names: string[]) => { + return client.callTool({ + name: 'get_component_batch', + arguments: {names}, + }) + } + + it('requires between 2 and 10 names', async () => { + const tooFew = await callBatch(['Button']) + const tooMany = await callBatch(Array.from({length: 11}, (_, index) => `Component${index}`)) + + expect(tooFew.isError).toBe(true) + expect(tooMany.isError).toBe(true) + expect(fetch).not.toHaveBeenCalled() + }) + + it('deduplicates names case-insensitively and preserves result order', async () => { + const fetchMock = vi.mocked(fetch) + fetchMock.mockImplementation(async input => { + const url = new URL(input.toString()) + return new Response(`${url.pathname} docs`) + }) + + const result = await callBatch(['Button', 'button', 'Dialog']) + + expect(fetchMock).toHaveBeenCalledTimes(2) + expect(result.content).toEqual([ + {type: 'text', text: '/product/components/button/llms.txt docs'}, + {type: 'text', text: '/product/components/dialog/llms.txt docs'}, + ]) + }) + + it('starts component fetches in parallel', async () => { + const fetchMock = vi.mocked(fetch) + const pendingResponses: Array<(response: Response) => void> = [] + fetchMock.mockImplementation( + () => + new Promise(resolve => { + pendingResponses.push(resolve) + }), + ) + + const resultPromise = callBatch(['Button', 'Dialog']) + await vi.waitFor(() => expect(fetchMock).toHaveBeenCalledTimes(2)) + + for (const resolve of pendingResponses) { + resolve(new Response('docs')) + } + + await expect(resultPromise).resolves.toMatchObject({ + content: [ + {type: 'text', text: 'docs'}, + {type: 'text', text: 'docs'}, + ], + }) + }) + + it('returns missing components and fetch failures in-band', async () => { + vi.mocked(fetch).mockResolvedValue(new Response(null, {status: 500})) + + const result = await callBatch(['MissingComponent', 'button']) + + expect(result.isError).not.toBe(true) + expect(result.content).toEqual([ + { + type: 'text', + text: '## MissingComponent\n\nStatus: not-found. No component named `MissingComponent` exists in @primer/react. Use `list_components` for valid names.', + }, + { + type: 'text', + text: '## Button\n\nStatus: fetch-error. Failed to load documentation for Button.', + }, + ]) + }) + + it('times out stalled requests without leaving timers active', async () => { + vi.useFakeTimers() + const fetchMock = vi.mocked(fetch) + fetchMock.mockImplementation( + (_input, init) => + new Promise((_resolve, reject) => { + init?.signal?.addEventListener('abort', () => reject(init.signal?.reason), {once: true}) + }), + ) + + const resultPromise = callBatch(['Button', 'Dialog']) + await vi.advanceTimersByTimeAsync(10_000) + + await expect(resultPromise).resolves.toMatchObject({ + content: [ + {type: 'text', text: '## Button\n\nStatus: fetch-error. Failed to load documentation for Button.'}, + {type: 'text', text: '## Dialog\n\nStatus: fetch-error. Failed to load documentation for Dialog.'}, + ], + }) + expect(vi.getTimerCount()).toBe(0) + }) +}) diff --git a/packages/mcp/src/server.ts b/packages/mcp/src/server.ts index 523dcd7540f..9dd19b6df11 100644 --- a/packages/mcp/src/server.ts +++ b/packages/mcp/src/server.ts @@ -161,6 +161,73 @@ server.registerTool( }, ) +server.registerTool( + 'get_component_batch', + { + description: + 'Retrieve documentation for multiple Primer React components in one call. ' + + 'Pass all component names you need at once. Docs are fetched in parallel to avoid repeated MCP round-trips when resolving several components. ' + + 'For a single component, prefer get_component instead.', + inputSchema: { + names: z + .array(z.string()) + .min(2) + .max(10) + .describe('Component names to retrieve (e.g. ["ActionMenu", "Button", "Pagination"])'), + }, + annotations: {readOnlyHint: true}, + }, + async ({names}) => { + const seenNames = new Set() + const deduped = names.filter(name => { + const normalizedName = name.toLowerCase() + if (seenNames.has(normalizedName)) { + return false + } + seenNames.add(normalizedName) + return true + }) + const components = listComponents() + + const results = await Promise.all( + deduped.map(async name => { + const match = components.find( + component => component.name === name || component.name.toLowerCase() === name.toLowerCase(), + ) + if (!match) { + return { + type: 'text' as const, + text: `## ${name}\n\nStatus: not-found. No component named \`${name}\` exists in @primer/react. Use \`list_components\` for valid names.`, + } + } + + const controller = new AbortController() + const timeout = setTimeout(() => controller.abort(), 10_000) + + try { + const llmsUrl = new URL(`/product/components/${match.slug}/llms.txt`, 'https://primer.style') + const llmsResponse = await fetch(llmsUrl, {signal: controller.signal}) + if (llmsResponse.ok) { + const text = await llmsResponse.text() + return {type: 'text' as const, text} + } + } catch (_: unknown) { + // Fall through to the in-band error result so other component requests can succeed. + } finally { + clearTimeout(timeout) + } + + return { + type: 'text' as const, + text: `## ${match.name}\n\nStatus: fetch-error. Failed to load documentation for ${match.name}.`, + } + }), + ) + + return {content: results} + }, +) + server.registerTool( 'get_component_examples', { diff --git a/packages/mcp/tsconfig.json b/packages/mcp/tsconfig.json index a84ca885fff..23093d59d49 100644 --- a/packages/mcp/tsconfig.json +++ b/packages/mcp/tsconfig.json @@ -1,4 +1,4 @@ { "extends": "../../tsconfig.base.json", - "include": ["src", "rolldown.config.ts"] + "include": ["src", "rolldown.config.ts", "vitest.config.ts"] } diff --git a/packages/mcp/vitest.config.ts b/packages/mcp/vitest.config.ts new file mode 100644 index 00000000000..ba6512fe614 --- /dev/null +++ b/packages/mcp/vitest.config.ts @@ -0,0 +1,8 @@ +import {defineConfig} from '@primer/vitest-config/config' + +export default defineConfig({ + test: { + environment: 'node', + detectAsyncLeaks: true, + }, +})