From bb6708927ee2fa0eb572f59b522729f4c0dda797 Mon Sep 17 00:00:00 2001 From: Csaba Kertesz Date: Mon, 27 Apr 2026 22:19:02 +0200 Subject: [PATCH] Implementated selective disable of tools to spare with the schema size --- src/infrastructure/config.ts | 1 + src/mcp/server.ts | 13 +++++- src/mcp/tool-registry.ts | 25 ++++++++--- src/types.ts | 1 + tests/unit/config.test.ts | 1 + tests/unit/mcp.test.ts | 87 ++++++++++++++++++++++++++++++++++++ 6 files changed, 121 insertions(+), 7 deletions(-) diff --git a/src/infrastructure/config.ts b/src/infrastructure/config.ts index 523c16af..46efb1a6 100644 --- a/src/infrastructure/config.ts +++ b/src/infrastructure/config.ts @@ -147,6 +147,7 @@ export const DEFAULTS = { implementations: 50, interfaces: 50, }, + disabledTools: [] as string[], }, } satisfies CodegraphConfig; diff --git a/src/mcp/server.ts b/src/mcp/server.ts index ae44c001..2fe5fd7a 100644 --- a/src/mcp/server.ts +++ b/src/mcp/server.ts @@ -163,11 +163,17 @@ function createCallToolHandler( customDbPath: string | undefined, allowedRepos: string[] | undefined, getQueries: () => Promise, + enabledToolNames: Set, ) { return async (request: any) => { const { name, arguments: args } = request.params; try { validateMultiRepoAccess(multiRepo, name, args); + + if (!enabledToolNames.has(name)) { + return { content: [{ type: 'text', text: `Unknown tool: ${name}` }], isError: true }; + } + const dbPath = await resolveDbPath(customDbPath, args, allowedRepos); const toolEntry = TOOL_HANDLERS.get(name); @@ -209,6 +215,9 @@ export async function startMCPServer( // Apply config-based MCP page-size overrides const config = options.config || loadConfig(); initMcpDefaults(config.mcp?.defaults ? { ...config.mcp.defaults } : undefined); + const disabledTools = config.mcp?.disabledTools ? [...config.mcp.disabledTools] : undefined; + const enabledTools = buildToolList(multiRepo, disabledTools); + const enabledToolNames = new Set(enabledTools.map((tool) => tool.name)); const { Server, StdioServerTransport, ListToolsRequestSchema, CallToolRequestSchema } = await loadMCPSdk(); @@ -225,12 +234,12 @@ export async function startMCPServer( ); server.setRequestHandler(ListToolsRequestSchema, async () => ({ - tools: buildToolList(multiRepo), + tools: enabledTools, })); server.setRequestHandler( CallToolRequestSchema, - createCallToolHandler(multiRepo, customDbPath, allowedRepos, getQueries), + createCallToolHandler(multiRepo, customDbPath, allowedRepos, getQueries, enabledToolNames), ); const transport = new (StdioServerTransport as any)(); diff --git a/src/mcp/tool-registry.ts b/src/mcp/tool-registry.ts index 07b71541..3f102e26 100644 --- a/src/mcp/tool-registry.ts +++ b/src/mcp/tool-registry.ts @@ -29,6 +29,14 @@ const PAGINATION_PROPS: Record = { offset: { type: 'number', description: 'Skip this many results (pagination, default: 0)' }, }; +function normalizeToolName(name: string): string { + return name.trim().toLowerCase().replace(/^codegraph\d+_/, ''); +} + +function buildDisabledToolSet(disabledTools?: string[]): Set { + return new Set((disabledTools || []).map((name) => normalizeToolName(name)).filter(Boolean)); +} + const BASE_TOOLS: ToolSchema[] = [ { name: 'query', @@ -849,18 +857,25 @@ const LIST_REPOS_TOOL: ToolSchema = { /** * Build the tool list based on multi-repo mode. */ -export function buildToolList(multiRepo: boolean): ToolSchema[] { - if (!multiRepo) return BASE_TOOLS; - return [ - ...BASE_TOOLS.map((tool) => ({ +export function buildToolList(multiRepo: boolean, disabledTools?: string[]): ToolSchema[] { + const disabled = buildDisabledToolSet(disabledTools); + const includeTool = (tool: ToolSchema): boolean => !disabled.has(normalizeToolName(tool.name)); + const baseTools = BASE_TOOLS.filter(includeTool); + + if (!multiRepo) return baseTools; + + const tools: ToolSchema[] = [ + ...baseTools.map((tool) => ({ ...tool, inputSchema: { ...tool.inputSchema, properties: { ...tool.inputSchema.properties, ...REPO_PROP }, }, })), - LIST_REPOS_TOOL, ]; + + if (includeTool(LIST_REPOS_TOOL)) tools.push(LIST_REPOS_TOOL); + return tools; } // Backward-compatible export: full multi-repo tool list diff --git a/src/types.ts b/src/types.ts index b5614c73..ebfe1368 100644 --- a/src/types.ts +++ b/src/types.ts @@ -1201,6 +1201,7 @@ export interface CodegraphConfig { mcp: { defaults: McpDefaults; + disabledTools: string[]; }; } diff --git a/tests/unit/config.test.ts b/tests/unit/config.test.ts index a2b50401..ebb87be4 100644 --- a/tests/unit/config.test.ts +++ b/tests/unit/config.test.ts @@ -137,6 +137,7 @@ describe('DEFAULTS', () => { it('has mcp defaults', () => { expect(DEFAULTS.mcp.defaults.list_functions).toBe(100); expect(DEFAULTS.mcp.defaults.fn_impact).toBe(5); + expect(DEFAULTS.mcp.disabledTools).toEqual([]); }); }); diff --git a/tests/unit/mcp.test.ts b/tests/unit/mcp.test.ts index c3a1a147..752c6935 100644 --- a/tests/unit/mcp.test.ts +++ b/tests/unit/mcp.test.ts @@ -255,6 +255,20 @@ describe('buildToolList', () => { expect(tool.inputSchema.properties.repo.type).toBe('string'); } }); + + it('removes disabled tools from schema in single-repo mode', () => { + const tools = buildToolList(false, ['execution_flow', 'module_map']); + const names = tools.map((t) => t.name); + expect(names).not.toContain('execution_flow'); + expect(names).not.toContain('module_map'); + }); + + it('supports prefixed disabled tool names and can disable list_repos', () => { + const tools = buildToolList(true, ['codegraph2_module_map', 'list_repos']); + const names = tools.map((t) => t.name); + expect(names).not.toContain('module_map'); + expect(names).not.toContain('list_repos'); + }); }); // ─── startMCPServer handler logic ──────────────────────────────────── @@ -335,6 +349,79 @@ describe('startMCPServer handler dispatch', () => { expect(unknownResult.content[0].text).toContain('Unknown tool'); }); + it('applies config.mcp.disabledTools to list and call handlers', async () => { + const handlers = {}; + + vi.doMock('@modelcontextprotocol/sdk/server/index.js', () => ({ + Server: class MockServer { + setRequestHandler(name, handler) { + handlers[name] = handler; + } + async connect() {} + }, + })); + vi.doMock('@modelcontextprotocol/sdk/server/stdio.js', () => ({ + StdioServerTransport: class MockTransport {}, + })); + vi.doMock('@modelcontextprotocol/sdk/types.js', () => ({ + ListToolsRequestSchema: 'tools/list', + CallToolRequestSchema: 'tools/call', + })); + + vi.doMock('../../src/infrastructure/config.js', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + loadConfig: vi.fn(() => ({ + ...actual.DEFAULTS, + mcp: { + ...actual.DEFAULTS.mcp, + disabledTools: ['module_map'], + }, + })), + }; + }); + + vi.doMock('../../src/domain/queries.js', () => ({ + EVERY_SYMBOL_KIND: [], + EVERY_EDGE_KIND: [], + VALID_ROLES: [], + diffImpactMermaid: vi.fn(), + impactAnalysisData: vi.fn(() => ({ file: 'test', sources: [] })), + moduleMapData: vi.fn(() => ({ topNodes: [], stats: {} })), + fileDepsData: vi.fn(() => ({ file: 'test', results: [] })), + fnDepsData: vi.fn(() => ({ name: 'test', results: [] })), + fnImpactData: vi.fn(() => ({ name: 'test', results: [] })), + contextData: vi.fn(() => ({ name: 'test', results: [] })), + childrenData: vi.fn(() => ({ name: 'test', results: [] })), + explainData: vi.fn(() => ({ target: 'test', kind: 'function', results: [] })), + exportsData: vi.fn(() => ({ + file: 'test', + results: [], + reexports: [], + totalExported: 0, + totalInternal: 0, + })), + whereData: vi.fn(() => ({ target: 'test', mode: 'symbol', results: [] })), + diffImpactData: vi.fn(() => ({ changedFiles: 0, affectedFunctions: [] })), + listFunctionsData: vi.fn(() => ({ count: 0, functions: [] })), + rolesData: vi.fn(() => ({ count: 0, summary: {}, symbols: [] })), + pathData: vi.fn(() => ({ from: 'a', to: 'b', found: false })), + })); + + const { startMCPServer } = await import('../../src/mcp/index.js'); + await startMCPServer('/tmp/test.db'); + + const toolsList = await handlers['tools/list'](); + expect(toolsList.tools.map((t) => t.name)).not.toContain('module_map'); + + const result = await handlers['tools/call']({ + params: { name: 'module_map', arguments: {} }, + }); + expect(result.isError).toBe(true); + expect(result.content[0].text).toContain('Unknown tool: module_map'); + }); + it('dispatches query deps mode to fnDepsData with options', async () => { const handlers = {};