Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
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
48 changes: 48 additions & 0 deletions __tests__/mcp-tool-allowlist.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,14 +5,19 @@
*/
import { describe, it, expect, afterEach } from 'vitest';
import { ToolHandler } from '../src/mcp/tools';
import type CodeGraph from '../src';

const ENV = 'CODEGRAPH_MCP_TOOLS';
const REQUIRE_PROJECT_PATH_ENV = 'CODEGRAPH_MCP_REQUIRE_PROJECT_PATH';

describe('CODEGRAPH_MCP_TOOLS allowlist', () => {
const original = process.env[ENV];
const originalRequireProjectPath = process.env[REQUIRE_PROJECT_PATH_ENV];
afterEach(() => {
if (original === undefined) delete process.env[ENV];
else process.env[ENV] = original;
if (originalRequireProjectPath === undefined) delete process.env[REQUIRE_PROJECT_PATH_ENV];
else process.env[REQUIRE_PROJECT_PATH_ENV] = originalRequireProjectPath;
});

const listed = () => new ToolHandler(null).getTools().map(t => t.name).sort();
Expand Down Expand Up @@ -61,3 +66,46 @@ describe('CODEGRAPH_MCP_TOOLS allowlist', () => {
expect(res.content[0].text).not.toMatch(/disabled via CODEGRAPH_MCP_TOOLS/);
});
});

describe('CODEGRAPH_MCP_REQUIRE_PROJECT_PATH', () => {
const original = process.env[REQUIRE_PROJECT_PATH_ENV];
const cgStub = {
getStats: () => ({ fileCount: 10 }),
} as unknown as CodeGraph;

afterEach(() => {
if (original === undefined) delete process.env[REQUIRE_PROJECT_PATH_ENV];
else process.env[REQUIRE_PROJECT_PATH_ENV] = original;
delete process.env[ENV];
});

it('marks projectPath required in exposed tool schemas', () => {
process.env[REQUIRE_PROJECT_PATH_ENV] = 'true';
const explore = new ToolHandler(null).getTools().find((tool) => tool.name === 'codegraph_explore');
expect(explore?.inputSchema.required).toContain('query');
expect(explore?.inputSchema.required).toContain('projectPath');
});

it('marks projectPath required even when a default project exists', () => {
process.env[REQUIRE_PROJECT_PATH_ENV] = 'true';
const explore = new ToolHandler(cgStub).getTools().find((tool) => tool.name === 'codegraph_explore');
expect(explore?.inputSchema.required).toContain('query');
expect(explore?.inputSchema.required).toContain('projectPath');
});

it('rejects tool calls that omit projectPath when required', async () => {
process.env[REQUIRE_PROJECT_PATH_ENV] = 'true';
const res = await new ToolHandler(null).execute('codegraph_explore', { query: 'alpha' });
expect(res.isError).toBe(true);
expect(res.content[0].text).toMatch(/projectPath is required/);
});

it('keeps normal validation after projectPath is provided', async () => {
process.env[REQUIRE_PROJECT_PATH_ENV] = 'true';
const res = await new ToolHandler(null).execute('codegraph_explore', {
query: 'alpha',
projectPath: '/tmp/project',
});
expect(res.content[0].text).not.toMatch(/projectPath is required/);
});
});
32 changes: 28 additions & 4 deletions src/mcp/tools.ts
Original file line number Diff line number Diff line change
Expand Up @@ -712,6 +712,24 @@ export function getStaticTools(): ToolDefinition[] {
* untouched, and `CODEGRAPH_MCP_TOOLS=explore,node,...` re-enables any of them.
*/
const DEFAULT_MCP_TOOLS = new Set(['explore']);
const REQUIRE_PROJECT_PATH_ENV = 'CODEGRAPH_MCP_REQUIRE_PROJECT_PATH';

function requiresProjectPath(): boolean {
return /^(1|true|yes|on)$/i.test(process.env[REQUIRE_PROJECT_PATH_ENV]?.trim() ?? '');
}

function withEnvRequiredProjectPath(tool: ToolDefinition): ToolDefinition {
if (!requiresProjectPath()) return tool;
const required = new Set(tool.inputSchema.required ?? []);
required.add('projectPath');
return {
...tool,
inputSchema: {
...tool.inputSchema,
required: [...required],
},
};
}

/**
* Tool handler that executes tools against a CodeGraph instance
Expand Down Expand Up @@ -916,15 +934,15 @@ export class ToolHandler {

return visible.map(tool => {
if (tool.name === 'codegraph_explore') {
return {
return withEnvRequiredProjectPath({
...tool,
description: `${tool.description} Budget: make at most ${budget} calls for this project (${stats.fileCount.toLocaleString()} files indexed).`,
};
});
}
return tool;
return withEnvRequiredProjectPath(tool);
});
} catch {
return visible;
return visible.map(withEnvRequiredProjectPath);
}
}

Expand Down Expand Up @@ -1281,6 +1299,12 @@ export class ToolHandler {
if (!this.isToolAllowed(toolName)) {
return this.errorResult(`Tool ${toolName} is disabled via CODEGRAPH_MCP_TOOLS`);
}
if (requiresProjectPath() && (args.projectPath === undefined || args.projectPath === null)) {
return this.errorResult(
`projectPath is required because ${REQUIRE_PROJECT_PATH_ENV}=true. ` +
'Pass an absolute project path, or unset the env var to allow the session default project.'
);
}
// Cross-cutting input validation. All tools accept an optional
// `projectPath` and most accept either `query`, `task`, or
// `symbol` — bound their lengths centrally so individual handlers
Expand Down