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:
@@ -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);
|
||||
});
|
||||
});
|
||||
@@ -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`/);
|
||||
});
|
||||
});
|
||||
@@ -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);
|
||||
});
|
||||
|
||||
@@ -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/);
|
||||
|
||||
@@ -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();
|
||||
});
|
||||
});
|
||||
@@ -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
@@ -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);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user