feat(mcp): codegraph_explore as the sole primary tool + store coverage + overload disambiguation (#647)

## Summary

Completes the explore-overhaul arc: `codegraph_explore` becomes the single primary tool an agent reaches for, and its coverage + output shape are tuned so flow/architecture questions resolve with near-zero Read/Grep.

### What changed
- **explore is the sole primary tool** — removed `codegraph_context` (the fuzzy-input Read-trigger) and `codegraph_trace` (under-picked by agents); explore already surfaces the call flow among the symbols you name. A plain natural-language question now works as the query.
- **Store/handler coverage** — functions defined inside object literals (Zustand `create((set, get) => ({ … }))`, Redux/Pinia/MobX, exported handler/route maps) are indexed as real symbols, including calls through `useStore.getState().fn()` and destructured `const { fn } = useStore.getState()`. A general AST rule, not a per-lib hack.
- **Overload disambiguation** — explore leads with the *right* definition when a method name is overloaded across types (a PascalCase type token in the query biases to that type's own def); `codegraph_node` returns *every* overload's body in one call, with an optional `file`/`line` selector to pin one.
- **Method-atomic render** — explore never returns half a method; at the size budget it drops whole methods/files (and lists what it dropped) instead of truncating a body mid-method.
- **Native-read-shaped output** — per-call output is capped to ~24K with a 25K hard ceiling and concentrated into ~150–250-line flow windows, mirroring how the agent natively reads; repo size scales the *call* budget, not the per-call size (a larger response just gets externalized to a file the host Reads back).
- **Blast radius** folded into explore (dependents + covering tests, locations only).

### Benchmark (refreshed on this build)
Re-validated the 7-repo A/B on 2026-06-02 (Opus 4.8, effort=high, median of 4). WITH arm re-measured on this build, WITHOUT reused:

**~16% cheaper · 47% fewer tokens · 22% faster · 58% fewer tool calls** — 0 file reads on 6 of 7 repos (Gin ~1).

The arc trades larger, cache-heavy explore responses for guaranteed near-zero reads, so cost/token margins soften vs the prior build (Excalidraw and Tokio land at cost break-even) while time and tool-calls stay clear wins everywhere — consistent with the project's stated optimization target (latency + tool-calls, not token cost).

### Validation
- Full suite green: **1112 passed, 2 skipped**.
- 28/28 plain WITH runs across the 7 README repos completed clean; reads median 0 on 6/7.

🤖 Generated with [Claude Code](https://claude.com/claude-code)
This commit is contained in:
Colby Mchenry
2026-06-02 10:15:27 -05:00
committed by GitHub
parent 8629f7ab4c
commit 68eaf0dbd8
27 changed files with 1471 additions and 1194 deletions
+123
View File
@@ -0,0 +1,123 @@
/**
* Context ranking: common-word precision + low-confidence handoff.
*
* Regression coverage for the failure where a prose query
* ("capture intro onboarding screen flat object") surfaced an unrelated
* constant named `FLAT` (in a download script) as a top entry point — because
* the descriptive word "flat" exact-matched it and the +exact-name bonus was
* exempt from single-term dampening. The fix: only distinctive identifiers earn
* that exemption; an isolated common-word exact match is demoted, and a query
* that resolves only to such weak matches is flagged low-confidence so the
* response hands off to explore/trace instead of bluffing.
*/
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import * as fs from 'fs';
import * as path from 'path';
import * as os from 'os';
import CodeGraph from '../src/index';
import { LOW_CONFIDENCE_MARKER } from '../src/context';
import { isDistinctiveIdentifier } from '../src/search/query-utils';
describe('isDistinctiveIdentifier', () => {
it('treats plain dictionary words as non-distinctive', () => {
for (const word of ['flat', 'object', 'screen', 'standing', 'capture']) {
expect(isDistinctiveIdentifier(word)).toBe(false);
}
});
it('treats leading-capital-only words (proper nouns / sentence start) as non-distinctive', () => {
expect(isDistinctiveIdentifier('Screen')).toBe(false);
expect(isDistinctiveIdentifier('Zustand')).toBe(false);
});
it('treats camelCase / PascalCase / snake_case / acronyms / digits as distinctive', () => {
expect(isDistinctiveIdentifier('setLastEmail')).toBe(true);
expect(isDistinctiveIdentifier('OrgUserStore')).toBe(true);
expect(isDistinctiveIdentifier('user_store')).toBe(true);
expect(isDistinctiveIdentifier('REST')).toBe(true);
expect(isDistinctiveIdentifier('v2')).toBe(true);
});
});
describe('Context ranking — common-word precision & confidence', () => {
let testDir: string;
let cg: CodeGraph;
beforeEach(async () => {
testDir = fs.mkdtempSync(path.join(os.tmpdir(), 'codegraph-ctxrank-'));
// The corroborated target: a capture-flow screen whose NAME alone matches
// three query terms (capture + intro + screen), and which lives under a
// matching directory.
const captureDir = path.join(testDir, 'src', 'app', 'capture');
fs.mkdirSync(captureDir, { recursive: true });
fs.writeFileSync(
path.join(captureDir, 'intro.tsx'),
`export function CaptureIntroScreen() {
// Onboarding screen shown before the user selects flat or standing object capture.
return null;
}
`
);
// The trap: an unrelated constant literally named FLAT, in a totally
// different area. "flat" in a prose query exact-matches it.
const scriptsDir = path.join(testDir, 'scripts', 'dataset');
fs.mkdirSync(scriptsDir, { recursive: true });
fs.writeFileSync(
path.join(scriptsDir, 'download.ts'),
`export const FLAT = 'freiburg_flat_dataset';
export function downloadDataset(name: string): string { return name; }
`
);
cg = CodeGraph.initSync(testDir, {
config: { include: ['**/*.ts', '**/*.tsx'], exclude: [] },
});
await cg.indexAll();
});
afterEach(() => {
if (cg) cg.destroy();
if (fs.existsSync(testDir)) fs.rmSync(testDir, { recursive: true, force: true });
});
it('does not let a common-word exact match (FLAT) outrank a corroborated symbol', async () => {
const sg = await cg.findRelevantContext(
'capture intro onboarding screen flat object'
);
const rootNames = sg.roots.map((id) => sg.nodes.get(id)?.name);
// The corroborated capture screen surfaces as an entry point...
expect(rootNames).toContain('CaptureIntroScreen');
// ...and the trap constant is never the lead result (the bug we fixed).
expect(rootNames[0]).not.toBe('FLAT');
const capIdx = rootNames.indexOf('CaptureIntroScreen');
const flatIdx = rootNames.indexOf('FLAT');
if (flatIdx >= 0) expect(capIdx).toBeLessThan(flatIdx);
// And it's confidently answered (we located a corroborated symbol).
expect(sg.confidence).toBe('high');
});
it('flags low confidence and emits the handoff when only common words match', async () => {
const query = 'flat object thing';
const sg = await cg.findRelevantContext(query);
expect(sg.confidence).toBe('low');
const md = await cg.buildContext(query, { format: 'markdown' });
expect(typeof md).toBe('string');
expect(md as string).toContain(LOW_CONFIDENCE_MARKER);
// The handoff routes to the precise tools rather than claiming completeness.
expect(md as string).toMatch(/codegraph_explore/);
});
it('does not emit the handoff for a precise, distinctive-symbol query', async () => {
const sg = await cg.findRelevantContext('CaptureIntroScreen');
expect(sg.confidence).toBe('high');
const md = await cg.buildContext('CaptureIntroScreen', { format: 'markdown' });
expect(md as string).not.toContain(LOW_CONFIDENCE_MARKER);
});
});
+73
View File
@@ -0,0 +1,73 @@
/**
* codegraph_explore blast-radius section.
*
* explore now appends a compact, always-on "Blast radius" for the entry
* symbols: who depends on each (locations only — no source) and which test
* files cover it, so the agent knows what to update/verify before editing
* without a separate impact call. Symbols with no dependents are skipped, and
* the section is omitted entirely when nothing qualifies.
*/
import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import * as fs from 'fs';
import * as path from 'path';
import * as os from 'os';
import CodeGraph from '../src/index';
import { ToolHandler } from '../src/mcp/tools';
describe('codegraph_explore — blast radius', () => {
let testDir: string;
let cg: CodeGraph;
let handler: ToolHandler;
beforeEach(async () => {
testDir = fs.mkdtempSync(path.join(os.tmpdir(), 'codegraph-blast-'));
const src = path.join(testDir, 'src');
fs.mkdirSync(src, { recursive: true });
// `target` is depended on by a sibling (caller) and a test file.
fs.writeFileSync(
path.join(src, 'feature.ts'),
`export function target() { return 1; }\n` +
`export function caller() { return target(); }\n`,
);
fs.writeFileSync(
path.join(src, 'feature.test.ts'),
`import { target } from './feature';\n` +
`export function checkTarget() { return target(); }\n`,
);
// A leaf with no dependents — must NOT show up in the blast radius.
fs.writeFileSync(
path.join(src, 'leaf.ts'),
`export function lonelyLeaf() { return 42; }\n`,
);
cg = CodeGraph.initSync(testDir, { config: { include: ['**/*.ts'], exclude: [] } });
await cg.indexAll();
handler = new ToolHandler(cg);
});
afterEach(() => {
if (cg) cg.destroy();
if (fs.existsSync(testDir)) fs.rmSync(testDir, { recursive: true, force: true });
});
it('lists dependents (locations only) and covering tests for an entry symbol', async () => {
const res = await handler.execute('codegraph_explore', { query: 'target' });
const text = res.content[0].text;
expect(text).toContain('### Blast radius');
expect(text).toContain('`target`');
expect(text).toMatch(/caller/); // a caller count is reported
// It names WHERE (the caller file) — not the caller's source body.
expect(text).toContain('feature.ts');
// Test coverage is surfaced (either the covering test file, or the warning).
expect(text).toMatch(/tests:.*feature\.test\.ts|no covering tests/);
});
it('omits symbols that have no dependents from the blast radius', async () => {
const res = await handler.execute('codegraph_explore', { query: 'lonelyLeaf' });
const text = res.content[0].text;
// lonelyLeaf has zero callers — it must never appear under a blast-radius bullet.
expect(text).not.toMatch(/Blast radius[\s\S]*`lonelyLeaf`/);
});
});
+14 -6
View File
@@ -27,9 +27,14 @@ describe('getExploreOutputBudget', () => {
expect(small.maxOutputChars).toBeLessThanOrEqual(20000);
});
it('keeps the historical 35k+ ceiling for medium-large projects so existing benchmarks do not regress', () => {
it('caps medium-large projects at the inline tool-result ceiling (~24k) so the result is never externalized', () => {
// A bigger single response gets externalized by the host to a file the agent
// Reads back (a 35k vscode explore did exactly that in the n=4 A/B) — adding a
// read AND cache-write cost. So large repos get MORE CALLS (getExploreBudget),
// not a fatter single response; the output cap stays under the inline limit.
const large = getExploreOutputBudget(10000);
expect(large.maxOutputChars).toBeGreaterThanOrEqual(35000);
expect(large.maxOutputChars).toBeLessThanOrEqual(25000);
expect(large.maxOutputChars).toBeGreaterThanOrEqual(20000);
});
it('uses tier breakpoints matching getExploreBudget so call-count and output-budget agree on a project', () => {
@@ -54,10 +59,13 @@ describe('getExploreOutputBudget', () => {
const tier3b = getExploreOutputBudget(14999);
expect(tier3a.maxOutputChars).toBe(tier3b.maxOutputChars);
// And crossing a breakpoint changes the cap.
expect(tier0a.maxOutputChars).not.toBe(tier1a.maxOutputChars);
expect(tier1a.maxOutputChars).not.toBe(tier2a.maxOutputChars);
expect(tier2a.maxOutputChars).not.toBe(tier3a.maxOutputChars);
// Small tiers step up (13k → 18k → 24k); medium and large SHARE the ~24k
// inline ceiling — scaling with repo size now lives in the CALL budget
// (getExploreBudget), not in a fatter single response.
expect(tier0a.maxOutputChars).not.toBe(tier1a.maxOutputChars); // <150 vs <500
expect(tier1a.maxOutputChars).not.toBe(tier2a.maxOutputChars); // <500 vs <5000
expect(tier2a.maxOutputChars).toBe(tier3a.maxOutputChars); // <5000 == <15000 (inline cap)
expect(getExploreBudget(5000)).toBeGreaterThan(getExploreBudget(4999)); // calls scale instead
});
it('gates off "Additional relevant files", completeness signal, and budget note on small projects', () => {
@@ -53,9 +53,9 @@ describe('MCP input size limits', () => {
expect(result.content[0]!.text).toMatch(/maximum length/i);
});
it('rejects an oversize task on codegraph_context', async () => {
it('rejects an oversize query on codegraph_explore', async () => {
const huge = 'b'.repeat(50_000);
const result = await handler.execute('codegraph_context', { task: huge });
const result = await handler.execute('codegraph_explore', { query: huge });
expect(result.isError).toBe(true);
expect(result.content[0]!.text).toMatch(/maximum length/i);
});
+9 -9
View File
@@ -21,28 +21,28 @@ describe('CODEGRAPH_MCP_TOOLS allowlist', () => {
delete process.env[ENV];
const all = listed();
expect(all).toContain('codegraph_explore');
expect(all).toContain('codegraph_context');
expect(all).toContain('codegraph_trace');
expect(all.length).toBeGreaterThanOrEqual(10);
expect(all).not.toContain('codegraph_context');
expect(all).not.toContain('codegraph_trace');
expect(all.length).toBeGreaterThanOrEqual(8);
});
it('filters ListTools to the allowlisted short names', () => {
process.env[ENV] = 'trace,search,node';
expect(listed()).toEqual(['codegraph_node', 'codegraph_search', 'codegraph_trace']);
process.env[ENV] = 'explore,search,node';
expect(listed()).toEqual(['codegraph_explore', 'codegraph_node', 'codegraph_search']);
});
it('accepts fully-qualified codegraph_ names and ignores whitespace', () => {
process.env[ENV] = ' codegraph_trace , search ';
expect(listed()).toEqual(['codegraph_search', 'codegraph_trace']);
process.env[ENV] = ' codegraph_explore , search ';
expect(listed()).toEqual(['codegraph_explore', 'codegraph_search']);
});
it('treats an empty/whitespace value as unset (full surface)', () => {
process.env[ENV] = ' ';
expect(listed().length).toBeGreaterThanOrEqual(10);
expect(listed().length).toBeGreaterThanOrEqual(8);
});
it('rejects a disabled tool on execute (defense in depth)', async () => {
process.env[ENV] = 'trace';
process.env[ENV] = 'node';
const res = await new ToolHandler(null).execute('codegraph_explore', {});
expect(res.isError).toBe(true);
expect(res.content[0].text).toMatch(/disabled via CODEGRAPH_MCP_TOOLS/);
+176
View File
@@ -0,0 +1,176 @@
/**
* Object-literal method extraction (general AST rule).
*
* The extractor pulls function-valued properties out of an object literal that
* is the value of an exported const — either DIRECTLY
* (`export const actions = { foo: () => {} }`) or RETURNED by an initializer
* call (`export const useStore = create((set, get) => ({ foo: () => {} }))`,
* incl. middleware wrappers). This makes store actions (Zustand/Redux/Pinia/
* MobX/handler maps) real nodes, so `codegraph_node`/`callers` on them resolve
* instead of returning "not found" and forcing the agent to Read the store.
*
* Keyed purely on AST shape — no library names in the implementation — so any
* same-shaped store is covered. Resolution then falls out of the existing
* exact-name matcher: every call form (`const {foo}=useStore.getState(); foo()`,
* `useStore.getState().foo()`, in-store `get().foo()`) reduces to a bare `foo`
* call that resolves to the action node once it exists.
*/
import { describe, it, expect, beforeAll, afterEach } from 'vitest';
import * as fs from 'fs';
import * as path from 'path';
import * as os from 'os';
import { CodeGraph } from '../src';
import { extractFromSource } from '../src/extraction';
import { initGrammars, loadAllGrammars } from '../src/extraction/grammars';
beforeAll(async () => {
await initGrammars();
await loadAllGrammars();
});
describe('object-literal method extraction', () => {
it('extracts Zustand store actions (object returned by create()) as function nodes', () => {
const code = `
import { create } from 'zustand'
interface Store {
count: number
fetchUser(): Promise<void>
switchOrganization(id: string): Promise<void>
reset(): void
}
export const useStore = create<Store>((set, get) => ({
count: 0,
fetchUser: async () => { await get().reset() },
switchOrganization: async (id: string) => { set({ count: 1 }) },
reset: () => set({ count: 0 }),
}))
`;
const result = extractFromSource('store.ts', code);
const fnNames = result.nodes.filter((n) => n.kind === 'function').map((n) => n.name);
expect(fnNames).toContain('fetchUser');
expect(fnNames).toContain('switchOrganization');
expect(fnNames).toContain('reset');
// Each action's body was walked: fetchUser references its sibling `reset`,
// so an in-store calls edge will resolve once the pipeline runs.
const fetchUser = result.nodes.find((n) => n.name === 'fetchUser')!;
const fetchUserRefs = result.unresolvedReferences.filter((r) => r.fromNodeId === fetchUser.id);
expect(fetchUserRefs.map((r) => r.referenceName)).toContain('reset');
// The action's body wasn't mis-attributed to the file scope (the reason we
// skip the generic body-visit for the store-factory call).
const fileNode = result.nodes.find((n) => n.kind === 'file')!;
const fileRefs = result.unresolvedReferences.filter((r) => r.fromNodeId === fileNode.id);
expect(fileRefs.map((r) => r.referenceName)).not.toContain('reset');
});
it('extracts actions through a middleware wrapper (create(persist(...)))', () => {
const code = `
import { create } from 'zustand'
import { persist } from 'zustand/middleware'
export const useCounter = create(
persist(
(set, get) => ({
value: 0,
increment: () => set({ value: get().value + 1 }),
}),
{ name: 'counter' }
)
)
`;
const result = extractFromSource('counter.ts', code);
const fnNames = result.nodes.filter((n) => n.kind === 'function').map((n) => n.name);
expect(fnNames).toContain('increment');
});
it('extracts actions when the initializer returns via a block (=> { return {...} })', () => {
const code = `
import { create } from 'zustand'
export const useThing = create((set) => {
const initial = 0
return {
value: initial,
bump: () => set({ value: 1 }),
}
})
`;
const result = extractFromSource('thing.ts', code);
const fnNames = result.nodes.filter((n) => n.kind === 'function').map((n) => n.name);
expect(fnNames).toContain('bump');
});
it('does NOT extract methods from a non-exported call-wrapped object (noise gate)', () => {
const code = `
function wrap(f: any) { return f }
const local = wrap(() => ({ shouldNotExtract: () => {} }))
`;
const result = extractFromSource('inline.ts', code);
const names = result.nodes.map((n) => n.name);
expect(names).not.toContain('shouldNotExtract');
});
it('still extracts the existing direct-object shape (export const actions = {...})', () => {
const code = `
export const actions = {
load: async () => { helper() },
}
function helper() {}
`;
const result = extractFromSource('actions.ts', code);
const fnNames = result.nodes.filter((n) => n.kind === 'function').map((n) => n.name);
expect(fnNames).toContain('load');
});
});
describe('object-literal method resolution (end-to-end)', () => {
let tmpDir: string | undefined;
afterEach(() => {
if (tmpDir) fs.rmSync(tmpDir, { recursive: true, force: true });
tmpDir = undefined;
});
it('resolves callers of store actions across files (destructured + chained getState())', async () => {
tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'cg-store-'));
fs.writeFileSync(path.join(tmpDir, 'package.json'), '{"name":"t","dependencies":{"zustand":"^4"}}\n');
fs.writeFileSync(
path.join(tmpDir, 'store.ts'),
`import { create } from 'zustand'\n` +
`interface S { fetchUser(): Promise<void>; reset(): void }\n` +
`export const useStore = create<S>((set, get) => ({\n` +
` fetchUser: async () => { get().reset() },\n` +
` reset: () => set({}),\n` +
`}))\n`
);
fs.writeFileSync(
path.join(tmpDir, 'caller.ts'),
`import { useStore } from './store'\n` +
`export async function loginFlow() {\n` +
` const { fetchUser } = useStore.getState()\n` +
` await fetchUser()\n` +
`}\n` +
`export function hardReset() {\n` +
` useStore.getState().reset()\n` +
`}\n`
);
const cg = CodeGraph.initSync(tmpDir);
await cg.indexAll();
const fns = cg.getNodesByKind('function');
const fetchUser = fns.find((n) => n.name === 'fetchUser' && n.filePath.endsWith('store.ts'));
const reset = fns.find((n) => n.name === 'reset' && n.filePath.endsWith('store.ts'));
expect(fetchUser).toBeDefined();
expect(reset).toBeDefined();
// Destructured-then-bare call: loginFlow -> fetchUser
const fetchUserCallers = cg.getCallers(fetchUser!.id).map((c) => c.node.name);
expect(fetchUserCallers).toContain('loginFlow');
// Chained getState() call: hardReset -> reset, AND in-store sibling: fetchUser -> reset
const resetCallers = cg.getCallers(reset!.id).map((c) => c.node.name);
expect(resetCallers).toContain('hardReset');
expect(resetCallers).toContain('fetchUser');
cg.close();
});
});
+19 -19
View File
@@ -501,10 +501,10 @@ describe('MCP Tool Improvements', () => {
expect(typeof ToolHandler).toBe('function');
});
it.skipIf(!HAS_SQLITE)('should have findSymbol and truncateOutput as private methods', async () => {
it.skipIf(!HAS_SQLITE)('should have findSymbolMatches and truncateOutput as private methods', async () => {
const { ToolHandler } = await import('../src/mcp/tools');
const proto = ToolHandler.prototype;
expect(typeof (proto as any).findSymbol).toBe('function');
expect(typeof (proto as any).findSymbolMatches).toBe('function');
expect(typeof (proto as any).truncateOutput).toBe('function');
});
@@ -567,20 +567,19 @@ export function getValueFromCache(): number { return 2; }
await cg.indexAll();
const handler = new ToolHandler(cg);
const findSymbol = (handler as any).findSymbol.bind(handler);
const findSymbolMatches = (handler as any).findSymbolMatches.bind(handler);
const match = findSymbol(cg, 'getValue');
expect(match).not.toBeNull();
expect(match.node.name).toBe('getValue');
// Should not have a disambiguation note for single exact match
expect(match.note).toBe('');
const matches = findSymbolMatches(cg, 'getValue');
// Exact-name match wins — a single result, not the partial getValueFromCache.
expect(matches.length).toBe(1);
expect(matches[0].name).toBe('getValue');
handler.closeAll();
cg.destroy();
cleanupTempDir(tmpDir);
});
it.skipIf(!HAS_SQLITE)('should note when multiple symbols share the same name', async () => {
it.skipIf(!HAS_SQLITE)('should return all definitions when multiple symbols share the same name', async () => {
const { ToolHandler } = await import('../src/mcp/tools');
const CodeGraph = (await import('../src/index')).default;
@@ -602,20 +601,21 @@ export function handle(): void {}
await cg.indexAll();
const handler = new ToolHandler(cg);
const findSymbol = (handler as any).findSymbol.bind(handler);
const findSymbolMatches = (handler as any).findSymbolMatches.bind(handler);
const match = findSymbol(cg, 'handle');
expect(match).not.toBeNull();
expect(match.node.name).toBe('handle');
// Should have a disambiguation note
expect(match.note).toContain('2 symbols named "handle"');
// Both same-named definitions are returned (no longer one + a dead-end
// note) so codegraph_node can hand back every overload and the agent never
// Reads to find the one it wanted.
const matches = findSymbolMatches(cg, 'handle');
expect(matches.length).toBe(2);
expect(matches.every((n: any) => n.name === 'handle')).toBe(true);
handler.closeAll();
cg.destroy();
cleanupTempDir(tmpDir);
});
it.skipIf(!HAS_SQLITE)('should return null when symbol is not found', async () => {
it.skipIf(!HAS_SQLITE)('should return no matches when symbol is not found', async () => {
const { ToolHandler } = await import('../src/mcp/tools');
const CodeGraph = (await import('../src/index')).default;
@@ -630,10 +630,10 @@ export function handle(): void {}
await cg.indexAll();
const handler = new ToolHandler(cg);
const findSymbol = (handler as any).findSymbol.bind(handler);
const findSymbolMatches = (handler as any).findSymbolMatches.bind(handler);
const match = findSymbol(cg, 'nonExistentSymbol');
expect(match).toBeNull();
const matches = findSymbolMatches(cg, 'nonExistentSymbol');
expect(matches.length).toBe(0);
handler.closeAll();
cg.destroy();
+19 -98
View File
@@ -263,23 +263,35 @@ describe('MCP Input Validation', () => {
expect(result.content[0].text).toContain('non-empty string');
});
it('should reject non-string task in codegraph_context', async () => {
const result = await handler.execute('codegraph_context', { task: undefined });
it('should reject non-string query in codegraph_explore', async () => {
const result = await handler.execute('codegraph_explore', { query: undefined });
expect(result.isError).toBe(true);
expect(result.content[0].text).toContain('non-empty string');
});
it('should truncate oversized codegraph_context output', async () => {
const oversizedContext = Array.from({ length: 400 }, (_, i) => `line-${i} ${'x'.repeat(80)}`).join('\n');
it('should truncate oversized tool output', async () => {
// Force a huge result set through codegraph_search; the response must be
// truncated with the sentinel rather than flooding the agent's context.
const many = Array.from({ length: 3000 }, (_, i) => ({
node: {
id: `n${i}`,
name: `symbol_${i}_${'x'.repeat(40)}`,
kind: 'function',
filePath: `src/very/deep/path/file_${i}.ts`,
startLine: 1,
endLine: 2,
language: 'typescript',
},
score: 1,
}));
const fakeCg = {
buildContext: async () => oversizedContext,
searchNodes: () => many,
};
const fakeHandler = new ToolHandler(fakeCg as unknown as CodeGraph);
const result = await fakeHandler.execute('codegraph_context', { task: 'find example' });
const result = await fakeHandler.execute('codegraph_search', { query: 'x' });
expect(result.isError).toBeFalsy();
expect(result.content[0].text.length).toBeLessThan(oversizedContext.length);
expect(result.content[0].text).toContain('... (output truncated)');
});
@@ -551,94 +563,3 @@ describe('Symlink Cycle Detection', () => {
expect(files).toContain('src/valid.ts');
});
});
describe('Session marker symlink resistance', () => {
// The marker write lives in src/mcp/tools.ts behind handleContext. We exercise
// it end-to-end via ToolHandler.execute so the test exercises the same code
// path Claude Code drives. The session id is per-test so other parallel test
// runs can't collide with the marker file we plant a symlink at.
const SESSION_ID = `cg-test-${process.pid}-${Date.now()}-${Math.random().toString(36).slice(2)}`;
const crypto = require('crypto') as typeof import('crypto');
const hash = crypto.createHash('md5').update(SESSION_ID).digest('hex').slice(0, 16);
const markerPath = path.join(os.tmpdir(), `codegraph-consulted-${hash}`);
let projectDir: string;
let victimDir: string;
let victimFile: string;
beforeEach(async () => {
projectDir = createTempDir();
victimDir = createTempDir();
victimFile = path.join(victimDir, 'private.txt');
fs.writeFileSync(victimFile, 'SECRET-DO-NOT-OVERWRITE\n');
if (fs.existsSync(markerPath)) fs.unlinkSync(markerPath);
// A real .codegraph/ has to exist for handleContext to get past the
// "not initialized" guard — index a tiny fixture so the call reaches the
// marker write step rather than short-circuiting on missing project state.
fs.writeFileSync(path.join(projectDir, 'a.ts'), 'export const x = 1;\n');
const cg = await CodeGraph.init(projectDir);
await cg.indexAll();
cg.close();
});
afterEach(() => {
if (fs.existsSync(markerPath)) fs.unlinkSync(markerPath);
cleanupTempDir(projectDir);
cleanupTempDir(victimDir);
});
it('does not follow a pre-planted symlink at the marker path', async () => {
// Skip on platforms where the user can't create symlinks (Windows without
// dev mode + admin). The CWE-59 risk we're guarding against doesn't apply
// when symlinks aren't creatable, so the skip is correct, not a gap.
try {
fs.symlinkSync(victimFile, markerPath);
} catch {
return;
}
const cg = await CodeGraph.open(projectDir);
const handler = new ToolHandler(cg);
process.env.CLAUDE_SESSION_ID = SESSION_ID;
try {
await handler.execute('codegraph_context', { task: 'find x' });
} finally {
delete process.env.CLAUDE_SESSION_ID;
cg.close();
}
// The victim file's contents must be untouched — the old writeFileSync
// path would have followed the symlink and written an ISO timestamp here.
expect(fs.readFileSync(victimFile, 'utf8')).toBe('SECRET-DO-NOT-OVERWRITE\n');
// And the marker path itself must still be the symlink we planted —
// no fallback path that quietly unlinked + recreated it (which would
// also work, but is a behavior we don't want to silently rely on).
expect(fs.lstatSync(markerPath).isSymbolicLink()).toBe(true);
});
it('writes the marker file with 0o600 perms on a clean path', async () => {
// No symlink planted — happy path. Verifies the new openSync(mode: 0o600)
// call is what actually lands on disk (regression guard for the perm
// tightening that came with the O_NOFOLLOW fix).
const cg = await CodeGraph.open(projectDir);
const handler = new ToolHandler(cg);
process.env.CLAUDE_SESSION_ID = SESSION_ID;
try {
await handler.execute('codegraph_context', { task: 'find x' });
} finally {
delete process.env.CLAUDE_SESSION_ID;
cg.close();
}
expect(fs.existsSync(markerPath)).toBe(true);
// chmod's low 9 bits — strip the file-type bits for a clean compare.
// Windows can't enforce 0o600 in the POSIX sense; skip the assertion
// there since the underlying OS will normalize the mode anyway.
if (process.platform !== 'win32') {
const mode = fs.statSync(markerPath).mode & 0o777;
expect(mode).toBe(0o600);
}
});
});
+58 -30
View File
@@ -75,7 +75,8 @@ describe.skipIf(!HAS_SQLITE)('matchesSymbol — module-qualified lookups (#173)'
let projectRoot: string;
let cg: any;
let handler: any;
let findSymbol: (cg: any, s: string) => { node: any; note: string } | null;
// findSymbolMatches returns ALL ranked matches; [0] is the resolved/picked one.
let findSymbolMatches: (cg: any, s: string) => any[];
let findAllSymbols: (cg: any, s: string) => { nodes: any[]; note: string };
beforeEach(async () => {
@@ -87,7 +88,7 @@ describe.skipIf(!HAS_SQLITE)('matchesSymbol — module-qualified lookups (#173)'
});
await cg.indexAll();
handler = new ToolHandler(cg);
findSymbol = (handler as any).findSymbol.bind(handler);
findSymbolMatches = (handler as any).findSymbolMatches.bind(handler);
findAllSymbols = (handler as any).findAllSymbols.bind(handler);
});
@@ -98,10 +99,11 @@ describe.skipIf(!HAS_SQLITE)('matchesSymbol — module-qualified lookups (#173)'
});
it('resolves `stage_apply::run` to the run in stage_apply.rs (not stage_detect.rs)', () => {
const match = findSymbol(cg, 'stage_apply::run');
expect(match).not.toBeNull();
expect(match!.node.name).toBe('run');
expect(match!.node.filePath).toMatch(/configurator\/stage_apply\.rs$/);
const matches = findSymbolMatches(cg, 'stage_apply::run');
expect(matches.length).toBeGreaterThan(0);
expect(matches[0]!.name).toBe('run');
// Every match must be in stage_apply.rs — never stage_detect.rs.
for (const n of matches) expect(n.filePath).toMatch(/configurator\/stage_apply\.rs$/);
});
it('rejects `stage_apply::run` for the same-named function in a different module', () => {
@@ -114,29 +116,29 @@ describe.skipIf(!HAS_SQLITE)('matchesSymbol — module-qualified lookups (#173)'
});
it('resolves `configurator::stage_apply::run` (multi-level qualifier)', () => {
const match = findSymbol(cg, 'configurator::stage_apply::run');
expect(match).not.toBeNull();
expect(match!.node.name).toBe('run');
expect(match!.node.filePath).toMatch(/configurator\/stage_apply\.rs$/);
const matches = findSymbolMatches(cg, 'configurator::stage_apply::run');
expect(matches.length).toBeGreaterThan(0);
expect(matches[0]!.name).toBe('run');
expect(matches[0]!.filePath).toMatch(/configurator\/stage_apply\.rs$/);
});
it('resolves `crate::configurator::stage_apply::run` (Rust path prefix stripped)', () => {
const match = findSymbol(cg, 'crate::configurator::stage_apply::run');
expect(match).not.toBeNull();
expect(match!.node.filePath).toMatch(/configurator\/stage_apply\.rs$/);
const matches = findSymbolMatches(cg, 'crate::configurator::stage_apply::run');
expect(matches.length).toBeGreaterThan(0);
expect(matches[0]!.filePath).toMatch(/configurator\/stage_apply\.rs$/);
});
it('resolves `configurator/stage_apply` (slash qualifier)', () => {
const match = findSymbol(cg, 'configurator/stage_apply/run');
expect(match).not.toBeNull();
expect(match!.node.filePath).toMatch(/configurator\/stage_apply\.rs$/);
const matches = findSymbolMatches(cg, 'configurator/stage_apply/run');
expect(matches.length).toBeGreaterThan(0);
expect(matches[0]!.filePath).toMatch(/configurator\/stage_apply\.rs$/);
});
it('does not silently collide bare `run` with `run_due_tasks`', () => {
const match = findSymbol(cg, 'run');
expect(match).not.toBeNull();
// Whatever it picks, it must be an exact-name match, not a partial.
expect(match!.node.name).toBe('run');
const matches = findSymbolMatches(cg, 'run');
expect(matches.length).toBeGreaterThan(0);
// Whatever it picks, every match must be an exact-name match, not a partial.
for (const n of matches) expect(n.name).toBe('run');
});
it('aggregates all bare-name `run` matches across modules', () => {
@@ -148,9 +150,22 @@ describe.skipIf(!HAS_SQLITE)('matchesSymbol — module-qualified lookups (#173)'
expect(all.note).toMatch(/Aggregated|symbols named "run"/);
});
it('still returns null for genuinely unknown qualified lookups', () => {
const match = findSymbol(cg, 'stage_apply::nonexistent_fn');
expect(match).toBeNull();
it('still returns nothing for genuinely unknown qualified lookups', () => {
const matches = findSymbolMatches(cg, 'stage_apply::nonexistent_fn');
expect(matches.length).toBe(0);
});
it('codegraph_node with a `file` hint pins an overloaded name to that file', async () => {
// `run` is defined in BOTH stage_apply.rs and stage_detect.rs. A bare lookup
// returns both; the `file` hint narrows to the one the caller saw in a trail.
const res = await handler.execute('codegraph_node', {
symbol: 'run',
includeCode: true,
file: 'stage_detect.rs',
});
const text = res.content?.[0]?.text ?? '';
expect(text).toMatch(/stage_detect\.rs/);
expect(text).not.toMatch(/stage_apply\.rs/);
});
});
@@ -158,7 +173,7 @@ describe.skipIf(!HAS_SQLITE)('matchesSymbol — dotted lookups (regression for #
let projectRoot: string;
let cg: any;
let handler: any;
let findSymbol: (cg: any, s: string) => { node: any; note: string } | null;
let findSymbolMatches: (cg: any, s: string) => any[];
beforeEach(async () => {
projectRoot = tmpRoot();
@@ -166,7 +181,7 @@ describe.skipIf(!HAS_SQLITE)('matchesSymbol — dotted lookups (regression for #
fs.mkdirSync(src, { recursive: true });
fs.writeFileSync(
path.join(src, 'session.ts'),
`export class Session {\n request(): void {}\n}\nexport function request(): void {}\n`
`export class Session {\n request(): void { fetch('x'); }\n}\nexport function request(): void {}\n`
);
const CodeGraph = (await import('../src/index')).default;
@@ -176,7 +191,7 @@ describe.skipIf(!HAS_SQLITE)('matchesSymbol — dotted lookups (regression for #
});
await cg.indexAll();
handler = new ToolHandler(cg);
findSymbol = (handler as any).findSymbol.bind(handler);
findSymbolMatches = (handler as any).findSymbolMatches.bind(handler);
});
afterEach(() => {
@@ -186,9 +201,22 @@ describe.skipIf(!HAS_SQLITE)('matchesSymbol — dotted lookups (regression for #
});
it('`Session.request` resolves to the method, not the bare function', () => {
const match = findSymbol(cg, 'Session.request');
expect(match).not.toBeNull();
expect(match!.node.kind).toBe('method');
expect(match!.node.qualifiedName).toContain('Session::request');
const matches = findSymbolMatches(cg, 'Session.request');
expect(matches.length).toBeGreaterThan(0);
expect(matches[0]!.kind).toBe('method');
expect(matches[0]!.qualifiedName).toContain('Session::request');
});
it('codegraph_node on an ambiguous bare name returns ALL overloads with bodies (no guess)', async () => {
// `request` is BOTH a method (Session.request) and a free function. The old
// behavior returned one + a dead-end "Others:" note, forcing a Read to get
// the other overload; now both bodies come back in one call.
const res = await handler.execute('codegraph_node', { symbol: 'request', includeCode: true });
const text = res.content?.[0]?.text ?? '';
expect(text).toContain('2 definitions named "request"');
// Both definitions are rendered (method + function), each with a Location.
expect(text).toMatch(/\(method\)/);
expect(text).toMatch(/\(function\)/);
expect((text.match(/\*\*Location:\*\*/g) || []).length).toBeGreaterThanOrEqual(2);
});
});
+1 -1
View File
@@ -180,7 +180,7 @@ describe('worktree mismatch surfaces on hot read tools (issue #155)', () => {
const savedPath = process.env.PATH;
process.env.PATH = '';
try {
const second = await handler.execute('codegraph_context', { task: 'mainOnly' });
const second = await handler.execute('codegraph_explore', { query: 'mainOnly' });
expect(second.content[0].text).toContain('different git worktree');
} finally {
process.env.PATH = savedPath;