Skip to content
Closed
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/infrastructure/config.ts
Original file line number Diff line number Diff line change
Expand Up @@ -147,6 +147,7 @@ export const DEFAULTS = {
implementations: 50,
interfaces: 50,
},
disabledTools: [] as string[],
},
} satisfies CodegraphConfig;

Expand Down
13 changes: 11 additions & 2 deletions src/mcp/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -163,11 +163,17 @@ function createCallToolHandler(
customDbPath: string | undefined,
allowedRepos: string[] | undefined,
getQueries: () => Promise<unknown>,
enabledToolNames: Set<string>,
) {
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 };
}
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 validateMultiRepoAccess runs before disabled-tool guard

validateMultiRepoAccess is called on line 171 before the enabledToolNames check on line 173. In multi-repo mode, if a caller omits the required repo argument when invoking a disabled tool, the handler will throw a multi-repo access error instead of returning the expected "Unknown tool: <name>" response. Swapping the two checks keeps the disabled-tool path clean.

Suggested change
validateMultiRepoAccess(multiRepo, name, args);
if (!enabledToolNames.has(name)) {
return { content: [{ type: 'text', text: `Unknown tool: ${name}` }], isError: true };
}
if (!enabledToolNames.has(name)) {
return { content: [{ type: 'text', text: `Unknown tool: ${name}` }], isError: true };
}
validateMultiRepoAccess(multiRepo, name, args);

Fix in Claude Code


const dbPath = await resolveDbPath(customDbPath, args, allowedRepos);

const toolEntry = TOOL_HANDLERS.get(name);
Expand Down Expand Up @@ -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();
Expand All @@ -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)();
Expand Down
25 changes: 20 additions & 5 deletions src/mcp/tool-registry.ts
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,14 @@ const PAGINATION_PROPS: Record<string, unknown> = {
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<string> {
return new Set((disabledTools || []).map((name) => normalizeToolName(name)).filter(Boolean));
}

const BASE_TOOLS: ToolSchema[] = [
{
name: 'query',
Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions src/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1201,6 +1201,7 @@ export interface CodegraphConfig {

mcp: {
defaults: McpDefaults;
disabledTools: string[];
};
Comment on lines 1202 to 1205
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 disabledTools typed as required but treated optionally at runtime

disabledTools: string[] is non-optional in CodegraphConfig, yet server.ts guards it with config.mcp?.disabledTools ? … : undefined. Any user-supplied config that includes an mcp block without disabledTools (e.g., configs predating this PR) will satisfy the type only because deep-merge with DEFAULTS fills the gap — but if the merge is ever a shallow spread the field can be absent at runtime while TypeScript believes it is always present. Making it optional (disabledTools?: string[]) would align the type with the actual runtime handling and avoid surprising TypeScript errors for callers that construct a partial CodegraphConfig.

Suggested change
mcp: {
defaults: McpDefaults;
disabledTools: string[];
};
mcp: {
defaults: McpDefaults;
disabledTools?: string[];
};

Fix in Claude Code

}

Expand Down
1 change: 1 addition & 0 deletions tests/unit/config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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([]);
});
});

Expand Down
87 changes: 87 additions & 0 deletions tests/unit/mcp.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 ────────────────────────────────────
Expand Down Expand Up @@ -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<typeof import('../../src/infrastructure/config.js')>();
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 = {};

Expand Down
Loading