From 4effb07f80f619c866ca854a6547c65044782a1f Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 01:21:41 +0530 Subject: [PATCH 01/13] Make truncated inline body hints actionable --- packages/core/src/repo.ts | 16 ++++++++++------ packages/core/test/core.test.ts | 29 +++++++++++++++++++++++++++-- 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/packages/core/src/repo.ts b/packages/core/src/repo.ts index b1f9e64..a26a120 100644 --- a/packages/core/src/repo.ts +++ b/packages/core/src/repo.ts @@ -145,7 +145,7 @@ const PREFIX_TOKEN_COST = 2 // The MCP compact (no-body) snippet renderer caps at this many lines; the token // estimate mirrors it so prefix cost is not over-counted for compact hits. const COMPACT_SNIPPET_MAX_LINES = 4 -const INLINE_BODY_TRUNCATION_MARKER = '… (truncated — read_chunk for full)' +const INLINE_BODY_TRUNCATION_MARKER = (locator: string) => `… (truncated — read_chunk ${locator} for full)` const INLINE_RANK2_SCORE_MARGIN = 0.6 // Below this lexical-row count, an over-constrained FTS query is progressively // relaxed (drop-rarest term, then full OR) so a concept phrased with a word the @@ -694,7 +694,7 @@ export class SqliteRepo implements Repo { const top = exactRows[0]! try { const source = await this.readRange(top.file_path, top.start_line, top.end_line) - const { body } = capInlineBody(source) + const { body } = capInlineBody(source, chunkLocator(top.file_path, top.start_line, top.end_line)) if (body) { topBody = body } @@ -1984,7 +1984,7 @@ async function tryInlineBody( return false } - const { body, tokens } = capInlineBody(source) + const { body, tokens } = capInlineBody(source, chunkLocator(hit.file, hit.range.startLine, hit.range.endLine)) if (!body) { return false } @@ -2017,7 +2017,7 @@ async function tryInlineBody( * stays recoverable from the line number, and that prefix cost is folded into * the returned token count. */ -function capInlineBody(source: string): { body: string; tokens: number } { +function capInlineBody(source: string, locator: string): { body: string; tokens: number } { if (!source) { return { body: '', tokens: 0 } } @@ -2041,10 +2041,14 @@ function capInlineBody(source: string): { body: string; tokens: number } { truncated = true } - const body = truncated ? `${kept.join('\n')}\n${INLINE_BODY_TRUNCATION_MARKER}` : kept.join('\n') + const body = truncated ? `${kept.join('\n')}\n${INLINE_BODY_TRUNCATION_MARKER(locator)}` : kept.join('\n') return { body, tokens: renderedBodyTokens(body) } } +function chunkLocator(file: string, startLine: number, endLine: number): string { + return `${file}:${startLine}-${endLine}` +} + /** Token count of a body INCLUDING the per-line `NN | ` render prefix. */ function renderedBodyTokens(text: string): number { if (!text) { @@ -3158,7 +3162,7 @@ function distinctExactSymbolDefinitions(exactRows: ChunkRow[], query: string): n const seen = new Set() for (const row of exactRows) { if (isExactSymbolMatch(row, query)) { - seen.add(`${row.file_path}${(row.symbol ?? '').toLowerCase()}${(row.parent ?? '').toLowerCase()}`) + seen.add(`${row.file_path}\u0000${(row.symbol ?? '').toLowerCase()}\u0000${(row.parent ?? '').toLowerCase()}`) } } return seen.size diff --git a/packages/core/test/core.test.ts b/packages/core/test/core.test.ts index 5f13b17..794d53a 100644 --- a/packages/core/test/core.test.ts +++ b/packages/core/test/core.test.ts @@ -609,8 +609,11 @@ describe('@codesift/core search body inlining', () => { const hits = await repo.search('huge retry backoff handler', { k: 1 }) expect(hits[0]?.body).toBeDefined() - const body = hits[0]!.body! - expect(body).toContain('… (truncated — read_chunk for full)') + const hit = hits[0]! + const body = hit.body! + const locator = `${hit.file}:${hit.range.startLine}-${hit.range.endLine}` + expect(body).toContain(`… (truncated — read_chunk ${locator} for full)`) + await expect(repo.readChunk(locator)).resolves.toContain('const huge79 = computeHugeRetryBackoff(79)') // Cap is whichever of ~50 lines / ~400 tokens is smaller; the body must be // shorter than the full 80+ line source on disk. expect(body.split('\n').length).toBeLessThanOrEqual(51) @@ -713,6 +716,28 @@ describe('@codesift/core find_symbol body inlining', () => { expect(body).not.toMatch(/\n[ \t]*\n[ \t]*\n/) expect(body).toContain('return base + 2') }) + + it('truncates oversized bodies with an actionable read_chunk marker', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-findsymbol-truncate-')) + temporaryDirectories.push(repoRoot) + await mkdir(join(repoRoot, 'src'), { recursive: true }) + + const longBody = Array.from({ length: 80 }, (_, index) => ` const huge${index} = computeHugeRetryBackoff(${index})`).join('\n') + await writeFile( + join(repoRoot, 'src', 'huge.ts'), + `export function hugeRetryBackoffHandler(): number {\n${longBody}\n return 0\n}\n`, + 'utf8' + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + + const [def] = await repo.findSymbol('hugeRetryBackoffHandler') + expect(def?.body).toBeDefined() + const locator = `${def!.file}:${def!.range.startLine}-${def!.range.endLine}` + expect(def!.body!).toContain(`… (truncated — read_chunk ${locator} for full)`) + await expect(repo.readChunk(locator)).resolves.toContain('const huge79 = computeHugeRetryBackoff(79)') + }) }) describe('@codesift/core import-resolved usages', () => { From 5420cac437107d7fa5ae35ba64da1d8da0fbe9b8 Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 01:28:34 +0530 Subject: [PATCH 02/13] Add efficiency guardrails for MCP output --- packages/eval/test/eval.test.ts | 30 +++++++++++++++++- packages/mcp/src/index.ts | 6 ++-- packages/mcp/test/mcp.test.ts | 55 +++++++++++++++++++++++++++------ 3 files changed, 78 insertions(+), 13 deletions(-) diff --git a/packages/eval/test/eval.test.ts b/packages/eval/test/eval.test.ts index 1bc78e1..1390ed4 100644 --- a/packages/eval/test/eval.test.ts +++ b/packages/eval/test/eval.test.ts @@ -4,9 +4,16 @@ import { join } from 'node:path' import { afterEach, describe, expect, it } from 'vitest' -import { DEFAULT_MANIFEST_PATH, candidateFromRipgrep, candidateMatchesExpected, createEmptyManifest, diffLossBudgets, evaluateManifest, formatSummary, isRerankEvalEnabled, loadManifest, proveRoutingPolicy, runRipgrep, summarizeEmptyRun, type GoldenQuery, type RipgrepHit } from '../src/index.js' +import { DEFAULT_MANIFEST_PATH, candidateFromRipgrep, candidateMatchesExpected, createEmptyManifest, diffLossBudgets, evaluateManifest, formatSummary, isRerankEvalEnabled, loadManifest, proveRoutingPolicy, runRipgrep, summarizeEmptyRun, type GoldenQuery, type GoldenQueryType, type RipgrepHit } from '../src/index.js' const temporaryDirectories: string[] = [] +const stableLocalTokenCeilings: Partial> = { + 'nl-concept': 180, + 'symbol-def': 125, + 'exact-identifier': 75, + 'string-literal': 40, + 'error-trace': 40 +} afterEach(async () => { await Promise.all( @@ -77,6 +84,27 @@ describe('@codesift/eval', () => { ) }, 30_000) + it('keeps stable local codesift runs at one-call rank-1 resolution with low token envelopes', async () => { + const manifest = await loadManifest(DEFAULT_MANIFEST_PATH) + const stableLocalManifest = { + repos: manifest.repos.filter((repo) => repo.id.startsWith('m3-') || repo.id === 'collision-ts' || repo.id === 'usages-ts'), + queries: manifest.queries.filter((query) => query.repoId.startsWith('m3-') || query.repoId === 'collision-ts' || query.repoId === 'usages-ts') + } + const summary = await evaluateManifest(stableLocalManifest, { resultLimit: 10, inspectionLimit: 5, latencyToleranceMs: Number.POSITIVE_INFINITY }) + const codesiftRuns = summary.runs.filter((run) => run.tool === 'codesift') + + expect(codesiftRuns).toHaveLength(stableLocalManifest.queries.length) + for (const run of codesiftRuns) { + expect(run.taskSuccess).toBe(true) + expect(run.callsToResolution).toBe(1) + expect(run.meanReciprocalRank).toBe(1) + + const tokenCeiling = stableLocalTokenCeilings[run.queryType] + expect(tokenCeiling).toBeDefined() + expect(run.tokensToResolution).toBeLessThanOrEqual(tokenCeiling!) + } + }, 30_000) + it('measures end-to-end calls-to-resolution truthfully across both tools', async () => { const manifest = await loadManifest(DEFAULT_MANIFEST_PATH) const m3Manifest = { diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 15224a2..49d7a58 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -542,7 +542,7 @@ export function formatMcpSearchHits(hits: SearchHit[]): string { return hit.usages?.length ? `${block}\n${formatUsageBlock(hit.usages)}` : block } - const snippet = compactHitSnippet(hit.snippet, hit.range.startLine ?? hit.snippetRange.startLine, 4) + const snippet = compactHitSnippet(hit.snippet, hit.snippetRange.startLine ?? hit.range.startLine, 4) const block = snippet ? `${header}\n${snippet}` : header return hit.usages?.length ? `${block}\n${formatUsageBlock(hit.usages)}` : block }) @@ -618,8 +618,8 @@ function compactSnippet(snippet: string, maxLines: number): string { } // Structure-preserving renderer for compact (no-body) search hits: keeps -// original indentation and emits `NN | code` line-number prefixes, matching the -// inlined-body block. Only trailing whitespace and a trailing newline are dropped. +// original indentation and emits `NN | code` prefixes from the centered snippet +// range. Only trailing whitespace and a trailing newline are dropped. function compactHitSnippet(snippet: string, startLine: number, maxLines: number): string { const lines = snippet.replace(/\n$/, '').split('\n').slice(0, maxLines) if (lines.length === 0) { diff --git a/packages/mcp/test/mcp.test.ts b/packages/mcp/test/mcp.test.ts index a17b734..c52b264 100644 --- a/packages/mcp/test/mcp.test.ts +++ b/packages/mcp/test/mcp.test.ts @@ -219,8 +219,12 @@ describe('@codesift/mcp server', () => { params: { name: 'search_code', arguments: { query: 'demoValue', k: 1 } } }) const searchResult = await waitForJsonRpcMessage(messages, (message) => rpcId(message) === 3, () => parseError, child) - const searchText = JSON.stringify(searchResult) + const searchText = mcpText(searchResult) expect(searchText).toContain('src/demo.ts') + expect(searchText).toContain('1 | export function demoValue(): string {') + expect(searchText).toContain("2 | return 'demo'") + expect(searchText).toContain('tokensReturned=') + expect(searchText).not.toContain(' ↩ ') const chunkId = /[=~+] ([^\s"]+)/.exec(searchText)?.[1] expect(chunkId).toBeTruthy() @@ -263,7 +267,7 @@ describe('formatMcpSearchHits structure-preserving output', () => { const lines = output.split('\n') // Header carries reason + id (hash stripped) + symbol path. expect(lines[0]).toBe('~ src/auth.ts:10-14 verify') - // Body block is line-numbered from startLine, preserving indentation, no ' ↩ ' flattening. + // Body block is line-numbered from range.startLine, preserving indentation, no ' ↩ ' flattening. expect(lines[1]).toBe('10 | function verify(token: string): boolean {') expect(lines[2]).toBe('11 | if (!token) {') expect(lines[3]).toBe('12 | return false') @@ -304,17 +308,22 @@ describe('formatMcpSearchHits structure-preserving output', () => { expect(output).toContain('3 | … (truncated — read_chunk src/auth.ts:10-14 for full)') }) - it('renders a compact (no body) hit with real newlines, preserved indentation, and NN | line-number prefixes', () => { + it('renders a compact (no body) hit from snippetRange with real newlines, preserved indentation, and NN | prefixes', () => { const output = formatMcpSearchHits([ - makeHit({ snippet: 'function verify() {\n return true\n}' }) + makeHit({ + range: { startLine: 10, endLine: 30 }, + snippet: 'function verify() {\n return true\n}', + snippetRange: { startLine: 18, endLine: 20 } + }) ]) expect(output).not.toContain(' ↩ ') const lines = output.split('\n') expect(lines[0]).toBe('~ src/auth.ts:10-14 verify') - // Compact snippet is line-numbered from range.startLine and keeps original indentation. - expect(lines[1]).toBe('10 | function verify() {') - expect(lines[2]).toBe('11 | return true') - expect(lines[3]).toBe('12 | }') + // Compact snippets are line-numbered from snippetRange.startLine; core may + // center snippets inside a wider hit range. + expect(lines[1]).toBe('18 | function verify() {') + expect(lines[2]).toBe('19 | return true') + expect(lines[3]).toBe('20 | }') expect(output).toContain('tokensReturned=42') }) @@ -322,7 +331,8 @@ describe('formatMcpSearchHits structure-preserving output', () => { const output = formatMcpSearchHits([ makeHit({ snippet: 'class Auth {\n verify() {\n return this.token != null\n }\n}', - range: { startLine: 20, endLine: 24 } + range: { startLine: 10, endLine: 30 }, + snippetRange: { startLine: 20, endLine: 24 } }) ]) const lines = output.split('\n') @@ -418,3 +428,30 @@ function rpcId(message: unknown): number | undefined { const id = (message as { id?: unknown }).id return typeof id === 'number' ? id : undefined } + +function mcpText(message: unknown): string { + if (typeof message !== 'object' || message === null || !('result' in message)) { + return '' + } + + const result = (message as { result?: unknown }).result + if (typeof result !== 'object' || result === null || !('content' in result)) { + return '' + } + + const content = (result as { content?: unknown }).content + if (!Array.isArray(content)) { + return '' + } + + return content + .map((entry) => { + if (typeof entry !== 'object' || entry === null || !('text' in entry)) { + return '' + } + + const text = (entry as { text?: unknown }).text + return typeof text === 'string' ? text : '' + }) + .join('\n') +} From 94e57f0f3d60ec15e6ad387a4602601278ad32b9 Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 02:02:52 +0530 Subject: [PATCH 03/13] Budget grep_code MCP output --- packages/eval/src/index.ts | 6 +- packages/eval/test/eval.test.ts | 34 ++++++- packages/mcp/src/index.ts | 163 +++++++++++++++++++++++++++++--- packages/mcp/test/mcp.test.ts | 107 +++++++++++++++++++-- 4 files changed, 286 insertions(+), 24 deletions(-) diff --git a/packages/eval/src/index.ts b/packages/eval/src/index.ts index 80ab593..840eac2 100644 --- a/packages/eval/src/index.ts +++ b/packages/eval/src/index.ts @@ -7,7 +7,7 @@ import { execFile as execFileCallback, spawn, type ChildProcessWithoutNullStream import { promisify } from 'node:util' import { VOYAGE_RERANK_PROVIDER_ID, openRepo, type FindSymbolOptions, type GrepHit, type GrepOptions, type Range, type SearchHit, type SearchOptions, type SymbolDefinition } from '@codesift/core' -import { formatMcpGrepHits, formatMcpSearchHits, formatMcpSymbols } from '@codesift/mcp' +import { DEFAULT_MCP_GREP_MAX_TOKENS, formatMcpGrepHits, formatMcpSearchHits, formatMcpSymbols } from '@codesift/mcp' const execFile = promisify(execFileCallback) @@ -686,7 +686,7 @@ function buildCodesiftToolCall(query: GoldenQuery, resultLimit: number): { name: case 'error-trace': return { name: 'grep_code', - arguments: { pattern: query.grepPattern ?? query.query, max_matches: resultLimit, ...pathGlob } + arguments: { pattern: query.grepPattern ?? query.query, max_matches: resultLimit, max_tokens: DEFAULT_MCP_GREP_MAX_TOKENS, ...pathGlob } } case 'nl-concept': return { @@ -869,7 +869,7 @@ async function runCodesiftPolicy(repo: Awaited>, que const pattern = query.grepPattern ?? query.query const hits = await repo.grep(pattern, grepOptions) const candidates = hits.map(candidateFromGrep) - const tokens = estimateTokenCount(formatMcpGrepHits(hits)) + const tokens = estimateTokenCount(formatMcpGrepHits(hits, { maxTokens: DEFAULT_MCP_GREP_MAX_TOKENS })) return evaluateRankedCandidates(candidates, query, tokens, inspectionLimit) } case 'nl-concept': { diff --git a/packages/eval/test/eval.test.ts b/packages/eval/test/eval.test.ts index 1390ed4..38cc7c2 100644 --- a/packages/eval/test/eval.test.ts +++ b/packages/eval/test/eval.test.ts @@ -1,4 +1,4 @@ -import { mkdtemp, rm, writeFile } from 'node:fs/promises' +import { mkdir, mkdtemp, rm, writeFile } from 'node:fs/promises' import { tmpdir } from 'node:os' import { join } from 'node:path' @@ -105,6 +105,38 @@ describe('@codesift/eval', () => { } }, 30_000) + it('accounts string-literal grep results with the MCP default token budget', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-eval-grep-budget-')) + temporaryDirectories.push(repoRoot) + + await mkdir(join(repoRoot, 'src'), { recursive: true }) + await writeFile( + join(repoRoot, 'src', 'noisy.ts'), + Array.from({ length: 160 }, (_, index) => `export const noisy${index} = 'LOUD_LITERAL_${index}_${'x'.repeat(80)}'`).join('\n'), + 'utf8' + ) + + const summary = await evaluateManifest({ + repos: [{ id: 'noisy-grep', language: 'typescript', repoPath: repoRoot }], + queries: [ + { + id: 'noisy-literal', + repoId: 'noisy-grep', + queryType: 'string-literal', + query: 'LOUD_LITERAL', + grepPattern: 'LOUD_LITERAL', + expected: [{ file: 'src/noisy.ts' }], + expectedLineRange: { startLine: 1, endLine: 1 } + } + ] + }, { resultLimit: 160, inspectionLimit: 5, latencyToleranceMs: Number.POSITIVE_INFINITY }) + + const codesiftRun = summary.runs.find((run) => run.tool === 'codesift') + expect(codesiftRun?.taskSuccess).toBe(true) + expect(codesiftRun?.meanReciprocalRank).toBe(1) + expect(codesiftRun?.tokensToResolution).toBeLessThanOrEqual(730) + }, 30_000) + it('measures end-to-end calls-to-resolution truthfully across both tools', async () => { const manifest = await loadManifest(DEFAULT_MANIFEST_PATH) const m3Manifest = { diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 49d7a58..e951d06 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -23,6 +23,7 @@ import { createHttpServerHandle } from './http.js' export { HttpMcpServerHandle, createHttpServerHandle } from './http.js' export const DEFAULT_SEARCH_K = CORE_DEFAULT_SEARCH_K +export const DEFAULT_MCP_GREP_MAX_TOKENS = 700 const SYMBOL_KINDS = [ 'class', @@ -79,6 +80,11 @@ export interface GrepCodeArgs { before_context_lines?: number | undefined after_context_lines?: number | undefined max_matches?: number | undefined + max_tokens?: number | undefined +} + +export interface FormatMcpGrepHitsOptions { + maxTokens?: number | undefined } export interface ReadChunkArgs { @@ -170,7 +176,8 @@ const grepCodeInputSchema = { context_lines: z.number().int().min(0).max(20).optional().describe('Symmetric context lines, like rg -C.'), before_context_lines: z.number().int().min(0).max(20).optional().describe('Lines before each match, like rg -B.'), after_context_lines: z.number().int().min(0).max(20).optional().describe('Lines after each match, like rg -A.'), - max_matches: z.number().int().positive().max(1000).optional().describe('Maximum matches to return. Default 1000.') + max_matches: z.number().int().positive().max(1000).optional().describe('Maximum matches to return. Default 1000.'), + max_tokens: z.number().int().positive().max(4000).optional().describe('Maximum approximate tokens to return across compact matches. Default 700.') } const readChunkInputSchema = { @@ -260,7 +267,8 @@ export function getToolDefinitions(): readonly McpToolDefinition[] { context_lines: { type: 'integer', minimum: 0, maximum: 20 }, before_context_lines: { type: 'integer', minimum: 0, maximum: 20 }, after_context_lines: { type: 'integer', minimum: 0, maximum: 20 }, - max_matches: { type: 'integer', minimum: 1, maximum: 1000 } + max_matches: { type: 'integer', minimum: 1, maximum: 1000 }, + max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: 700 } }) }, { @@ -390,7 +398,10 @@ export async function callMcpTool(repo: Repo, name: McpToolName, args: unknown): case 'find_symbol': return formatMcpSymbols(await router.findSymbol(findSymbolArgsSchema.parse(args))) case 'grep_code': - return formatMcpGrepHits(await router.grepCode(grepCodeArgsSchema.parse(args))) + { + const parsed = grepCodeArgsSchema.parse(args) + return formatMcpGrepHits(await router.grepCode(parsed), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_GREP_MAX_TOKENS }) + } case 'read_chunk': return router.readChunk(readChunkArgsSchema.parse(args)) case 'index_status': @@ -466,7 +477,7 @@ export function createSdkServer(repo: Repo): McpServer { textResult(formatMcpSymbols(await router.findSymbol(args))) ) server.registerTool('grep_code', { description: toolDescription('grep_code'), inputSchema: grepCodeInputSchema }, async (args) => - textResult(formatMcpGrepHits(await router.grepCode(args))) + textResult(formatMcpGrepHits(await router.grepCode(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_GREP_MAX_TOKENS })) ) server.registerTool('read_chunk', { description: toolDescription('read_chunk'), inputSchema: readChunkInputSchema }, async (args) => textResult(await router.readChunk(args)) @@ -590,18 +601,146 @@ export function formatMcpSymbols(definitions: SymbolDefinition[]): string { .join('\n') } -export function formatMcpGrepHits(hits: GrepHit[]): string { +export function formatMcpGrepHits(hits: GrepHit[], options: FormatMcpGrepHitsOptions = {}): string { if (hits.length === 0) { return 'no_matches' } - return hits - .map((hit) => { - const range = formatRange(hit.range.startLine, hit.range.endLine) - const snippet = compactSnippet(hit.snippet, 5).split('\n').join(' ↩ ') - return `${hit.file}:${range}:${hit.column} | ${snippet}` - }) - .join('\n') + const maxTokens = options.maxTokens + if (maxTokens === undefined) { + return hits.map((hit) => formatMcpGrepHit(hit)).join('\n') + } + + const maxChars = maxTokens * 4 + const rendered: string[] = [] + let usedChars = 0 + + for (let index = 0; index < hits.length; index += 1) { + const hit = hits[index]! + const line = formatMcpGrepHit(hit) + const separatorChars = rendered.length === 0 ? 0 : 1 + if (usedChars + separatorChars + line.length <= maxChars) { + rendered.push(line) + usedChars += separatorChars + line.length + continue + } + + if (rendered.length === 0) { + const omitted = hits.length - 1 + const marker = omitted > 0 ? grepOmissionMarker(omitted, Math.max(0, maxChars - 2)) : undefined + const firstHit = formatMcpGrepHit(hit, maxCharsForFirstHit(maxChars, marker)) + return marker ? withGrepOmissionMarker([firstHit], marker) : firstHit + } + + return fitGrepOutputWithMarker(rendered, hits.length - index, maxChars) + } + + return rendered.join('\n') +} + +function formatMcpGrepHit(hit: GrepHit, maxChars?: number): string { + const range = formatRange(hit.range.startLine, hit.range.endLine) + const suffix = `:${range}:${hit.column} | ` + const prefix = maxChars === undefined ? `${hit.file}${suffix}` : formatGrepPrefix(hit.file, suffix, maxChars) + const snippet = compactSnippet(hit.snippet, 5).split('\n').join(' ↩ ') + + if (maxChars === undefined || prefix.length + snippet.length <= maxChars) { + return `${prefix}${snippet}` + } + + return `${prefix}${truncateWithEllipsis(snippet, maxChars - prefix.length)}` +} + +function formatGrepPrefix(file: string, suffix: string, maxChars: number): string { + const prefix = `${file}${suffix}` + if (prefix.length <= maxChars) { + return prefix + } + + const fileChars = Math.max(0, maxChars - suffix.length - 1) + if (fileChars > 0) { + return `…${file.slice(-fileChars)}${suffix}` + } + + return suffix.length <= maxChars ? suffix : suffix.slice(-maxChars) +} + +function fitGrepOutputWithMarker(rendered: string[], omitted: number, maxChars: number): string { + const kept = [...rendered] + let omittedCount = omitted + let marker = grepOmissionMarker(omittedCount, Math.max(0, maxChars - 2)) + + while (kept.length > 1 && joinedLength(kept, marker) > maxChars) { + kept.pop() + omittedCount += 1 + marker = grepOmissionMarker(omittedCount, Math.max(0, maxChars - 2)) + } + + if (!marker) { + return formatMcpGrepLineForBudget(kept[0]!, maxChars) + } + + if (joinedLength(kept, marker) <= maxChars) { + return withGrepOmissionMarker(kept, marker) + } + + const line = formatMcpGrepLineForBudget(kept[0]!, maxCharsForFirstHit(maxChars, marker)) + return marker ? withGrepOmissionMarker([line], marker) : line +} + +function formatMcpGrepLineForBudget(line: string, maxChars: number): string { + const separator = ' | ' + const separatorIndex = line.indexOf(separator) + if (separatorIndex < 0 || line.length <= maxChars) { + return line + } + + let prefix = line.slice(0, separatorIndex + separator.length) + const snippet = line.slice(separatorIndex + separator.length) + if (prefix.length > maxChars) { + prefix = `…${prefix.slice(-(Math.max(0, maxChars - 1)))}` + } + return `${prefix}${truncateWithEllipsis(snippet, maxChars - prefix.length)}` +} + +function maxCharsForFirstHit(maxChars: number, marker: string | undefined): number { + return marker === undefined ? maxChars : Math.max(0, maxChars - marker.length - 1) +} + +function joinedLength(rendered: string[], marker: string): number { + return rendered.join('\n').length + 1 + marker.length +} + +function grepOmissionMarker(omitted: number, maxChars?: number): string { + const full = `matches_omitted=${omitted}; refine path_glob/max_matches or raise max_tokens.` + if (maxChars === undefined || full.length <= maxChars) { + return full + } + + const compact = `matches_omitted=${omitted}; raise max_tokens` + if (compact.length <= maxChars) { + return compact + } + + const bare = `matches_omitted=${omitted}` + return bare.length <= maxChars ? bare : '' +} + +function withGrepOmissionMarker(rendered: string[], marker: string): string { + return `${rendered.join('\n')}\n${marker}` +} + +function truncateWithEllipsis(text: string, maxChars: number): string { + if (maxChars <= 0) { + return '' + } + if (text.length <= maxChars) { + return text + } + if (maxChars === 1) { + return '…' + } + return `${text.slice(0, maxChars - 1).trimEnd()}…` } function formatChunkId(id: string): string { diff --git a/packages/mcp/test/mcp.test.ts b/packages/mcp/test/mcp.test.ts index c52b264..76794cf 100644 --- a/packages/mcp/test/mcp.test.ts +++ b/packages/mcp/test/mcp.test.ts @@ -13,6 +13,7 @@ import { createHttpServer, createRouter, createStdioServer, + callMcpTool, formatMcpGrepHits, formatMcpSearchHits, formatMcpSymbols, @@ -46,6 +47,9 @@ describe('@codesift/mcp server', () => { 'index_status' ]) expect(getToolDefinitions().find((tool) => tool.name === 'grep_code')?.inputSchema.required).toEqual(['pattern']) + expect(getToolDefinitions().find((tool) => tool.name === 'grep_code')?.inputSchema.properties).toMatchObject({ + max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: 700 } + }) expect(getToolDefinitions().find((tool) => tool.name === 'search_code')?.inputSchema.properties).toMatchObject({ context: { type: 'string', enum: ['sig', 'body'] }, with_usages: { type: 'boolean' } @@ -371,22 +375,109 @@ describe('formatMcpSearchHits structure-preserving output', () => { describe('formatMcpGrepHits keeps its single-line ↩-joined format', () => { it('flattens the snippet with the ↩ arrow and trims each line (unchanged from search-hit rendering)', () => { - const hit: GrepHit = { - file: 'src/auth.ts', - range: { startLine: 10, endLine: 12 }, - line: 10, - column: 3, - match: 'verify', - snippet: 'function verify() {\n return true\n}' - } + const hit = makeGrepHit({ snippet: 'function verify() {\n return true\n}' }) expect(formatMcpGrepHits([hit])).toBe('src/auth.ts:10-12:3 | function verify() { ↩ return true ↩ }') }) it('returns no_matches for an empty result set', () => { expect(formatMcpGrepHits([])).toBe('no_matches') }) + + it('leaves grep formatting unchanged when the rendered hits fit under the token budget', () => { + const output = formatMcpGrepHits([ + makeGrepHit({ snippet: 'function verify() {\n return true\n}' }) + ], { maxTokens: 80 }) + + expect(output).toBe('src/auth.ts:10-12:3 | function verify() { ↩ return true ↩ }') + }) + + it('preserves first hits and appends a compact omission marker when maxTokens truncates matches', () => { + const hits = Array.from({ length: 6 }, (_, index) => + makeGrepHit({ + file: `src/file-${index}.ts`, + range: { startLine: index + 1, endLine: index + 1 }, + line: index + 1, + snippet: `const repeated${index} = 'needle-${index}'` + }) + ) + const output = formatMcpGrepHits(hits, { maxTokens: 32 }) + + expect(output).toContain("src/file-0.ts:1:3 | const repeated0 = 'needle-0'") + expect(output).not.toContain('src/file-5.ts') + expect(output).toContain('matches_omitted=') + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(32) + }) + + it('truncates only snippet text when the first grep hit alone exceeds the budget', () => { + const output = formatMcpGrepHits([ + makeGrepHit({ + file: 'src/verbose.ts', + range: { startLine: 7, endLine: 7 }, + line: 7, + column: 11, + snippet: `const message = '${'x'.repeat(400)}'` + }), + makeGrepHit({ file: 'src/tail.ts' }) + ], { maxTokens: 25 }) + + expect(output.startsWith('src/verbose.ts:7:11 | ')).toBe(true) + expect(output).toContain('…') + expect(output).toContain('matches_omitted=1; refine path_glob/max_matches or raise max_tokens.') + expect(output).not.toContain('src/tail.ts') + }) + + it('keeps long-prefix tiny-budget output within the approximate budget', () => { + const output = formatMcpGrepHits([ + makeGrepHit({ + file: `src/${'very-long-directory/'.repeat(8)}verbose.ts`, + range: { startLine: 1, endLine: 1 }, + line: 1, + column: 1, + snippet: 'const value = "needle"' + }), + makeGrepHit({ file: 'src/tail.ts' }) + ], { maxTokens: 8 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(8) + expect(output).toContain('matches_omitted=1') + expect(output).not.toContain('src/tail.ts') + }) +}) + +describe('grep_code MCP output budgeting', () => { + it('budgets noisy literal output through callMcpTool while preserving the first match', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-mcp-grep-budget-')) + temporaryDirectories.push(repoRoot) + + await mkdir(join(repoRoot, 'src'), { recursive: true }) + await writeFile( + join(repoRoot, 'src', 'noisy.ts'), + Array.from({ length: 20 }, (_, index) => `export const value${index} = 'NEEDLE_${index}'`).join('\n'), + 'utf8' + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + + const output = await callMcpTool(repo, 'grep_code', { pattern: 'NEEDLE', path_glob: 'src/**', max_tokens: 40 }) + expect(output).toContain("src/noisy.ts:1:24 | export const value0 = 'NEEDLE_0'") + expect(output).toContain('matches_omitted=') + expect(output).not.toContain('NEEDLE_19') + }) }) +function makeGrepHit(overrides: Partial = {}): GrepHit { + return { + file: 'src/auth.ts', + range: { startLine: 10, endLine: 12 }, + line: 10, + column: 3, + match: 'verify', + snippet: 'function verify() {\n return true\n}', + ...overrides + } +} + function sendJsonRpc(child: ChildProcessWithoutNullStreams, message: unknown): void { child.stdin.write(`${JSON.stringify(message)}\n`) } From dee44b486ae0a35e6ce9eb3545908587dd609525 Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 02:13:37 +0530 Subject: [PATCH 04/13] Budget read_chunk MCP output --- packages/eval/src/index.ts | 4 +- packages/mcp/src/index.ts | 73 +++++++++++++++++++++++++++++++++-- packages/mcp/test/mcp.test.ts | 52 +++++++++++++++++++++++++ 3 files changed, 123 insertions(+), 6 deletions(-) diff --git a/packages/eval/src/index.ts b/packages/eval/src/index.ts index 840eac2..60b3a9a 100644 --- a/packages/eval/src/index.ts +++ b/packages/eval/src/index.ts @@ -7,7 +7,7 @@ import { execFile as execFileCallback, spawn, type ChildProcessWithoutNullStream import { promisify } from 'node:util' import { VOYAGE_RERANK_PROVIDER_ID, openRepo, type FindSymbolOptions, type GrepHit, type GrepOptions, type Range, type SearchHit, type SearchOptions, type SymbolDefinition } from '@codesift/core' -import { DEFAULT_MCP_GREP_MAX_TOKENS, formatMcpGrepHits, formatMcpSearchHits, formatMcpSymbols } from '@codesift/mcp' +import { DEFAULT_MCP_GREP_MAX_TOKENS, DEFAULT_MCP_READ_CHUNK_MAX_TOKENS, formatMcpGrepHits, formatMcpReadChunk, formatMcpSearchHits, formatMcpSymbols } from '@codesift/mcp' const execFile = promisify(execFileCallback) @@ -1307,7 +1307,7 @@ function bodyOverlapsExpected(hit: SearchHit, expected: Range): boolean { async function readChunkSafely(repo: Awaited>, id: string): Promise { try { - return await repo.readChunk(id) + return formatMcpReadChunk(await repo.readChunk(id), { maxTokens: DEFAULT_MCP_READ_CHUNK_MAX_TOKENS }) } catch { return null } diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index e951d06..fdea6da 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -24,6 +24,9 @@ export { HttpMcpServerHandle, createHttpServerHandle } from './http.js' export const DEFAULT_SEARCH_K = CORE_DEFAULT_SEARCH_K export const DEFAULT_MCP_GREP_MAX_TOKENS = 700 +export const DEFAULT_MCP_READ_CHUNK_MAX_TOKENS = 1000 +export const MIN_MCP_READ_CHUNK_MAX_TOKENS = 6 +export const MAX_MCP_READ_CHUNK_MAX_TOKENS = 4000 const SYMBOL_KINDS = [ 'class', @@ -87,9 +90,14 @@ export interface FormatMcpGrepHitsOptions { maxTokens?: number | undefined } +export interface FormatMcpReadChunkOptions { + maxTokens?: number | undefined +} + export interface ReadChunkArgs { id: string context_lines?: number | undefined + max_tokens?: number | undefined } export interface McpToolDefinition { @@ -182,7 +190,8 @@ const grepCodeInputSchema = { const readChunkInputSchema = { id: z.string().min(1).describe('Stable chunk id returned by search_code.'), - context_lines: z.number().int().min(0).max(50).optional().describe('Extra lines before/after the chunk.') + context_lines: z.number().int().min(0).max(50).optional().describe('Extra lines before/after the chunk.'), + max_tokens: z.number().int().min(MIN_MCP_READ_CHUNK_MAX_TOKENS).max(MAX_MCP_READ_CHUNK_MAX_TOKENS).optional().describe('Maximum approximate tokens to return from this chunk. Default 1000.') } const indexStatusInputSchema = {} @@ -276,7 +285,8 @@ export function getToolDefinitions(): readonly McpToolDefinition[] { description: 'Expand an ADDITIONAL search_code hit id into its source chunk, or widen context around one. The top search_code result is already returned inline, so this is rarely needed for the best hit; use it for secondary hits or extra surrounding lines, not as a first step.', inputSchema: jsonSchema(['id'], { id: { type: 'string' }, - context_lines: { type: 'integer', minimum: 0, maximum: 50 } + context_lines: { type: 'integer', minimum: 0, maximum: 50 }, + max_tokens: { type: 'integer', minimum: MIN_MCP_READ_CHUNK_MAX_TOKENS, maximum: MAX_MCP_READ_CHUNK_MAX_TOKENS, default: DEFAULT_MCP_READ_CHUNK_MAX_TOKENS } }) }, { @@ -403,7 +413,10 @@ export async function callMcpTool(repo: Repo, name: McpToolName, args: unknown): return formatMcpGrepHits(await router.grepCode(parsed), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_GREP_MAX_TOKENS }) } case 'read_chunk': - return router.readChunk(readChunkArgsSchema.parse(args)) + { + const parsed = readChunkArgsSchema.parse(args) + return formatMcpReadChunk(await router.readChunk(parsed), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_READ_CHUNK_MAX_TOKENS }) + } case 'index_status': return JSON.stringify(await router.indexStatus()) } @@ -480,7 +493,7 @@ export function createSdkServer(repo: Repo): McpServer { textResult(formatMcpGrepHits(await router.grepCode(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_GREP_MAX_TOKENS })) ) server.registerTool('read_chunk', { description: toolDescription('read_chunk'), inputSchema: readChunkInputSchema }, async (args) => - textResult(await router.readChunk(args)) + textResult(formatMcpReadChunk(await router.readChunk(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_READ_CHUNK_MAX_TOKENS })) ) server.registerTool('index_status', { description: toolDescription('index_status'), inputSchema: indexStatusInputSchema }, async () => textResult(JSON.stringify(await router.indexStatus())) @@ -638,6 +651,58 @@ export function formatMcpGrepHits(hits: GrepHit[], options: FormatMcpGrepHitsOpt return rendered.join('\n') } +export function formatMcpReadChunk(content: string, options: FormatMcpReadChunkOptions = {}): string { + const maxTokens = options.maxTokens + if (maxTokens === undefined) { + return content + } + + const maxChars = maxTokens * 4 + if (content.length <= maxChars) { + return content + } + + const marker = readChunkTruncationMarker(Math.max(0, maxChars - 1)) + if (!marker) { + return truncateWithEllipsis('content_truncated=true', maxChars) + } + + const contentBudget = Math.max(0, maxChars - marker.length - 1) + const truncatedContent = truncateReadChunkContent(content, contentBudget) + if (!truncatedContent) { + return marker.length <= maxChars ? marker : truncateWithEllipsis(marker, maxChars) + } + + return `${truncatedContent}\n${marker}` +} + +function truncateReadChunkContent(content: string, maxChars: number): string { + if (maxChars <= 0) { + return '' + } + if (content.length <= maxChars) { + return content + } + + const truncated = truncateWithEllipsis(content, maxChars) + return truncated.replace(/\n+$/g, '') +} + +function readChunkTruncationMarker(maxChars?: number): string { + const full = 'content_truncated=true; raise max_tokens or narrow the chunk/range.' + if (maxChars === undefined || full.length <= maxChars) { + return full + } + + const compact = 'content_truncated=true; raise max_tokens' + if (compact.length <= maxChars) { + return compact + } + + const bare = 'content_truncated=true' + return bare.length <= maxChars ? bare : '' +} + function formatMcpGrepHit(hit: GrepHit, maxChars?: number): string { const range = formatRange(hit.range.startLine, hit.range.endLine) const suffix = `:${range}:${hit.column} | ` diff --git a/packages/mcp/test/mcp.test.ts b/packages/mcp/test/mcp.test.ts index 76794cf..e7c7583 100644 --- a/packages/mcp/test/mcp.test.ts +++ b/packages/mcp/test/mcp.test.ts @@ -8,13 +8,16 @@ import { afterEach, describe, expect, it } from 'vitest' import { openRepo, registerEmbeddingProvider, type GrepHit, type SearchHit } from '@codesift/core' import { + DEFAULT_MCP_READ_CHUNK_MAX_TOKENS, DEFAULT_SEARCH_K, + MIN_MCP_READ_CHUNK_MAX_TOKENS, MCP_SERVER_INSTRUCTIONS, createHttpServer, createRouter, createStdioServer, callMcpTool, formatMcpGrepHits, + formatMcpReadChunk, formatMcpSearchHits, formatMcpSymbols, getToolDefinitions @@ -50,6 +53,9 @@ describe('@codesift/mcp server', () => { expect(getToolDefinitions().find((tool) => tool.name === 'grep_code')?.inputSchema.properties).toMatchObject({ max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: 700 } }) + expect(getToolDefinitions().find((tool) => tool.name === 'read_chunk')?.inputSchema.properties).toMatchObject({ + max_tokens: { type: 'integer', minimum: MIN_MCP_READ_CHUNK_MAX_TOKENS, maximum: 4000, default: DEFAULT_MCP_READ_CHUNK_MAX_TOKENS } + }) expect(getToolDefinitions().find((tool) => tool.name === 'search_code')?.inputSchema.properties).toMatchObject({ context: { type: 'string', enum: ['sig', 'body'] }, with_usages: { type: 'boolean' } @@ -466,6 +472,52 @@ describe('grep_code MCP output budgeting', () => { }) }) +describe('read_chunk MCP output budgeting', () => { + it('leaves under-budget chunk content byte-for-byte unchanged', () => { + const content = 'export function demoValue(): string {\n return "demo"\n}\n' + + expect(formatMcpReadChunk(content, { maxTokens: 80 })).toBe(content) + }) + + it('preserves the first source content and appends a compact truncation marker', () => { + const content = Array.from({ length: 30 }, (_, index) => `export const noisyValue${index} = '${'x'.repeat(40)}'`).join('\n') + const output = formatMcpReadChunk(content, { maxTokens: 45 }) + + expect(output).toContain("export const noisyValue0 = '") + expect(output).toContain('content_truncated=true; raise max_tokens or narrow the chunk/range.') + expect(output).not.toContain('noisyValue29') + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(45) + }) + + it('keeps tiny-budget output within the approximate budget with a recognizable marker', () => { + const output = formatMcpReadChunk('const value = "' + 'x'.repeat(400) + '"', { maxTokens: 6 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(6) + expect(output).toContain('content_truncated') + }) + + it('budgets noisy read_chunk output through callMcpTool while preserving the first lines', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-mcp-read-budget-')) + temporaryDirectories.push(repoRoot) + + await mkdir(join(repoRoot, 'src'), { recursive: true }) + await writeFile( + join(repoRoot, 'src', 'noisy.ts'), + Array.from({ length: 120 }, (_, index) => `export const noisyValue${index} = '${'x'.repeat(40)}'`).join('\n'), + 'utf8' + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + + const output = await callMcpTool(repo, 'read_chunk', { id: 'src/noisy.ts:1-120', max_tokens: 60 }) + expect(output).toContain("export const noisyValue0 = '") + expect(output).toContain('content_truncated=true') + expect(output).not.toContain('noisyValue119') + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(60) + }) +}) + function makeGrepHit(overrides: Partial = {}): GrepHit { return { file: 'src/auth.ts', From 29c1aef27eebf4fc385fccf141e2596217b3fbde Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 02:29:33 +0530 Subject: [PATCH 05/13] Budget MCP find_symbol output --- packages/eval/src/index.ts | 6 +- packages/mcp/src/index.ts | 138 ++++++++++++++++++++++++++++++---- packages/mcp/test/mcp.test.ts | 110 ++++++++++++++++++++++++++- 3 files changed, 234 insertions(+), 20 deletions(-) diff --git a/packages/eval/src/index.ts b/packages/eval/src/index.ts index 60b3a9a..e157bb1 100644 --- a/packages/eval/src/index.ts +++ b/packages/eval/src/index.ts @@ -7,7 +7,7 @@ import { execFile as execFileCallback, spawn, type ChildProcessWithoutNullStream import { promisify } from 'node:util' import { VOYAGE_RERANK_PROVIDER_ID, openRepo, type FindSymbolOptions, type GrepHit, type GrepOptions, type Range, type SearchHit, type SearchOptions, type SymbolDefinition } from '@codesift/core' -import { DEFAULT_MCP_GREP_MAX_TOKENS, DEFAULT_MCP_READ_CHUNK_MAX_TOKENS, formatMcpGrepHits, formatMcpReadChunk, formatMcpSearchHits, formatMcpSymbols } from '@codesift/mcp' +import { DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS, DEFAULT_MCP_GREP_MAX_TOKENS, DEFAULT_MCP_READ_CHUNK_MAX_TOKENS, formatMcpGrepHits, formatMcpReadChunk, formatMcpSearchHits, formatMcpSymbols } from '@codesift/mcp' const execFile = promisify(execFileCallback) @@ -681,7 +681,7 @@ function buildCodesiftToolCall(query: GoldenQuery, resultLimit: number): { name: switch (query.queryType) { case 'symbol-def': case 'exact-identifier': - return { name: 'find_symbol', arguments: { name: query.query, ...pathGlob } } + return { name: 'find_symbol', arguments: { name: query.query, max_tokens: DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS, ...pathGlob } } case 'string-literal': case 'error-trace': return { @@ -829,7 +829,7 @@ async function runCodesiftPolicy(repo: Awaited>, que // non-top match) still forces a follow-up read — count it as 2 calls. A // MULTI-target lookup ("where are all the Xs") is answered by the location set // in a single call, so it is never charged a per-body follow-up. - let tokensToResolution = estimateTokenCount(formatMcpSymbols(definitions)) + let tokensToResolution = estimateTokenCount(formatMcpSymbols(definitions, { maxTokens: DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS })) let callsToResolution = 1 if (matchingRank !== null && query.expected.length === 1) { diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index fdea6da..bcffa2e 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -24,6 +24,7 @@ export { HttpMcpServerHandle, createHttpServerHandle } from './http.js' export const DEFAULT_SEARCH_K = CORE_DEFAULT_SEARCH_K export const DEFAULT_MCP_GREP_MAX_TOKENS = 700 +export const DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS = 700 export const DEFAULT_MCP_READ_CHUNK_MAX_TOKENS = 1000 export const MIN_MCP_READ_CHUNK_MAX_TOKENS = 6 export const MAX_MCP_READ_CHUNK_MAX_TOKENS = 4000 @@ -69,6 +70,7 @@ export interface FindSymbolArgs { kind?: FindSymbolOptions['kind'] | undefined path_glob?: string | undefined with_body?: boolean | undefined + max_tokens?: number | undefined } export interface GrepCodeArgs { @@ -94,6 +96,10 @@ export interface FormatMcpReadChunkOptions { maxTokens?: number | undefined } +export interface FormatMcpSymbolsOptions { + maxTokens?: number | undefined +} + export interface ReadChunkArgs { id: string context_lines?: number | undefined @@ -170,7 +176,8 @@ const findSymbolInputSchema = { name: z.string().min(1).describe('Exact or partial symbol name, e.g. TokenVerifier or verifyJwtToken.'), kind: kindFilterSchema.optional().describe('Optional symbol kind filter.'), path_glob: z.string().min(1).optional().describe('Repo-relative glob, e.g. "src/auth/**".'), - with_body: z.boolean().optional().describe('Inline the top exact match\'s full definition body so it resolves in one call. Default true; set false for a compact name→location list.') + with_body: z.boolean().optional().describe('Inline the top exact match\'s full definition body so it resolves in one call. Default true; set false for a compact name→location list.'), + max_tokens: z.number().int().positive().max(4000).optional().describe('Maximum approximate tokens to return across symbol definitions. Default 700.') } const grepCodeInputSchema = { @@ -259,7 +266,8 @@ export function getToolDefinitions(): readonly McpToolDefinition[] { name: { type: 'string' }, kind: kindJsonSchema(), path_glob: { type: 'string' }, - with_body: { type: 'boolean', default: true } + with_body: { type: 'boolean', default: true }, + max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS } }) }, { @@ -406,7 +414,10 @@ export async function callMcpTool(repo: Repo, name: McpToolName, args: unknown): case 'search_code': return formatMcpSearchHits(await router.searchCode(searchCodeArgsSchema.parse(args))) case 'find_symbol': - return formatMcpSymbols(await router.findSymbol(findSymbolArgsSchema.parse(args))) + { + const parsed = findSymbolArgsSchema.parse(args) + return formatMcpSymbols(await router.findSymbol(parsed), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS }) + } case 'grep_code': { const parsed = grepCodeArgsSchema.parse(args) @@ -487,7 +498,7 @@ export function createSdkServer(repo: Repo): McpServer { textResult(formatMcpSearchHits(await router.searchCode(args))) ) server.registerTool('find_symbol', { description: toolDescription('find_symbol'), inputSchema: findSymbolInputSchema }, async (args) => - textResult(formatMcpSymbols(await router.findSymbol(args))) + textResult(formatMcpSymbols(await router.findSymbol(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS })) ) server.registerTool('grep_code', { description: toolDescription('grep_code'), inputSchema: grepCodeInputSchema }, async (args) => textResult(formatMcpGrepHits(await router.grepCode(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_GREP_MAX_TOKENS })) @@ -596,22 +607,117 @@ function formatHitSymbol(hit: SearchHit): string { return ` ${[hit.parent, hit.symbol].filter(Boolean).join(' > ')}` } -export function formatMcpSymbols(definitions: SymbolDefinition[]): string { +export function formatMcpSymbols(definitions: SymbolDefinition[], options: FormatMcpSymbolsOptions = {}): string { if (definitions.length === 0) { return 'no_symbols' } - return definitions - .map((definition, index) => { - const header = `#${index + 1} ${definition.kind} ${definition.name} ${definition.file}:${formatRange(definition.range.startLine, definition.range.endLine)}` - // The top exact match carries a paste-ready body so the identifier resolves - // in a single call; render it with the same line-numbered block as search. - if (definition.body !== undefined) { - return `${header}\n${formatBodyBlock(definition.body, definition.range.startLine)}` - } - return header - }) - .join('\n') + const rendered = definitions.map(formatMcpSymbolDefinition) + const output = rendered.join('\n') + const maxTokens = options.maxTokens + if (maxTokens === undefined || output.length <= maxTokens * 4) { + return output + } + + return fitSymbolOutputForBudget(definitions, rendered, maxTokens * 4) +} + +function formatMcpSymbolDefinition(definition: SymbolDefinition, index: number): string { + const header = formatMcpSymbolHeader(definition, index) + // The top exact match carries a paste-ready body so the identifier resolves + // in a single call; render it with the same line-numbered block as search. + if (definition.body !== undefined) { + return `${header}\n${formatBodyBlock(definition.body, definition.range.startLine)}` + } + return header +} + +function formatMcpSymbolHeader(definition: SymbolDefinition, index: number): string { + return `#${index + 1} ${definition.kind} ${definition.name} ${definition.file}:${formatRange(definition.range.startLine, definition.range.endLine)}` +} + +function fitSymbolOutputForBudget(definitions: SymbolDefinition[], rendered: string[], maxChars: number): string { + const omittedAfterFirst = definitions.length - 1 + const firstOnlyMarker = omittedAfterFirst > 0 + ? symbolOmissionMarker(omittedAfterFirst, Math.max(0, maxChars - 2)) + : symbolBodyTruncationMarker(Math.max(0, maxChars - 2)) + const firstOnlyBudget = maxCharsForFirstHit(maxChars, firstOnlyMarker) + if (rendered[0]!.length > firstOnlyBudget) { + const first = formatMcpSymbolDefinitionForBudget(definitions[0]!, 0, firstOnlyBudget) + return firstOnlyMarker ? [first, firstOnlyMarker].filter(Boolean).join('\n') : first + } + + const parts = [rendered[0]!] + for (let index = 1; index < rendered.length; index += 1) { + const candidate = [...parts, rendered[index]!] + const omitted = definitions.length - 1 - (candidate.length - 1) + const marker = omitted > 0 ? symbolOmissionMarker(omitted, Math.max(0, maxChars - 2)) : undefined + const candidateOutput = marker ? [...candidate, marker].join('\n') : candidate.join('\n') + if (candidateOutput.length <= maxChars) { + parts.push(rendered[index]!) + continue + } + break + } + + const omitted = definitions.length - parts.length + const marker = omitted > 0 ? symbolOmissionMarker(omitted, Math.max(0, maxChars - 2)) : undefined + const outputParts = marker ? [...parts, marker] : parts + + const output = outputParts.join('\n') + return output.length <= maxChars ? output : truncateWithEllipsis(output, maxChars) +} + +function formatMcpSymbolDefinitionForBudget(definition: SymbolDefinition, index: number, maxChars: number): string { + const header = formatMcpSymbolHeader(definition, index) + if (definition.body === undefined || maxChars <= header.length + 1) { + return truncateWithEllipsis(header, maxChars) + } + + const bodyBudget = maxChars - header.length - 1 + const body = truncateSymbolBodyBlock(formatBodyBlock(definition.body, definition.range.startLine), bodyBudget) + return body ? `${header}\n${body}` : truncateWithEllipsis(header, maxChars) +} + +function truncateSymbolBodyBlock(body: string, maxChars: number): string { + if (maxChars <= 0) { + return '' + } + if (body.length <= maxChars) { + return body + } + + return truncateWithEllipsis(body, maxChars).replace(/\n+$/g, '') +} + +function symbolOmissionMarker(omitted: number, maxChars?: number): string { + const full = `symbols_omitted=${omitted}; refine name/kind/path_glob or raise max_tokens.` + if (maxChars === undefined || full.length <= maxChars) { + return full + } + + const compact = `symbols_omitted=${omitted}; raise max_tokens` + if (compact.length <= maxChars) { + return compact + } + + const bare = `symbols_omitted=${omitted}` + return bare.length <= maxChars ? bare : '' +} + +function symbolBodyTruncationMarker(maxChars?: number): string { + const full = 'symbol_body_truncated=true; refine name/kind/path_glob or raise max_tokens.' + if (maxChars === undefined || full.length <= maxChars) { + return full + } + + const compact = 'symbol_body_truncated=true; raise max_tokens' + if (compact.length <= maxChars) { + return compact + } + + const bare = 'symbol_body_truncated=true' + return bare.length <= maxChars ? bare : '' } export function formatMcpGrepHits(hits: GrepHit[], options: FormatMcpGrepHitsOptions = {}): string { diff --git a/packages/mcp/test/mcp.test.ts b/packages/mcp/test/mcp.test.ts index e7c7583..3c01f4c 100644 --- a/packages/mcp/test/mcp.test.ts +++ b/packages/mcp/test/mcp.test.ts @@ -5,9 +5,10 @@ import { join } from 'node:path' import { afterEach, describe, expect, it } from 'vitest' -import { openRepo, registerEmbeddingProvider, type GrepHit, type SearchHit } from '@codesift/core' +import { openRepo, registerEmbeddingProvider, type GrepHit, type SearchHit, type SymbolDefinition } from '@codesift/core' import { + DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS, DEFAULT_MCP_READ_CHUNK_MAX_TOKENS, DEFAULT_SEARCH_K, MIN_MCP_READ_CHUNK_MAX_TOKENS, @@ -53,6 +54,9 @@ describe('@codesift/mcp server', () => { expect(getToolDefinitions().find((tool) => tool.name === 'grep_code')?.inputSchema.properties).toMatchObject({ max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: 700 } }) + expect(getToolDefinitions().find((tool) => tool.name === 'find_symbol')?.inputSchema.properties).toMatchObject({ + max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS } + }) expect(getToolDefinitions().find((tool) => tool.name === 'read_chunk')?.inputSchema.properties).toMatchObject({ max_tokens: { type: 'integer', minimum: MIN_MCP_READ_CHUNK_MAX_TOKENS, maximum: 4000, default: DEFAULT_MCP_READ_CHUNK_MAX_TOKENS } }) @@ -144,6 +148,30 @@ describe('@codesift/mcp server', () => { expect(DEFAULT_SEARCH_K).toBe(8) }) + it('budgets find_symbol output through the MCP call surface', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-mcp-symbol-budget-')) + temporaryDirectories.push(repoRoot) + + await mkdir(join(repoRoot, 'src'), { recursive: true }) + await Promise.all( + Array.from({ length: 6 }, async (_, index) => { + await writeFile( + join(repoRoot, 'src', `shared${index}.ts`), + `export function sharedName(): string {\n const value = 'shared-${index}-${'x'.repeat(40)}'\n return value\n}\n`, + 'utf8' + ) + }) + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + const output = await callMcpTool(repo, 'find_symbol', { name: 'sharedName', max_tokens: 35 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(35) + expect(output).toContain('#1 function sharedName') + expect(output).toContain('symbols_omitted=') + }) + it('creates server handles', async () => { const repo = await openRepo(process.cwd()) const stdio = createStdioServer(repo) @@ -269,6 +297,86 @@ function makeHit(overrides: Partial): SearchHit { } } +function makeSymbol(overrides: Partial> & { body?: string | undefined } = {}): SymbolDefinition { + const { body, ...rest } = overrides + const definition: SymbolDefinition = { + id: 'src/auth.ts:10-13@abcdef01', + name: 'verifyToken', + kind: 'function', + file: 'src/auth.ts', + range: { startLine: 10, endLine: 13 }, + body: 'export function verifyToken(token: string): boolean {\n return token.length > 0\n}', + ...rest + } + + if ('body' in overrides) { + if (body === undefined) { + delete definition.body + } else { + definition.body = body + } + } + + return definition +} + +describe('formatMcpSymbols budgeted output', () => { + it('preserves no_symbols and under-budget output byte-for-byte', () => { + const definitions = [makeSymbol()] + const expected = [ + '#1 function verifyToken src/auth.ts:10-13', + '10 | export function verifyToken(token: string): boolean {', + '11 | return token.length > 0', + '12 | }' + ].join('\n') + + expect(formatMcpSymbols([])).toBe('no_symbols') + expect(formatMcpSymbols(definitions)).toBe(expected) + expect(formatMcpSymbols(definitions, { maxTokens: 500 })).toBe(expected) + }) + + it('preserves the first definition and omits later rows when over budget', () => { + const definitions = [ + makeSymbol({ + body: `export function verifyToken(token: string): boolean {\n const padded = '${'x'.repeat(180)}'\n return token + padded !== ''\n}` + }), + makeSymbol({ name: 'verifyToken', file: 'src/legacy.ts', range: { startLine: 1, endLine: 1 }, body: undefined }), + makeSymbol({ name: 'verifyToken', file: 'src/other.ts', range: { startLine: 5, endLine: 5 }, body: undefined }) + ] + + const output = formatMcpSymbols(definitions, { maxTokens: 45 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(45) + expect(output).toContain('#1 function verifyToken src/auth.ts:10-13') + expect(output).toContain('10 | export function verifyToken') + expect(output).toContain('symbols_omitted=') + expect(output).not.toContain('#2 function verifyToken') + }) + + it('marks a single over-budget symbol body as truncated instead of omitted', () => { + const output = formatMcpSymbols([ + makeSymbol({ + body: `export function verifyToken(token: string): boolean {\n const padded = '${'x'.repeat(180)}'\n return token + padded !== ''\n}` + }) + ], { maxTokens: 35 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(35) + expect(output).toContain('#1 function verifyToken src/auth.ts:10-13') + expect(output).toContain('symbol_body_truncated=true') + expect(output).not.toContain('symbols_omitted=0') + }) + + it('keeps a recognizable omission marker for tiny budgets when it fits', () => { + const output = formatMcpSymbols([ + makeSymbol({ body: `export function verifyToken() {\n return '${'x'.repeat(100)}'\n}` }), + makeSymbol({ file: 'src/other.ts', body: undefined }) + ], { maxTokens: 5 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(5) + expect(output).toContain('symbols_omitted=1') + }) +}) + describe('formatMcpSearchHits structure-preserving output', () => { it('renders an inlined body as a line-numbered block starting at range.startLine with original indentation', () => { const body = 'function verify(token: string): boolean {\n if (!token) {\n return false\n }\n return true\n}' From ae28c8be602c38d3ab2c37e86618e47e12d9769f Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 03:27:29 +0530 Subject: [PATCH 06/13] Budget search_code MCP output --- packages/eval/src/index.ts | 23 +- packages/eval/test/eval.test.ts | 2 + packages/mcp/src/index.ts | 512 ++++++++++++++++++++++++++++++-- packages/mcp/test/mcp.test.ts | 167 +++++++++++ 4 files changed, 675 insertions(+), 29 deletions(-) diff --git a/packages/eval/src/index.ts b/packages/eval/src/index.ts index e157bb1..cb7d45c 100644 --- a/packages/eval/src/index.ts +++ b/packages/eval/src/index.ts @@ -7,7 +7,16 @@ import { execFile as execFileCallback, spawn, type ChildProcessWithoutNullStream import { promisify } from 'node:util' import { VOYAGE_RERANK_PROVIDER_ID, openRepo, type FindSymbolOptions, type GrepHit, type GrepOptions, type Range, type SearchHit, type SearchOptions, type SymbolDefinition } from '@codesift/core' -import { DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS, DEFAULT_MCP_GREP_MAX_TOKENS, DEFAULT_MCP_READ_CHUNK_MAX_TOKENS, formatMcpGrepHits, formatMcpReadChunk, formatMcpSearchHits, formatMcpSymbols } from '@codesift/mcp' +import { + DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS, + DEFAULT_MCP_GREP_MAX_TOKENS, + DEFAULT_MCP_READ_CHUNK_MAX_TOKENS, + DEFAULT_MCP_SEARCH_MAX_TOKENS, + formatMcpGrepHits, + formatMcpReadChunk, + formatMcpSearchHits, + formatMcpSymbols +} from '@codesift/mcp' const execFile = promisify(execFileCallback) @@ -129,6 +138,7 @@ export interface UsageReportEntry { missingUsageFiles: string[] withoutUsagesCallsToResolution: number withUsagesCallsToResolution: number + withUsagesTokensToResolution: number } export interface RerankReportEntry { @@ -366,7 +376,7 @@ export function formatSummary(summary: EvalSummary): string { lines.push('with_usages report-only') for (const entry of summary.usagesReport) { lines.push( - ` - ${entry.queryId}: usageRecall=${formatRatio(entry.usageRecall)} savedCallsΔ=${formatSignedNumber(entry.withUsagesCallsDelta)} bundled=${formatNumber(entry.withUsagesCallsToResolution)} baseline=${formatNumber(entry.withoutUsagesCallsToResolution)} matched=${entry.matchedUsageFiles.join(', ') || 'none'} missing=${entry.missingUsageFiles.join(', ') || 'none'}` + ` - ${entry.queryId}: usageRecall=${formatRatio(entry.usageRecall)} savedCallsΔ=${formatSignedNumber(entry.withUsagesCallsDelta)} bundledCalls=${formatNumber(entry.withUsagesCallsToResolution)} baselineCalls=${formatNumber(entry.withoutUsagesCallsToResolution)} bundledTokens=${formatNumber(entry.withUsagesTokensToResolution)} matched=${entry.matchedUsageFiles.join(', ') || 'none'} missing=${entry.missingUsageFiles.join(', ') || 'none'}` ) } lines.push('') @@ -579,7 +589,8 @@ async function measureWithUsagesReport( matchedUsageFiles, missingUsageFiles, withoutUsagesCallsToResolution, - withUsagesCallsToResolution: withUsagesRun.callsToResolution + withUsagesCallsToResolution: withUsagesRun.callsToResolution, + withUsagesTokensToResolution: withUsagesRun.tokensToResolution } } @@ -691,7 +702,7 @@ function buildCodesiftToolCall(query: GoldenQuery, resultLimit: number): { name: case 'nl-concept': return { name: 'search_code', - arguments: { query: query.query, k: resultLimit, max_tokens: 700, ...pathGlob } + arguments: { query: query.query, k: resultLimit, max_tokens: DEFAULT_MCP_SEARCH_MAX_TOKENS, ...pathGlob } } } } @@ -741,7 +752,7 @@ function rpcId(message: unknown): number | undefined { function buildSearchOptions(query: GoldenQuery, resultLimit: number, overrides: Partial = {}): SearchOptions { const options: SearchOptions = { k: resultLimit, - maxTokens: 700, + maxTokens: DEFAULT_MCP_SEARCH_MAX_TOKENS, ...overrides } @@ -764,7 +775,7 @@ async function runSearchPolicy( const score = scoreCandidates(candidates, query, inspectionLimit) const matchingRank = score.matchingRank - let tokensToResolution = estimateTokenCount(formatMcpSearchHits(hits)) + let tokensToResolution = estimateTokenCount(formatMcpSearchHits(hits, { maxTokens: DEFAULT_MCP_SEARCH_MAX_TOKENS })) let callsToResolution = 1 if (matchingRank !== null && query.expected.length === 1) { diff --git a/packages/eval/test/eval.test.ts b/packages/eval/test/eval.test.ts index 38cc7c2..4509ef9 100644 --- a/packages/eval/test/eval.test.ts +++ b/packages/eval/test/eval.test.ts @@ -204,10 +204,12 @@ describe('@codesift/eval', () => { missingUsageFiles: [] }) expect(summary.usagesReport[0]?.withUsagesCallsDelta).toBeGreaterThanOrEqual(1) + expect(summary.usagesReport[0]?.withUsagesTokensToResolution).toBeGreaterThan(0) expect(summary.usagesReport[0]?.withUsagesCallsToResolution).toBeLessThan(summary.usagesReport[0]?.withoutUsagesCallsToResolution ?? 0) const rendered = formatSummary(summary) expect(rendered).toContain('with_usages report-only') + expect(rendered).toContain('bundledTokens=') expect(rendered).toContain('usageRecall=1.00') }, 30_000) diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index bcffa2e..3e3c49c 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -23,6 +23,7 @@ import { createHttpServerHandle } from './http.js' export { HttpMcpServerHandle, createHttpServerHandle } from './http.js' export const DEFAULT_SEARCH_K = CORE_DEFAULT_SEARCH_K +export const DEFAULT_MCP_SEARCH_MAX_TOKENS = 700 export const DEFAULT_MCP_GREP_MAX_TOKENS = 700 export const DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS = 700 export const DEFAULT_MCP_READ_CHUNK_MAX_TOKENS = 1000 @@ -92,6 +93,10 @@ export interface FormatMcpGrepHitsOptions { maxTokens?: number | undefined } +export interface FormatMcpSearchHitsOptions { + maxTokens?: number | undefined +} + export interface FormatMcpReadChunkOptions { maxTokens?: number | undefined } @@ -253,7 +258,7 @@ export function getToolDefinitions(): readonly McpToolDefinition[] { lang: { type: 'array', items: { type: 'string' } }, path_glob: { type: 'string' }, kind: kindJsonSchema(), - max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: 700 }, + max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: DEFAULT_MCP_SEARCH_MAX_TOKENS }, single_best: { type: 'boolean' }, context: { type: 'string', enum: ['sig', 'body'] }, with_usages: { type: 'boolean' } @@ -324,7 +329,7 @@ export function createRouter(repo: Repo): McpRouter { options.kind = args.kind } - options.maxTokens = args.max_tokens ?? 700 + options.maxTokens = args.max_tokens ?? DEFAULT_MCP_SEARCH_MAX_TOKENS if (args.single_best !== undefined) { options.singleBest = args.single_best } @@ -412,7 +417,10 @@ export async function callMcpTool(repo: Repo, name: McpToolName, args: unknown): switch (name) { case 'search_code': - return formatMcpSearchHits(await router.searchCode(searchCodeArgsSchema.parse(args))) + { + const parsed = searchCodeArgsSchema.parse(args) + return formatMcpSearchHits(await router.searchCode(parsed), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_SEARCH_MAX_TOKENS }) + } case 'find_symbol': { const parsed = findSymbolArgsSchema.parse(args) @@ -495,7 +503,7 @@ export function createSdkServer(repo: Repo): McpServer { ) server.registerTool('search_code', { description: toolDescription('search_code'), inputSchema: searchCodeInputSchema }, async (args) => - textResult(formatMcpSearchHits(await router.searchCode(args))) + textResult(formatMcpSearchHits(await router.searchCode(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_SEARCH_MAX_TOKENS })) ) server.registerTool('find_symbol', { description: toolDescription('find_symbol'), inputSchema: findSymbolInputSchema }, async (args) => textResult(formatMcpSymbols(await router.findSymbol(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS })) @@ -553,37 +561,314 @@ function toolDescription(name: McpToolName): string { return tool?.description ?? name } -export function formatMcpSearchHits(hits: SearchHit[]): string { +export function formatMcpSearchHits(hits: SearchHit[], options: FormatMcpSearchHitsOptions = {}): string { if (hits.length === 0) { return 'no_hits' } + const ambiguityHint = formatSearchAmbiguityHint(hits[0]) + const tokensLine = formatSearchTokensLine(hits) + const renderedHits = hits.map((hit) => formatMcpSearchHit(hit)) + const output = joinSections([ambiguityHint, ...renderedHits, tokensLine]) + const maxTokens = options.maxTokens + if (maxTokens === undefined || output.length <= maxTokens * 4) { + return output + } + + return fitSearchOutputForBudget(hits, renderedHits, ambiguityHint, tokensLine, maxTokens * 4) +} + +function formatSearchAmbiguityHint(hit: SearchHit | undefined): string { // An identifier that collides across multiple definitions is returned as a candidate // set, not a single confident answer — lead with the collision count so the caller // disambiguates instead of trusting the first row. - const ambiguousDefCount = hits[0]?.ambiguousDefCount - const ambiguityHint = ambiguousDefCount && ambiguousDefCount >= 2 ? `ambiguous: ${ambiguousDefCount} defs\n` : '' + const ambiguousDefCount = hit?.ambiguousDefCount + return ambiguousDefCount && ambiguousDefCount >= 2 ? `ambiguous: ${ambiguousDefCount} defs` : '' +} +function formatSearchTokensLine(hits: SearchHit[]): string { const tokensReturned = hits.reduce((sum, hit) => sum + hit.tokensReturned, 0) - const body = hits - .map((hit) => { - const symbol = formatHitSymbol(hit) - const generated = hit.generated ? ' [generated]' : '' - const stale = hit.stale ? ' [stale]' : '' - const header = `${hit.reason} ${formatChunkId(hit.id)}${symbol}${generated}${stale}` + return `tokensReturned=${tokensReturned}` +} + +function formatMcpSearchHit(hit: SearchHit): string { + const sections = [formatSearchHitHeader(hit)] + const content = formatSearchHitContent(hit) + if (content) { + sections.push(content) + } + if (hit.usages?.length) { + sections.push(formatUsageBlock(hit.usages)) + } + return joinSections(sections) +} + +function formatSearchHitHeader(hit: SearchHit): string { + const symbol = formatHitSymbol(hit) + const generated = hit.generated ? ' [generated]' : '' + const stale = hit.stale ? ' [stale]' : '' + return `${hit.reason} ${formatChunkId(hit.id)}${symbol}${generated}${stale}` +} + +function formatSearchHitContent(hit: SearchHit): string { + if (hit.body !== undefined) { + return formatBodyBlock(hit.body, hit.range.startLine ?? hit.snippetRange.startLine) + } + + return compactHitSnippet(hit.snippet, hit.snippetRange.startLine ?? hit.range.startLine, 4) +} + +function fitSearchOutputForBudget( + hits: SearchHit[], + renderedHits: string[], + ambiguityHint: string, + tokensLine: string, + maxChars: number +): string { + const leadingSections = ambiguityHint ? [ambiguityHint] : [] + const hitsOnlyOutput = joinSections([...leadingSections, ...renderedHits]) + if (hitsOnlyOutput.length <= maxChars) { + return appendSearchTokensLineIfFits(hitsOnlyOutput, tokensLine, maxChars) + } - if (hit.body !== undefined) { - const block = `${header}\n${formatBodyBlock(hit.body, hit.range.startLine ?? hit.snippetRange.startLine)}` - return hit.usages?.length ? `${block}\n${formatUsageBlock(hit.usages)}` : block + const keptHits: string[] = [] + for (let index = 0; index < renderedHits.length; index += 1) { + const candidateHits = [...keptHits, renderedHits[index]!] + const omitted = renderedHits.length - candidateHits.length + const marker = omitted > 0 + ? searchHitOmissionMarker(omitted, remainingCharsForTrailingLine([...leadingSections, ...candidateHits], maxChars)) + : '' + const candidateOutput = joinSections([...leadingSections, ...candidateHits, marker]) + if (candidateOutput.length <= maxChars) { + keptHits.push(renderedHits[index]!) + continue + } + + if (keptHits.length === 0) { + const output = fitSearchSingleHitOutput(hits[0]!, leadingSections, omitted, maxChars) + return appendSearchTokensLineIfFits(output, tokensLine, maxChars) + } + + break + } + + const outputHits = [...keptHits] + let omitted = renderedHits.length - outputHits.length + let marker = omitted > 0 + ? searchHitOmissionMarker(omitted, remainingCharsForTrailingLine([...leadingSections, ...outputHits], maxChars)) + : '' + + while (omitted > 0 && !marker && outputHits.length > 1) { + outputHits.pop() + omitted += 1 + marker = searchHitOmissionMarker(omitted, remainingCharsForTrailingLine([...leadingSections, ...outputHits], maxChars)) + } + + if (omitted > 0 && !marker && outputHits.length === 1) { + const output = fitSearchSingleHitOutput(hits[0]!, leadingSections, omitted, maxChars) + return appendSearchTokensLineIfFits(output, tokensLine, maxChars) + } + + const output = joinSections([...leadingSections, ...outputHits, marker]) + return appendSearchTokensLineIfFits(output, tokensLine, maxChars) +} + +function fitSearchSingleHitOutput(hit: SearchHit, leadingSections: string[], omittedHits: number, maxChars: number): string { + const leadingOnly = joinSections(leadingSections) + const header = formatSearchHitHeader(hit) + const trailingVariants = omittedHits > 0 ? searchHitOmissionMarkerVariants(omittedHits) : [''] + + for (const trailingLine of trailingVariants) { + const reservedHitBudget = maxChars + - (leadingOnly ? leadingOnly.length + 1 : 0) + - (trailingLine ? trailingLine.length + 1 : 0) + if (reservedHitBudget <= 0) { + continue + } + + const hitOutput = formatSearchHitForBudget(hit, reservedHitBudget) + if (hitOutput) { + const candidate = joinSections([...leadingSections, hitOutput, trailingLine]) + if (candidate.length <= maxChars) { + return candidate } + } + } - const snippet = compactHitSnippet(hit.snippet, hit.snippetRange.startLine ?? hit.range.startLine, 4) - const block = snippet ? `${header}\n${snippet}` : header - return hit.usages?.length ? `${block}\n${formatUsageBlock(hit.usages)}` : block - }) - .join('\n') + const hitBudget = maxChars - (leadingOnly ? leadingOnly.length + 1 : 0) + if (omittedHits === 0 && hitBudget > 0 && hitBudget >= header.length) { + const hitOnly = formatSearchHitForBudget(hit, hitBudget) + if (hitOnly) { + const candidate = joinSections([...leadingSections, hitOnly]) + if (candidate.length <= maxChars) { + return candidate + } + } + } + + if (omittedHits > 0) { + const markerWithLeading = fitMarkerWithLeadingSections(leadingSections, trailingVariants, maxChars) + if (markerWithLeading) { + return markerWithLeading + } + + return bestEffortOmissionMarker(trailingVariants, maxChars) + } + + for (const trailingLine of trailingVariants) { + if (trailingLine && trailingLine.length <= maxChars) { + return trailingLine + } + } + + if (leadingOnly && leadingOnly.length <= maxChars) { + return leadingOnly + } + + return formatSearchHitForBudget(hit, maxChars) +} + +function formatSearchHitForBudget(hit: SearchHit, maxChars: number): string { + const full = formatMcpSearchHit(hit) + if (full.length <= maxChars) { + return full + } + + const header = formatSearchHitHeader(hit) + const content = formatSearchHitContent(hit) + const usageCount = hit.usages?.length ?? 0 + if (maxChars <= header.length) { + if (usageCount > 0) { + return reserveUsageMarkerForSearchHit(header, content || undefined, maxChars, usageCount) + } + return truncateWithEllipsis(header, maxChars) + } + + const sections = [header] + + if (content) { + const remainingAfterHeader = maxChars - header.length - 1 + if (content.length + 1 <= remainingAfterHeader) { + sections.push(content) + if (usageCount > 0) { + const usageBlock = formatUsageBlockForBudget(hit.usages!, maxChars - joinSections(sections).length - 1) + if (usageBlock) { + sections.push(usageBlock) + return joinSections(sections) + } + + return reserveUsageMarkerForSearchHit(header, content, maxChars, usageCount) + } + return joinSections(sections) + } + + if (usageCount > 0) { + return reserveUsageMarkerForSearchHit(header, content, maxChars, usageCount) + } + + const truncatedContent = truncateSearchBlock(content, remainingAfterHeader) + if (truncatedContent) { + sections.push(truncatedContent) + } + const fallback = joinSections(sections) + return fallback.length <= maxChars ? fallback : truncateWithEllipsis(fallback, maxChars) + } + + if (usageCount > 0) { + const usageBlock = formatUsageBlockForBudget(hit.usages!, maxChars - header.length - 1) + if (usageBlock) { + return joinSections([header, usageBlock]) + } + + return reserveUsageMarkerForSearchHit(header, undefined, maxChars, usageCount) + } + + return truncateWithEllipsis(header, maxChars) +} + +function reserveUsageMarkerForSearchHit(header: string, content: string | undefined, maxChars: number, omittedUsages: number): string { + const markers = usageOmissionMarkerVariants(omittedUsages) + + if (content !== undefined) { + for (const marker of markers) { + const contentBudget = maxChars - header.length - 1 - marker.length - 1 + if (contentBudget <= 0) { + continue + } + + const truncatedContent = truncateSearchBlock(content, contentBudget) + if (!truncatedContent) { + continue + } + + const candidate = joinSections([header, truncatedContent, marker]) + if (candidate.length <= maxChars) { + return candidate + } + } + } + + const headerWithMarker = fitMarkerWithLeadingSections([header], markers, maxChars) + if (headerWithMarker) { + return headerWithMarker + } + + if (content !== undefined) { + for (const marker of markers) { + const contentBudget = maxChars - marker.length - 1 + if (contentBudget <= 0) { + continue + } + + const truncatedContent = truncateSearchBlock(content, contentBudget) + if (!truncatedContent) { + continue + } + + const candidate = joinSections([truncatedContent, marker]) + if (candidate.length <= maxChars) { + return candidate + } + } + } + + return bestEffortOmissionMarker(markers, maxChars) +} + +function fitMarkerWithLeadingSections(leadingSections: string[], markers: string[], maxChars: number): string { + for (const marker of markers) { + const markerWithLeading = joinSections([...leadingSections, marker]) + if (markerWithLeading && markerWithLeading.length <= maxChars) { + return markerWithLeading + } + } + + const leadingOutput = joinSections(leadingSections) + const remainingForMarker = maxChars - (leadingOutput ? leadingOutput.length + 1 : 0) + if (remainingForMarker > 0) { + const truncatedMarker = bestEffortOmissionMarker(markers, remainingForMarker) + if (truncatedMarker) { + const candidate = joinSections([...leadingSections, truncatedMarker]) + if (candidate.length <= maxChars) { + return candidate + } + } + } - return `${ambiguityHint}${body}\ntokensReturned=${tokensReturned}` + return '' +} + +function bestEffortOmissionMarker(markers: string[], maxChars: number): string { + if (maxChars <= 0 || markers.length === 0) { + return '' + } + + const fittingMarker = firstMarkerThatFits(markers, maxChars) + if (fittingMarker) { + return fittingMarker + } + + return truncateWithEllipsis(markers[0]!, maxChars) } function formatBodyBlock(body: string, startLine: number): string { @@ -594,11 +879,99 @@ function formatBodyBlock(body: string, startLine: number): string { function formatUsageBlock(usages: SymbolUsage[]): string { const lines = ['usages (import-resolved):'] for (const usage of usages) { - lines.push(`- ${usage.file}:${usage.line} | ${usage.snippet}`) + lines.push(formatUsageLineRaw(usage)) } return lines.join('\n') } +function formatUsageBlockForBudget(usages: SymbolUsage[], maxChars: number): string { + if (maxChars <= 0) { + return '' + } + + const full = formatUsageBlock(usages) + if (full.length <= maxChars) { + return full + } + + const markerOnly = bestEffortOmissionMarker(usageOmissionMarkerVariants(usages.length), maxChars) + const header = 'usages (import-resolved):' + if (header.length > maxChars) { + return markerOnly + } + + const sections = [header] + for (let index = 0; index < usages.length; index += 1) { + const usage = usages[index]! + const fullLine = formatUsageLineForBudget(usage) + const omittedAfterLine = usages.length - index - 1 + const marker = omittedAfterLine > 0 + ? usageOmissionMarker(omittedAfterLine, remainingCharsForTrailingLine([...sections, fullLine], maxChars)) + : '' + const fullCandidate = joinSections([...sections, fullLine, marker]) + if (fullCandidate.length <= maxChars) { + sections.push(fullLine) + continue + } + + if (omittedAfterLine > 0) { + for (const markerVariant of usageOmissionMarkerVariants(omittedAfterLine)) { + const lineBudget = maxChars + - joinSections(sections).length + - 1 + - markerVariant.length + - 1 + const truncatedLine = formatUsageLineForBudget(usage, lineBudget) + if (!truncatedLine) { + continue + } + + const truncatedCandidate = joinSections([...sections, truncatedLine, markerVariant]) + if (truncatedCandidate.length <= maxChars) { + return truncatedCandidate + } + } + } else { + const truncatedLine = formatUsageLineForBudget(usage, maxChars - joinSections(sections).length - 1) + if (truncatedLine) { + const truncatedCandidate = joinSections([...sections, truncatedLine]) + if (truncatedCandidate.length <= maxChars) { + return truncatedCandidate + } + } + } + + const omitted = usages.length - (sections.length - 1) + for (const fallbackMarker of usageOmissionMarkerVariants(omitted)) { + const fallback = joinSections([...sections, fallbackMarker]) + if (fallback.length <= maxChars) { + return fallback + } + } + + return markerOnly + } + + return joinSections(sections) +} + +function formatUsageLineRaw(usage: SymbolUsage): string { + return `- ${usage.file}:${usage.line} | ${usage.snippet}` +} + +function formatUsageLineForBudget(usage: SymbolUsage, maxChars?: number): string { + const prefix = `- ${usage.file}:${usage.line} | ` + const snippet = usage.snippet.replace(/\n+/g, ' ').trimEnd() + if (maxChars === undefined || prefix.length + snippet.length <= maxChars) { + return `${prefix}${snippet}` + } + if (prefix.length > maxChars) { + return '' + } + + return `${prefix}${truncateWithEllipsis(snippet, maxChars - prefix.length)}` +} + function formatHitSymbol(hit: SearchHit): string { if (!hit.symbol) { return '' @@ -882,6 +1255,14 @@ function joinedLength(rendered: string[], marker: string): number { return rendered.join('\n').length + 1 + marker.length } +function searchHitOmissionMarker(omitted: number, maxChars?: number): string { + return firstMarkerThatFits(searchHitOmissionMarkerVariants(omitted), maxChars) +} + +function usageOmissionMarker(omitted: number, maxChars?: number): string { + return firstMarkerThatFits(usageOmissionMarkerVariants(omitted), maxChars) +} + function grepOmissionMarker(omitted: number, maxChars?: number): string { const full = `matches_omitted=${omitted}; refine path_glob/max_matches or raise max_tokens.` if (maxChars === undefined || full.length <= maxChars) { @@ -901,6 +1282,91 @@ function withGrepOmissionMarker(rendered: string[], marker: string): string { return `${rendered.join('\n')}\n${marker}` } +function searchHitOmissionMarkerVariants(omitted: number): string[] { + return [ + `hits_omitted=${omitted}; refine query/path_glob/k or raise max_tokens.`, + `hits_omitted=${omitted}; raise max_tokens`, + `hits_omitted=${omitted}` + ] +} + +function usageOmissionMarkerVariants(omitted: number): string[] { + return [ + `usages_omitted=${omitted}; raise max_tokens`, + `usages_omitted=${omitted}` + ] +} + +function firstMarkerThatFits(markers: string[], maxChars?: number): string { + if (maxChars === undefined) { + return markers[0] ?? '' + } + + for (const marker of markers) { + if (marker.length <= maxChars) { + return marker + } + } + + return '' +} + +function joinSections(sections: Array): string { + return sections.filter((section): section is string => Boolean(section)).join('\n') +} + +function remainingCharsForTrailingLine(sections: string[], maxChars: number): number { + const output = joinSections(sections) + return Math.max(0, maxChars - output.length - 1) +} + +function appendSearchTokensLineIfFits(output: string, tokensLine: string, maxChars: number): string { + if (!output) { + return tokensLine.length <= maxChars ? tokensLine : truncateWithEllipsis(tokensLine, maxChars) + } + + const withTokens = `${output}\n${tokensLine}` + return withTokens.length <= maxChars ? withTokens : output +} + +function truncateSearchBlock(block: string, maxChars: number): string { + if (maxChars <= 0) { + return '' + } + if (block.length <= maxChars) { + return block + } + + const lines = block.split('\n') + const kept: string[] = [] + let used = 0 + + for (let index = 0; index < lines.length; index += 1) { + const line = lines[index]! + const separator = kept.length === 0 ? 0 : 1 + if (used + separator + line.length <= maxChars) { + kept.push(line) + used += separator + line.length + continue + } + + const lineBudget = maxChars - used - separator + if (lineBudget > 0) { + const truncatedLine = truncateWithEllipsis(line, lineBudget) + if (truncatedLine) { + kept.push(truncatedLine) + } + } + break + } + + if (kept.length === 0) { + return truncateWithEllipsis(block, maxChars) + } + + return kept.join('\n').replace(/\n+$/g, '') +} + function truncateWithEllipsis(text: string, maxChars: number): string { if (maxChars <= 0) { return '' diff --git a/packages/mcp/test/mcp.test.ts b/packages/mcp/test/mcp.test.ts index 3c01f4c..e529e7b 100644 --- a/packages/mcp/test/mcp.test.ts +++ b/packages/mcp/test/mcp.test.ts @@ -9,6 +9,7 @@ import { openRepo, registerEmbeddingProvider, type GrepHit, type SearchHit, type import { DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS, + DEFAULT_MCP_SEARCH_MAX_TOKENS, DEFAULT_MCP_READ_CHUNK_MAX_TOKENS, DEFAULT_SEARCH_K, MIN_MCP_READ_CHUNK_MAX_TOKENS, @@ -61,6 +62,7 @@ describe('@codesift/mcp server', () => { max_tokens: { type: 'integer', minimum: MIN_MCP_READ_CHUNK_MAX_TOKENS, maximum: 4000, default: DEFAULT_MCP_READ_CHUNK_MAX_TOKENS } }) expect(getToolDefinitions().find((tool) => tool.name === 'search_code')?.inputSchema.properties).toMatchObject({ + max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: DEFAULT_MCP_SEARCH_MAX_TOKENS }, context: { type: 'string', enum: ['sig', 'body'] }, with_usages: { type: 'boolean' } }) @@ -172,6 +174,46 @@ describe('@codesift/mcp server', () => { expect(output).toContain('symbols_omitted=') }) + it('budgets search_code output through the MCP call surface, including bundled usages', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-mcp-search-budget-')) + temporaryDirectories.push(repoRoot) + + await mkdir(join(repoRoot, 'src'), { recursive: true }) + await writeFile( + join(repoRoot, 'src', 'demo.ts'), + [ + 'export function demoValue(input: string): string {', + ` const message = \`${'x'.repeat(160)}:${'y'.repeat(160)}\``, + ' return `${input}:${message}`', + '}', + '' + ].join('\n'), + 'utf8' + ) + await Promise.all(([ + ['app.ts', "import { demoValue } from './demo'\nexport const appResult = demoValue('app' + '-'.repeat(80))\n"], + ['worker.ts', "import { demoValue } from './demo'\nexport const workerResult = demoValue('worker' + '-'.repeat(80))\n"], + ['job.ts', "import { demoValue } from './demo'\nexport const jobResult = demoValue('job' + '-'.repeat(80))\n"], + ['cli.ts', "import { demoValue } from './demo'\nexport const cliResult = demoValue('cli' + '-'.repeat(80))\n"] + ] satisfies Array<[string, string]>).map(async ([file, content]) => { + await writeFile(join(repoRoot, 'src', file), content, 'utf8') + })) + + const repo = await openRepo(repoRoot) + await repo.sync() + const output = await callMcpTool(repo, 'search_code', { + query: 'demoValue', + k: 3, + context: 'body', + with_usages: true, + max_tokens: 45 + }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(45) + expect(output).toContain('src/demo.ts') + expect(output).toContain('usages_omitted=') + }) + it('creates server handles', async () => { const repo = await openRepo(process.cwd()) const stdio = createStdioServer(repo) @@ -487,6 +529,131 @@ describe('formatMcpSearchHits structure-preserving output', () => { }) }) +describe('formatMcpSearchHits budgeted output', () => { + it('preserves no_hits and under-budget output byte-for-byte', () => { + expect(formatMcpSearchHits([], { maxTokens: DEFAULT_MCP_SEARCH_MAX_TOKENS })).toBe('no_hits') + + const usageSnippet = " return verify(token)\n // keep trailing spaces " + const hits = [makeHit({ + body: 'export function verify() {\n return true\n}', + usages: [ + { file: 'src/server.ts', range: { startLine: 4, endLine: 5 }, line: 4, snippet: usageSnippet, resolution: 'import-resolved' } + ] + })] + const expected = [ + '~ src/auth.ts:10-14 verify', + '10 | export function verify() {', + '11 | return true', + '12 | }', + 'usages (import-resolved):', + `- src/server.ts:4 | ${usageSnippet}`, + 'tokensReturned=42' + ].join('\n') + + expect(formatMcpSearchHits(hits)).toBe(expected) + expect(formatMcpSearchHits(hits, { maxTokens: 200 })).toBe(expected) + }) + + it('preserves the first hit and appends a hits_omitted marker when later hits are dropped', () => { + const hits = Array.from({ length: 5 }, (_, index) => + makeHit({ + id: `src/file-${index}.ts:${index + 1}-${index + 4}@abcdef0${index}`, + file: `src/file-${index}.ts`, + range: { startLine: index + 1, endLine: index + 4 }, + snippetRange: { startLine: index + 1, endLine: index + 4 }, + body: `export function repeated${index}() {\n return '${'x'.repeat(70)}-${index}'\n}`, + tokensReturned: 30 + index + }) + ) + + const output = formatMcpSearchHits(hits, { maxTokens: 40 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(40) + expect(output).toContain('src/file-0.ts:1-4') + expect(output).not.toContain('src/file-4.ts') + expect(output).toContain('hits_omitted=') + }) + + it('truncates usage snippets without losing file:line locators and reports omitted usages', () => { + const output = formatMcpSearchHits([ + makeHit({ + body: 'export function verify() {\n return true\n}', + usages: Array.from({ length: 5 }, (_, index) => ({ + file: `src/callers/${index}.ts`, + range: { startLine: index + 40, endLine: index + 40 }, + line: index + 40, + snippet: ` return verify(token, '${'x'.repeat(90)}', '${'y'.repeat(90)}')`, + resolution: 'import-resolved' as const + })) + }) + ], { maxTokens: 42 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(42) + expect(output).toContain('usages (import-resolved):') + expect(output).toMatch(/- src\/callers\/0\.ts:40 \| /) + expect(output).toContain('usages_omitted=') + }) + + it('keeps the bare usage omission marker when that is the only variant that fits', () => { + const output = formatMcpSearchHits([ + makeHit({ + body: 'export function verify() {\n const normalized = token.trim()\n return normalized.length > 0\n}', + usages: Array.from({ length: 5 }, (_, index) => ({ + file: `src/usage/${index}.ts`, + range: { startLine: index + 10, endLine: index + 10 }, + line: index + 10, + snippet: ` return verify(token, '${'x'.repeat(80)}')`, + resolution: 'import-resolved' as const + })) + }) + ], { maxTokens: 12 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(12) + expect(output).toContain('usages_omitted=5') + expect(output).not.toContain('raise max_tokens') + }) + + it('keeps a recognizable omission marker for tiny budgets when it fits', () => { + const output = formatMcpSearchHits([ + makeHit({ body: `export function verify() {\n return '${'x'.repeat(120)}'\n}` }), + makeHit({ file: 'src/other.ts', id: 'src/other.ts:1-1@beef0002' }) + ], { maxTokens: 5 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(5) + expect(output).toContain('hits_omitted=1') + }) + + it('falls back to a truncated hit omission marker instead of a misleading bare hit when the budget is microscopic', () => { + const output = formatMcpSearchHits([ + makeHit({ body: `export function verify() {\n return '${'x'.repeat(120)}'\n}` }), + makeHit({ file: 'src/other.ts', id: 'src/other.ts:1-1@beef0002' }) + ], { maxTokens: 3 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(3) + expect(output).toContain('hits_omit') + expect(output).not.toContain('src/auth.ts:10-14 verify') + }) + + it('falls back to a truncated usage omission marker instead of silently dropping usages', () => { + const output = formatMcpSearchHits([ + makeHit({ + body: 'export function verify() {\n const normalized = token.trim()\n return normalized.length > 0\n}', + usages: Array.from({ length: 5 }, (_, index) => ({ + file: `src/usage/${index}.ts`, + range: { startLine: index + 10, endLine: index + 10 }, + line: index + 10, + snippet: ` return verify(token, '${'x'.repeat(80)}')`, + resolution: 'import-resolved' as const + })) + }) + ], { maxTokens: 4 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(4) + expect(output).toContain('usages_omit') + expect(output).not.toContain('usages (import-resolved):') + }) +}) + describe('formatMcpGrepHits keeps its single-line ↩-joined format', () => { it('flattens the snippet with the ↩ arrow and trims each line (unchanged from search-hit rendering)', () => { const hit = makeGrepHit({ snippet: 'function verify() {\n return true\n}' }) From bd37841e5173659d7aeb3a39b91b5cc1b2ba0c0e Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 03:40:42 +0530 Subject: [PATCH 07/13] Compact MCP control-plane metadata --- packages/mcp/src/index.ts | 76 +++++++++++++++++------------------ packages/mcp/test/mcp.test.ts | 42 ++++++++++++++----- 2 files changed, 70 insertions(+), 48 deletions(-) diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 3e3c49c..61711ec 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -155,55 +155,53 @@ export type McpJsonRpcResponse = | { jsonrpc: '2.0'; id: string | number | null; error: { code: number; message: string; data?: unknown } } export const MCP_SERVER_INSTRUCTIONS = [ - 'codesift is the repo search tool. Prefer it before host grep/read-file flows because results are compact and include stable ids for follow-up reads.', - 'Routing policy: use find_symbol for exact identifiers/definitions; use grep_code for literal strings, env vars, error messages, operators, or regex; use search_code for concepts/behaviors/natural language.', - 'search_code returns the complete top result inline (the full enclosing symbol body); no follow-up read is normally needed. Use read_chunk only to expand an ADDITIONAL hit beyond the top result or to widen context.', - 'For search_code, start with k=5-8 and set max_tokens when context is tight. For grep_code, keep context_lines small unless the user asks for surrounding code.', - 'If index_status reports no index, stale data, or a running/failed/aborted sync, suggest running codesift index before relying on results.' + 'Use codesift before host grep/read. Route identifiers/definitions->find_symbol; literals/env/errors/operators/regex->grep_code; concepts/unknown names->search_code.', + 'Top search_code body is inline; read_chunk only for non-top hits/wider context. Broad search_code: k=5-8. Keep grep_code context small.', + 'Check index_status; if missing/stale/running/failed/aborted, warn and suggest codesift index/sync.' ].join('\n') const symbolKindSchema = z.enum(SYMBOL_KINDS) const kindFilterSchema = z.union([symbolKindSchema, z.array(symbolKindSchema)]) const searchCodeInputSchema = { - query: z.string().min(1).describe('Natural-language, behavior, or symbol-aware query. Use grep_code instead for exact literals/regex.'), - k: z.number().int().positive().max(50).optional().describe('Maximum hits to return. Default is 8.'), - lang: z.array(z.string().min(1)).optional().describe('Language filter, e.g. ["typescript"], ["python"].'), - path_glob: z.string().min(1).optional().describe('Repo-relative glob, e.g. "src/**".'), - kind: kindFilterSchema.optional().describe('Symbol kind filter for matching chunks.'), - max_tokens: z.number().int().positive().max(4000).optional().describe('Maximum approximate tokens to return across compact hits. Default 700.'), - single_best: z.boolean().optional().describe('Return only the highest-confidence answer, useful for identifier-exact lookups.'), - context: z.enum(['sig', 'body']).optional().describe('Inline policy: sig=compact signatures/snippets only, body=inline full bodies wherever the budget allows.'), - with_usages: z.boolean().optional().describe('Bundle top-N import-resolved/local usage sites for the top definition hit (TS/JS + Python only).') + query: z.string().min(1).describe('Concept, behavior, or fuzzy name; not exact literals/regex.'), + k: z.number().int().positive().max(50).optional().describe('Max hits. Default 8.'), + lang: z.array(z.string().min(1)).optional().describe('Language filter, e.g. ["typescript"].'), + path_glob: z.string().min(1).optional().describe('Repo glob, e.g. "src/**".'), + kind: kindFilterSchema.optional().describe('Symbol kind filter.'), + max_tokens: z.number().int().positive().max(4000).optional().describe('Approx output tokens. Default 700.'), + single_best: z.boolean().optional().describe('Return only best hit.'), + context: z.enum(['sig', 'body']).optional().describe('sig=snippets; body=inline bodies within budget.'), + with_usages: z.boolean().optional().describe('Add top usage sites for the top definition hit.') } const findSymbolInputSchema = { - name: z.string().min(1).describe('Exact or partial symbol name, e.g. TokenVerifier or verifyJwtToken.'), - kind: kindFilterSchema.optional().describe('Optional symbol kind filter.'), - path_glob: z.string().min(1).optional().describe('Repo-relative glob, e.g. "src/auth/**".'), - with_body: z.boolean().optional().describe('Inline the top exact match\'s full definition body so it resolves in one call. Default true; set false for a compact name→location list.'), - max_tokens: z.number().int().positive().max(4000).optional().describe('Maximum approximate tokens to return across symbol definitions. Default 700.') + name: z.string().min(1).describe('Exact/partial symbol name.'), + kind: kindFilterSchema.optional().describe('Symbol kind filter.'), + path_glob: z.string().min(1).optional().describe('Repo glob, e.g. "src/auth/**".'), + with_body: z.boolean().optional().describe('Inline top exact body. Default true.'), + max_tokens: z.number().int().positive().max(4000).optional().describe('Approx output tokens. Default 700.') } const grepCodeInputSchema = { - pattern: z.string().min(1).describe('Literal byte/string by default. Set regex=true for regular expressions.'), - regex: z.boolean().optional().describe('Treat pattern as a JavaScript regular expression. Default false for byte-exact literal search.'), - ignore_case: z.boolean().optional().describe('Case-insensitive matching, like rg -i.'), - whole_word: z.boolean().optional().describe('Match only whole identifier/word occurrences, like rg -w.'), - multiline: z.boolean().optional().describe('Allow regex dot to span newlines and report multi-line matches.'), + pattern: z.string().min(1).describe('Literal by default; set regex=true for regex.'), + regex: z.boolean().optional().describe('Use JavaScript regex. Default false.'), + ignore_case: z.boolean().optional().describe('Case-insensitive match.'), + whole_word: z.boolean().optional().describe('Whole word/identifier match.'), + multiline: z.boolean().optional().describe('Allow multi-line regex matches.'), lang: z.array(z.string().min(1)).optional().describe('Language filter, e.g. ["typescript"].'), - path_glob: z.string().min(1).optional().describe('Repo-relative glob, e.g. "packages/core/**".'), - context_lines: z.number().int().min(0).max(20).optional().describe('Symmetric context lines, like rg -C.'), - before_context_lines: z.number().int().min(0).max(20).optional().describe('Lines before each match, like rg -B.'), - after_context_lines: z.number().int().min(0).max(20).optional().describe('Lines after each match, like rg -A.'), - max_matches: z.number().int().positive().max(1000).optional().describe('Maximum matches to return. Default 1000.'), - max_tokens: z.number().int().positive().max(4000).optional().describe('Maximum approximate tokens to return across compact matches. Default 700.') + path_glob: z.string().min(1).optional().describe('Repo glob, e.g. "packages/core/**".'), + context_lines: z.number().int().min(0).max(20).optional().describe('Lines before/after each match.'), + before_context_lines: z.number().int().min(0).max(20).optional().describe('Lines before each match.'), + after_context_lines: z.number().int().min(0).max(20).optional().describe('Lines after each match.'), + max_matches: z.number().int().positive().max(1000).optional().describe('Max matches. Default 1000.'), + max_tokens: z.number().int().positive().max(4000).optional().describe('Approx output tokens. Default 700.') } const readChunkInputSchema = { - id: z.string().min(1).describe('Stable chunk id returned by search_code.'), - context_lines: z.number().int().min(0).max(50).optional().describe('Extra lines before/after the chunk.'), - max_tokens: z.number().int().min(MIN_MCP_READ_CHUNK_MAX_TOKENS).max(MAX_MCP_READ_CHUNK_MAX_TOKENS).optional().describe('Maximum approximate tokens to return from this chunk. Default 1000.') + id: z.string().min(1).describe('Stable search_code hit id.'), + context_lines: z.number().int().min(0).max(50).optional().describe('Extra surrounding lines.'), + max_tokens: z.number().int().min(MIN_MCP_READ_CHUNK_MAX_TOKENS).max(MAX_MCP_READ_CHUNK_MAX_TOKENS).optional().describe('Approx output tokens. Default 1000.') } const indexStatusInputSchema = {} @@ -245,8 +243,8 @@ class StdioMcpServerHandle implements McpServerHandle { export function getToolDefinitions(): readonly McpToolDefinition[] { const provider = getDefaultEmbeddingProvider() const searchDescription = isLearnedEmbeddingProvider(provider) - ? 'Concept/behavior search over the current repo using hybrid lexical + semantic retrieval. Returns the complete top result inline (full enclosing symbol body); no follow-up read is normally needed. Use for natural-language questions; prefer grep_code for exact literals/regex and find_symbol for definitions.' - : 'Concept/behavior search over the current repo using lexical retrieval. Returns the complete top result inline (full enclosing symbol body); no follow-up read is normally needed. Use for natural-language questions; prefer grep_code for exact literals/regex and find_symbol for definitions.' + ? 'Hybrid lexical + semantic repo search for concepts/unknown names. Top body inline; read_chunk for other hits. Use grep_code for literals/regex, find_symbol for definitions.' + : 'Lexical repo search for concepts/unknown names. Top body inline; read_chunk for other hits. Use grep_code for literals/regex, find_symbol for definitions.' return [ { @@ -266,7 +264,7 @@ export function getToolDefinitions(): readonly McpToolDefinition[] { }, { name: 'find_symbol', - description: 'Exact identifier/definition lookup from the symbols table. Returns the top match\'s full definition body inline (no follow-up read normally needed). Use before search_code for class/function/type names.', + description: 'Exact identifier/definition lookup. Top unambiguous body inline; use before search_code for names.', inputSchema: jsonSchema(['name'], { name: { type: 'string' }, kind: kindJsonSchema(), @@ -277,7 +275,7 @@ export function getToolDefinitions(): readonly McpToolDefinition[] { }, { name: 'grep_code', - description: 'Literal, byte-exact, or regex search over indexed repo files. Use instead of host grep for env vars, error strings, operators, and regex.', + description: 'Literal/regex indexed-file search for env vars, errors, operators, exact strings; keep context small.', inputSchema: jsonSchema(['pattern'], { pattern: { type: 'string' }, regex: { type: 'boolean', default: false }, @@ -295,7 +293,7 @@ export function getToolDefinitions(): readonly McpToolDefinition[] { }, { name: 'read_chunk', - description: 'Expand an ADDITIONAL search_code hit id into its source chunk, or widen context around one. The top search_code result is already returned inline, so this is rarely needed for the best hit; use it for secondary hits or extra surrounding lines, not as a first step.', + description: 'Read non-top hit/wider context by id. Not needed for top search_code/find_symbol hits returned inline.', inputSchema: jsonSchema(['id'], { id: { type: 'string' }, context_lines: { type: 'integer', minimum: 0, maximum: 50 }, @@ -304,7 +302,7 @@ export function getToolDefinitions(): readonly McpToolDefinition[] { }, { name: 'index_status', - description: 'Inspect index freshness, sync/crash state, counts, provider, and vector availability before relying on search results.', + description: 'Inspect index freshness, sync state, counts, provider, vectors.', inputSchema: jsonSchema([], {}) } ] diff --git a/packages/mcp/test/mcp.test.ts b/packages/mcp/test/mcp.test.ts index e529e7b..7db3a73 100644 --- a/packages/mcp/test/mcp.test.ts +++ b/packages/mcp/test/mcp.test.ts @@ -66,25 +66,49 @@ describe('@codesift/mcp server', () => { context: { type: 'string', enum: ['sig', 'body'] }, with_usages: { type: 'boolean' } }) - expect(MCP_SERVER_INSTRUCTIONS).toContain('use grep_code for literal strings') + expect(MCP_SERVER_INSTRUCTIONS).toContain('literals/env/errors/operators/regex->grep_code') }) it('asserts single-call sufficiency in instructions and tool descriptions', () => { - expect(MCP_SERVER_INSTRUCTIONS).toContain('search_code returns the complete top result inline') - expect(MCP_SERVER_INSTRUCTIONS).toContain('no follow-up read is normally needed') - expect(MCP_SERVER_INSTRUCTIONS.toLowerCase()).toContain('read_chunk only to expand an additional hit') + const lowerInstructions = MCP_SERVER_INSTRUCTIONS.toLowerCase() + expect(lowerInstructions).toContain('top search_code body is inline') + expect(lowerInstructions).toContain('read_chunk only for non-top hits/wider context') const tools = getToolDefinitions() const searchDescription = tools.find((tool) => tool.name === 'search_code')?.description ?? '' - expect(searchDescription).toContain('complete top result inline') - expect(searchDescription.toLowerCase()).toContain('no follow-up read is normally needed') + expect(searchDescription.toLowerCase()).toContain('top body inline') + + const findDescription = tools.find((tool) => tool.name === 'find_symbol')?.description ?? '' + expect(findDescription.toLowerCase()).toContain('top unambiguous body inline') const readDescription = tools.find((tool) => tool.name === 'read_chunk')?.description ?? '' - expect(readDescription).toContain('ADDITIONAL') - expect(readDescription.toLowerCase()).toContain('already returned inline') + expect(readDescription.toLowerCase()).toContain('not needed for top search_code/find_symbol hits') + expect(readDescription.toLowerCase()).toContain('returned inline') // The routing guidance for find_symbol/grep_code is preserved. - expect(MCP_SERVER_INSTRUCTIONS).toContain('use find_symbol for exact identifiers/definitions') + expect(MCP_SERVER_INSTRUCTIONS).toContain('identifiers/definitions->find_symbol') + expect(MCP_SERVER_INSTRUCTIONS).toContain('concepts/unknown names->search_code') + expect(MCP_SERVER_INSTRUCTIONS).toContain('Broad search_code: k=5-8') + expect(MCP_SERVER_INSTRUCTIONS).toContain('Check index_status') + }) + + it('keeps control-plane metadata under the MCP budget', () => { + const tools = getToolDefinitions() + const payload = JSON.stringify({ instructions: MCP_SERVER_INSTRUCTIONS, tools }) + const toolDescriptionCeilings = new Map([ + ['search_code', 190], + ['find_symbol', 130], + ['grep_code', 120], + ['read_chunk', 120], + ['index_status', 90] + ]) + + expect(payload.length).toBeLessThanOrEqual(3500) + expect(Math.ceil(payload.length / 4)).toBeLessThanOrEqual(875) + expect(MCP_SERVER_INSTRUCTIONS.length).toBeLessThanOrEqual(430) + for (const tool of tools) { + expect(tool.description.length).toBeLessThanOrEqual(toolDescriptionCeilings.get(tool.name) ?? 120) + } }) it('uses honest lexical wording by default and semantic wording with a learned provider', () => { From 6255ac790ea56f6fd84f1e797231256d25ee7eb9 Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 03:55:02 +0530 Subject: [PATCH 08/13] Budget MCP index_status output --- packages/mcp/src/index.ts | 241 +++++++++++++++++++++++++++++++++- packages/mcp/test/mcp.test.ts | 146 +++++++++++++++++++- 2 files changed, 379 insertions(+), 8 deletions(-) diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 61711ec..6f8ee13 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -27,8 +27,10 @@ export const DEFAULT_MCP_SEARCH_MAX_TOKENS = 700 export const DEFAULT_MCP_GREP_MAX_TOKENS = 700 export const DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS = 700 export const DEFAULT_MCP_READ_CHUNK_MAX_TOKENS = 1000 +export const DEFAULT_MCP_INDEX_STATUS_MAX_TOKENS = 200 export const MIN_MCP_READ_CHUNK_MAX_TOKENS = 6 export const MAX_MCP_READ_CHUNK_MAX_TOKENS = 4000 +export const MAX_MCP_INDEX_STATUS_MAX_TOKENS = 1000 const SYMBOL_KINDS = [ 'class', @@ -111,6 +113,14 @@ export interface ReadChunkArgs { max_tokens?: number | undefined } +export interface IndexStatusArgs { + max_tokens?: number | undefined +} + +export interface FormatMcpIndexStatusOptions { + maxTokens?: number | undefined +} + export interface McpToolDefinition { name: McpToolName description: string @@ -204,12 +214,15 @@ const readChunkInputSchema = { max_tokens: z.number().int().min(MIN_MCP_READ_CHUNK_MAX_TOKENS).max(MAX_MCP_READ_CHUNK_MAX_TOKENS).optional().describe('Approx output tokens. Default 1000.') } -const indexStatusInputSchema = {} +const indexStatusInputSchema = { + max_tokens: z.number().int().positive().max(MAX_MCP_INDEX_STATUS_MAX_TOKENS).optional().describe('Approx output tokens. Default 200.') +} const searchCodeArgsSchema = z.object(searchCodeInputSchema).strict() const findSymbolArgsSchema = z.object(findSymbolInputSchema).strict() const grepCodeArgsSchema = z.object(grepCodeInputSchema).strict() const readChunkArgsSchema = z.object(readChunkInputSchema).strict() +const indexStatusArgsSchema = z.object(indexStatusInputSchema).strict() const toolCallParamsSchema = z.object({ name: z.enum(MCP_TOOL_NAMES), arguments: z.unknown().optional() @@ -275,7 +288,7 @@ export function getToolDefinitions(): readonly McpToolDefinition[] { }, { name: 'grep_code', - description: 'Literal/regex indexed-file search for env vars, errors, operators, exact strings; keep context small.', + description: 'Literal/regex search for env vars, errors, operators, exact strings.', inputSchema: jsonSchema(['pattern'], { pattern: { type: 'string' }, regex: { type: 'boolean', default: false }, @@ -302,8 +315,10 @@ export function getToolDefinitions(): readonly McpToolDefinition[] { }, { name: 'index_status', - description: 'Inspect index freshness, sync state, counts, provider, vectors.', - inputSchema: jsonSchema([], {}) + description: 'Inspect index health.', + inputSchema: jsonSchema([], { + max_tokens: { type: 'integer', minimum: 1, maximum: MAX_MCP_INDEX_STATUS_MAX_TOKENS, default: DEFAULT_MCP_INDEX_STATUS_MAX_TOKENS } + }) } ] } @@ -435,7 +450,10 @@ export async function callMcpTool(repo: Repo, name: McpToolName, args: unknown): return formatMcpReadChunk(await router.readChunk(parsed), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_READ_CHUNK_MAX_TOKENS }) } case 'index_status': - return JSON.stringify(await router.indexStatus()) + { + const parsed = indexStatusArgsSchema.parse(args) + return formatMcpIndexStatus(await router.indexStatus(), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_INDEX_STATUS_MAX_TOKENS }) + } } } @@ -512,8 +530,8 @@ export function createSdkServer(repo: Repo): McpServer { server.registerTool('read_chunk', { description: toolDescription('read_chunk'), inputSchema: readChunkInputSchema }, async (args) => textResult(formatMcpReadChunk(await router.readChunk(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_READ_CHUNK_MAX_TOKENS })) ) - server.registerTool('index_status', { description: toolDescription('index_status'), inputSchema: indexStatusInputSchema }, async () => - textResult(JSON.stringify(await router.indexStatus())) + server.registerTool('index_status', { description: toolDescription('index_status'), inputSchema: indexStatusInputSchema }, async (args) => + textResult(formatMcpIndexStatus(await router.indexStatus(), { maxTokens: args.max_tokens ?? DEFAULT_MCP_INDEX_STATUS_MAX_TOKENS })) ) return server @@ -970,6 +988,215 @@ function formatUsageLineForBudget(usage: SymbolUsage, maxChars?: number): string return `${prefix}${truncateWithEllipsis(snippet, maxChars - prefix.length)}` } +export function formatMcpIndexStatus(status: RepoStatus, options: FormatMcpIndexStatusOptions = {}): string { + const lines = [ + formatIndexStatusPrimaryLine(status), + ...formatIndexStatusDetailLines(status) + ] + const action = suggestedIndexStatusAction(status) + if (action) { + lines.push(`action=${action}`) + } + + const output = lines.join('\n') + const maxTokens = options.maxTokens + if (maxTokens === undefined || output.length <= maxTokens * 4) { + return output + } + + return fitIndexStatusOutputForBudget(lines[0]!, action ? `action=${action}` : undefined, lines.slice(1, action ? -1 : undefined), maxTokens * 4) +} + +function formatIndexStatusPrimaryLine(status: RepoStatus): string { + return [ + `indexed=${status.indexed ? 'yes' : 'no'}`, + `stale=${status.stale ? 'yes' : 'no'}`, + `sync=${status.sync.state}`, + `chunks=${status.chunkCount}`, + `symbols=${status.symbolCount}`, + `gen=${status.indexGeneration}`, + `generated=${status.generatedFileCount}/${status.generatedChunkCount}`, + formatIndexStatusProvider(status), + `compat=${status.compatibility.ok ? 'ok' : status.compatibility.code ?? 'mismatch'}`, + `vector=${status.vectorSearch.state}` + ].filter(Boolean).join(' ') +} + +function formatIndexStatusProvider(status: RepoStatus): string { + if (!status.provider) { + return 'provider=unconfigured' + } + + const parts = [`provider=${compactStatusValue(status.provider.id)}`] + if (status.provider.dims !== undefined) { + parts.push(`dims=${status.provider.dims}`) + } + if (status.provider.modelVersion) { + parts.push(`model=${compactStatusValue(status.provider.modelVersion, 48)}`) + } + return parts.join(' ') +} + +function formatIndexStatusDetailLines(status: RepoStatus): string[] { + const lines: string[] = [] + + for (const reason of status.staleReasons ?? []) { + const parts = [`stale_reasons=${reason.code}`] + if (reason.count !== undefined) { + parts.push(`count=${reason.count}`) + } + if (reason.files?.length) { + parts.push(`files=${formatStatusFileList(reason.files)}`) + } + if (reason.message) { + parts.push(`message=${compactStatusValue(reason.message, 100)}`) + } + if (reason.indexed) { + parts.push(`indexed=${compactStatusValue(reason.indexed, 48)}`) + } + if (reason.current) { + parts.push(`current=${compactStatusValue(reason.current, 48)}`) + } + lines.push(parts.join(' ')) + } + + if (status.sync.error) { + lines.push(`sync_error=${compactStatusValue(status.sync.error, 140)}`) + } + + if (!status.compatibility.ok) { + if (status.compatibility.message) { + lines.push(`compat_message=${compactStatusValue(status.compatibility.message, 160)}`) + } else if (status.compatibility.code) { + lines.push(`compat_code=${status.compatibility.code}`) + } + } + + if (status.vectorSearch.reason) { + lines.push(`vector_reason=${status.vectorSearch.reason}`) + } + if (status.vectorSearch.state === 'unavailable' && status.vectorSearch.message) { + lines.push(`vector_message=${compactStatusValue(status.vectorSearch.message, 120)}`) + } + + return lines +} + +function suggestedIndexStatusAction(status: RepoStatus): string { + if (status.indexed && !status.compatibility.ok) { + return 'codesift index --rebuild' + } + if (!status.indexed) { + return 'codesift index' + } + if (status.sync.state === 'failed' || status.sync.state === 'aborted') { + return 'codesift sync' + } + if (status.stale) { + return 'codesift sync' + } + if (status.sync.state === 'running') { + return 'wait for sync' + } + + return '' +} + +function fitIndexStatusOutputForBudget(primaryLine: string, actionLine: string | undefined, detailLines: string[], maxChars: number): string { + const marker = indexStatusTruncationMarker(Math.max(0, maxChars - 1)) + const prioritized = actionLine ? [primaryLine, actionLine] : [primaryLine] + const output = prioritized.join('\n') + + if (!marker) { + return truncateWithEllipsis(output, maxChars) + } + + const fitted = [...prioritized] + for (const detail of detailLines) { + const candidate = [...fitted, detail, marker].join('\n') + if (candidate.length <= maxChars) { + fitted.push(detail) + continue + } + + const remainingForDetail = maxChars - fitted.join('\n').length - 1 - marker.length - 1 + if (remainingForDetail > 0) { + const truncatedDetail = truncateWithEllipsis(detail, remainingForDetail) + const truncatedCandidate = [...fitted, truncatedDetail, marker].join('\n') + if (truncatedCandidate.length <= maxChars) { + fitted.push(truncatedDetail) + } + } + break + } + + const fittedWithMarker = [...fitted, marker].join('\n') + if (fittedWithMarker.length <= maxChars) { + return fittedWithMarker + } + + const primaryWithMarker = [primaryLine, marker].join('\n') + if (primaryWithMarker.length <= maxChars) { + return primaryWithMarker + } + + let compactPrimaryOnlyFallback = '' + for (const compactPrimary of compactIndexStatusPrimaryLineVariants(primaryLine)) { + if (compactPrimary === primaryLine) { + continue + } + const compactPrioritized = actionLine ? [compactPrimary, actionLine] : [compactPrimary] + const compactWithMarker = [...compactPrioritized, marker].join('\n') + if (compactWithMarker.length <= maxChars) { + return compactWithMarker + } + + const compactPrimaryWithMarker = [compactPrimary, marker].join('\n') + if (compactPrimaryWithMarker.length <= maxChars) { + compactPrimaryOnlyFallback ||= compactPrimaryWithMarker + } + } + + if (compactPrimaryOnlyFallback) { + return compactPrimaryOnlyFallback + } + + return truncateWithEllipsis(primaryLine, maxChars) +} + +function compactIndexStatusPrimaryLineVariants(primaryLine: string): string[] { + const fields = primaryLine.split(' ') + return [ + ['indexed=', 'stale=', 'sync=', 'compat=', 'vector='], + ['indexed=', 'stale=', 'sync='], + ['indexed=', 'sync='] + ].map((prefixes) => fields.filter((field) => prefixes.some((prefix) => field.startsWith(prefix))).join(' ')) +} + +function indexStatusTruncationMarker(maxChars?: number): string { + const full = 'status_truncated=true; raise max_tokens' + const bare = 'status_truncated=true' + if (maxChars !== undefined && bare.length <= maxChars) { + return bare + } + if (maxChars === undefined || full.length <= maxChars) { + return full + } + + return bare.length <= maxChars ? bare : '' +} + +function formatStatusFileList(files: string[]): string { + const kept = files.slice(0, 2) + const omitted = files.length - kept.length + return [...kept.map((file) => compactStatusValue(file, 60)), ...(omitted > 0 ? [`+${omitted}`] : [])].join(',') +} + +function compactStatusValue(value: string, maxChars = 80): string { + const compact = value.replace(/\s+/g, ' ').trim() + return compact.length <= maxChars ? compact : truncateWithEllipsis(compact, maxChars) +} + function formatHitSymbol(hit: SearchHit): string { if (!hit.symbol) { return '' diff --git a/packages/mcp/test/mcp.test.ts b/packages/mcp/test/mcp.test.ts index 7db3a73..76756f8 100644 --- a/packages/mcp/test/mcp.test.ts +++ b/packages/mcp/test/mcp.test.ts @@ -5,10 +5,11 @@ import { join } from 'node:path' import { afterEach, describe, expect, it } from 'vitest' -import { openRepo, registerEmbeddingProvider, type GrepHit, type SearchHit, type SymbolDefinition } from '@codesift/core' +import { openRepo, registerEmbeddingProvider, type GrepHit, type RepoStatus, type SearchHit, type SymbolDefinition } from '@codesift/core' import { DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS, + DEFAULT_MCP_INDEX_STATUS_MAX_TOKENS, DEFAULT_MCP_SEARCH_MAX_TOKENS, DEFAULT_MCP_READ_CHUNK_MAX_TOKENS, DEFAULT_SEARCH_K, @@ -19,6 +20,7 @@ import { createStdioServer, callMcpTool, formatMcpGrepHits, + formatMcpIndexStatus, formatMcpReadChunk, formatMcpSearchHits, formatMcpSymbols, @@ -66,6 +68,9 @@ describe('@codesift/mcp server', () => { context: { type: 'string', enum: ['sig', 'body'] }, with_usages: { type: 'boolean' } }) + expect(getToolDefinitions().find((tool) => tool.name === 'index_status')?.inputSchema.properties).toMatchObject({ + max_tokens: { type: 'integer', minimum: 1, maximum: 1000, default: DEFAULT_MCP_INDEX_STATUS_MAX_TOKENS } + }) expect(MCP_SERVER_INSTRUCTIONS).toContain('literals/env/errors/operators/regex->grep_code') }) @@ -386,6 +391,145 @@ function makeSymbol(overrides: Partial> & { body? return definition } +function makeStatus(overrides: Partial = {}): RepoStatus { + return { + root: '/tmp/repo', + indexPath: '/tmp/repo/.codesift/index.db', + indexed: true, + stale: false, + sync: { state: 'completed', completedAt: '2026-06-20T00:00:00.000Z' }, + chunkCount: 123, + symbolCount: 45, + generatedFileCount: 0, + generatedChunkCount: 0, + indexGeneration: 3, + provider: { id: 'local-hash' }, + compatibility: { ok: true }, + vectorSearch: { available: true, state: 'lazy' }, + ...overrides + } +} + +describe('formatMcpIndexStatus compact output', () => { + it('renders healthy status as compact lines without raw paths or JSON wrappers', () => { + const output = formatMcpIndexStatus(makeStatus()) + + expect(output).toBe('indexed=yes stale=no sync=completed chunks=123 symbols=45 gen=3 generated=0/0 provider=local-hash compat=ok vector=lazy') + expect(output).not.toContain('/tmp/repo') + expect(output).not.toContain('indexPath') + expect(output).not.toContain('{') + }) + + it('reports a missing index with an indexing action', () => { + const output = formatMcpIndexStatus(makeStatus({ + indexed: false, + sync: { state: 'idle' }, + chunkCount: 0, + symbolCount: 0, + indexGeneration: 0, + provider: null + })) + + expect(output).toContain('indexed=no') + expect(output).toContain('provider=unconfigured') + expect(output).toContain('action=codesift index') + }) + + it('reports stale reason codes, clipped files, and sync action', () => { + const output = formatMcpIndexStatus(makeStatus({ + stale: true, + staleReasons: [ + { code: 'file_modified', message: '4 files modified', count: 4, files: ['src/a.ts', 'src/b.ts', 'src/c.ts', 'src/d.ts'] } + ] + })) + + expect(output).toContain('stale=yes') + expect(output).toContain('stale_reasons=file_modified count=4 files=src/a.ts,src/b.ts,+2') + expect(output).toContain('message=4 files modified') + expect(output).toContain('action=codesift sync') + }) + + it('keeps sync errors and a truncation marker under small budgets when possible', () => { + const output = formatMcpIndexStatus(makeStatus({ + sync: { state: 'failed', error: `planned failure ${'x'.repeat(300)}` } + }), { maxTokens: 45 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(45) + expect(output).toContain('sync=failed') + expect(output).toContain('action=codesift sync') + expect(output).toContain('sync_error=') + expect(output).toContain('status_truncated=true') + }) + + it('keeps action visible by compacting the primary line for tight actionable budgets', () => { + const output = formatMcpIndexStatus(makeStatus({ + provider: { id: `provider-${'x'.repeat(80)}`, dims: 1536, modelVersion: `model-${'y'.repeat(80)}` }, + compatibility: { + ok: false, + code: 'model_version_mismatch', + message: `provider model changed ${'z'.repeat(200)}` + } + }), { maxTokens: 32 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(32) + expect(output).toContain('indexed=yes') + expect(output).toContain('action=codesift index --rebuild') + expect(output).toContain('status_truncated=true') + }) + + it('reports compatibility mismatch without raw expected or actual snapshots', () => { + const output = formatMcpIndexStatus(makeStatus({ + compatibility: { + ok: false, + code: 'provider_mismatch', + message: 'Provider changed; run codesift index --rebuild', + expected: { providerId: 'local-hash' }, + actual: { providerId: 'cloud-provider' } + } + })) + + expect(output).toContain('compat=provider_mismatch') + expect(output).toContain('compat_message=Provider changed; run codesift index --rebuild') + expect(output).toContain('action=codesift index --rebuild') + expect(output).not.toContain('expected') + expect(output).not.toContain('actual') + }) + + it('reports vector unavailability reason without dumping long detail', () => { + const output = formatMcpIndexStatus(makeStatus({ + vectorSearch: { + available: false, + state: 'unavailable', + reason: 'native-dependency-unavailable', + message: 'vector search unavailable (native dep), lexical/symbol still works', + detail: `missing sqlite-vec prebuild ${'x'.repeat(200)}` + } + })) + + expect(output).toContain('vector=unavailable') + expect(output).toContain('vector_reason=native-dependency-unavailable') + expect(output).toContain('vector_message=vector search unavailable') + expect(output).not.toContain('sqlite-vec prebuild') + }) + + it('budgets index_status output through callMcpTool', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-mcp-status-budget-')) + temporaryDirectories.push(repoRoot) + + await mkdir(join(repoRoot, 'src'), { recursive: true }) + await writeFile(join(repoRoot, 'src', 'demo.ts'), `export const demoValue = '${'x'.repeat(200)}'\n`, 'utf8') + + const repo = await openRepo(repoRoot) + await repo.sync() + + const output = await callMcpTool(repo, 'index_status', { max_tokens: 40 }) + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(40) + expect(output).toContain('indexed=yes') + expect(output).not.toContain('indexPath') + expect(output).not.toContain('{') + }) +}) + describe('formatMcpSymbols budgeted output', () => { it('preserves no_symbols and under-budget output byte-for-byte', () => { const definitions = [makeSymbol()] From 7440eef08648f5d9c31af5e24c8df58d29490d99 Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 04:11:57 +0530 Subject: [PATCH 09/13] Merge overlapping grep context windows --- packages/core/src/repo.ts | 3 +- packages/core/src/types.ts | 1 + packages/core/test/grep.parity.test.ts | 28 ++++ packages/mcp/src/index.ts | 199 ++++++++++++++++++++++--- packages/mcp/test/mcp.test.ts | 147 ++++++++++++++++++ 5 files changed, 357 insertions(+), 21 deletions(-) diff --git a/packages/core/src/repo.ts b/packages/core/src/repo.ts index a26a120..cd327b6 100644 --- a/packages/core/src/repo.ts +++ b/packages/core/src/repo.ts @@ -623,7 +623,8 @@ export class SqliteRepo implements Repo { line: startLine, column: startOffset - lineStarts[startLine - 1]! + 1, match: matchedText, - snippet: lines.slice(snippetStartLine - 1, snippetEndLine).join('\n') + snippet: lines.slice(snippetStartLine - 1, snippetEndLine).join('\n'), + snippetRange: { startLine: snippetStartLine, endLine: snippetEndLine } } if (file.language) { diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index 4ff7a69..c9e2ace 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -78,6 +78,7 @@ export interface GrepHit { column: number match: string snippet: string + snippetRange?: Range language?: string } diff --git a/packages/core/test/grep.parity.test.ts b/packages/core/test/grep.parity.test.ts index 85b6682..a27e8d4 100644 --- a/packages/core/test/grep.parity.test.ts +++ b/packages/core/test/grep.parity.test.ts @@ -61,6 +61,34 @@ describe('@codesift/core grep parity', () => { expect(codesiftHits).toEqual(ripgrepHits) } }) + + it('returns the snippet range used to build contextual grep snippets', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-grep-snippet-range-')) + temporaryDirectories.push(repoRoot) + + await mkdir(join(repoRoot, 'src'), { recursive: true }) + await writeFile( + join(repoRoot, 'src', 'demo.ts'), + [ + 'const before = true', + "const value = 'NEEDLE'", + 'const after = true' + ].join('\n'), + 'utf8' + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + + await expect(repo.grep('NEEDLE', { pathGlob: 'src/**', contextLines: 1 })).resolves.toMatchObject([ + { + file: 'src/demo.ts', + range: { startLine: 2, endLine: 2 }, + snippet: "const before = true\nconst value = 'NEEDLE'\nconst after = true", + snippetRange: { startLine: 1, endLine: 3 } + } + ]) + }) }) async function runRipgrep(repoRoot: string, pattern: string, options?: GrepOptions): Promise> { diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 6f8ee13..1d420d4 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -1323,36 +1323,34 @@ export function formatMcpGrepHits(hits: GrepHit[], options: FormatMcpGrepHitsOpt return 'no_matches' } + const entries = buildMcpGrepEntries(hits) const maxTokens = options.maxTokens if (maxTokens === undefined) { - return hits.map((hit) => formatMcpGrepHit(hit)).join('\n') + return entries.map((entry) => entry.text).join('\n') } const maxChars = maxTokens * 4 - const rendered: string[] = [] + const rendered: McpRenderedGrepEntry[] = [] let usedChars = 0 - for (let index = 0; index < hits.length; index += 1) { - const hit = hits[index]! - const line = formatMcpGrepHit(hit) + for (let index = 0; index < entries.length; index += 1) { + const entry = entries[index]! + const line = entry.text const separatorChars = rendered.length === 0 ? 0 : 1 if (usedChars + separatorChars + line.length <= maxChars) { - rendered.push(line) + rendered.push({ text: line, hits: entry.hits, matchCount: entry.hits.length }) usedChars += separatorChars + line.length continue } if (rendered.length === 0) { - const omitted = hits.length - 1 - const marker = omitted > 0 ? grepOmissionMarker(omitted, Math.max(0, maxChars - 2)) : undefined - const firstHit = formatMcpGrepHit(hit, maxCharsForFirstHit(maxChars, marker)) - return marker ? withGrepOmissionMarker([firstHit], marker) : firstHit + return formatMcpGrepEntryForBudget(entry, hits.length - entry.hits.length, maxChars) } - return fitGrepOutputWithMarker(rendered, hits.length - index, maxChars) + return fitGrepOutputWithMarker(rendered, countGrepEntryMatches(entries, index), maxChars) } - return rendered.join('\n') + return rendered.map((entry) => entry.text).join('\n') } export function formatMcpReadChunk(content: string, options: FormatMcpReadChunkOptions = {}): string { @@ -1420,6 +1418,162 @@ function formatMcpGrepHit(hit: GrepHit, maxChars?: number): string { return `${prefix}${truncateWithEllipsis(snippet, maxChars - prefix.length)}` } +interface McpGrepEntry { + hits: GrepHit[] + text: string +} + +interface McpRenderedGrepEntry { + text: string + hits: GrepHit[] + matchCount: number +} + +function buildMcpGrepEntries(hits: GrepHit[]): McpGrepEntry[] { + const entries: McpGrepEntry[] = [] + + for (let index = 0; index < hits.length;) { + const group = collectOverlappingGrepHits(hits, index) + if (group.length > 1) { + entries.push({ hits: group, text: formatMcpGrepGroup(group) }) + index += group.length + continue + } + + const hit = hits[index]! + entries.push({ hits: [hit], text: formatMcpGrepHit(hit) }) + index += 1 + } + + return entries +} + +function collectOverlappingGrepHits(hits: GrepHit[], startIndex: number): GrepHit[] { + const first = hits[startIndex]! + if (!isGroupableGrepHit(first)) { + return [first] + } + + const group = [first] + let mergedEndLine = first.snippetRange.endLine + + for (let index = startIndex + 1; index < hits.length; index += 1) { + const hit = hits[index]! + if (!isGroupableGrepHit(hit) || hit.file !== first.file || hit.snippetRange.startLine > mergedEndLine + 1) { + break + } + + group.push(hit) + mergedEndLine = Math.max(mergedEndLine, hit.snippetRange.endLine) + } + + return group +} + +function isGroupableGrepHit(hit: GrepHit): hit is GrepHit & { snippetRange: NonNullable } { + return hit.snippetRange !== undefined && + (hit.snippetRange.startLine < hit.range.startLine || hit.snippetRange.endLine > hit.range.endLine) +} + +function formatMcpGrepGroup(hits: GrepHit[]): string { + return [ + ...hits.map((hit) => formatMcpGrepLocator(hit)), + formatMcpGrepMergedSnippet(hits) + ].join('\n') +} + +function formatMcpGrepLocator(hit: GrepHit): string { + return `${hit.file}:${formatRange(hit.range.startLine, hit.range.endLine)}:${hit.column}` +} + +function formatMcpGrepMergedSnippet(hits: GrepHit[]): string { + const sourceLines = new Map() + + for (const hit of hits) { + if (!hit.snippetRange) { + continue + } + + const lines = hit.snippet.split('\n') + for (let offset = 0; offset < lines.length; offset += 1) { + const lineNumber = hit.snippetRange.startLine + offset + if (lineNumber <= hit.snippetRange.endLine && !sourceLines.has(lineNumber)) { + sourceLines.set(lineNumber, lines[offset]!) + } + } + } + + return [...sourceLines.entries()] + .sort(([left], [right]) => left - right) + .map(([lineNumber, line]) => `${lineNumber} | ${line}`) + .join('\n') +} + +function countGrepEntryMatches(entries: McpGrepEntry[], startIndex: number): number { + let count = 0 + for (let index = startIndex; index < entries.length; index += 1) { + count += entries[index]!.hits.length + } + return count +} + +function formatMcpGrepEntryForBudget(entry: McpGrepEntry, omittedAfter: number, maxChars: number): string { + if (entry.hits.length === 1) { + const marker = omittedAfter > 0 ? grepOmissionMarker(omittedAfter, Math.max(0, maxChars - 2)) : undefined + const firstHit = formatMcpGrepHit(entry.hits[0]!, maxCharsForFirstHit(maxChars, marker)) + return marker ? withGrepOmissionMarker([firstHit], marker) : firstHit + } + + return formatMcpGrepGroupForBudget(entry.hits, omittedAfter, maxChars) +} + +function formatMcpGrepGroupForBudget(hits: GrepHit[], omittedAfter: number, maxChars: number): string { + for (let locatorCount = hits.length; locatorCount >= 1; locatorCount -= 1) { + const omitted = omittedAfter + hits.length - locatorCount + const marker = omitted > 0 ? grepOmissionMarker(omitted, Math.max(0, maxChars - 2)) : undefined + const locators = hits.slice(0, locatorCount).map((hit) => formatMcpGrepLocator(hit)) + const locatorText = locators.join('\n') + const locatorOutput = marker ? withGrepOmissionMarker(locators, marker) : locatorText + + if (locatorOutput.length > maxChars) { + continue + } + + const snippetBudget = marker === undefined + ? maxChars - locatorText.length - 1 + : maxChars - locatorText.length - marker.length - 2 + if (snippetBudget <= 0) { + return locatorOutput + } + + const snippet = truncateGrepGroupSnippet(formatMcpGrepMergedSnippet(hits.slice(0, locatorCount)), snippetBudget) + if (!snippet) { + return locatorOutput + } + + const withSnippet = marker + ? withGrepOmissionMarker([...locators, snippet], marker) + : `${locatorText}\n${snippet}` + if (withSnippet.length <= maxChars) { + return withSnippet + } + } + + const marker = grepOmissionMarker(Math.max(0, omittedAfter + hits.length - 1), Math.max(0, maxChars - 2)) + const firstHit = formatMcpGrepHit(hits[0]!, maxCharsForFirstHit(maxChars, marker)) + return marker ? withGrepOmissionMarker([firstHit], marker) : firstHit +} + +function truncateGrepGroupSnippet(snippet: string, maxChars: number): string { + if (maxChars <= 0) { + return '' + } + if (snippet.length <= maxChars) { + return snippet + } + return truncateWithEllipsis(snippet, maxChars).replace(/\n+$/g, '') +} + function formatGrepPrefix(file: string, suffix: string, maxChars: number): string { const prefix = `${file}${suffix}` if (prefix.length <= maxChars) { @@ -1434,26 +1588,31 @@ function formatGrepPrefix(file: string, suffix: string, maxChars: number): strin return suffix.length <= maxChars ? suffix : suffix.slice(-maxChars) } -function fitGrepOutputWithMarker(rendered: string[], omitted: number, maxChars: number): string { +function fitGrepOutputWithMarker(rendered: McpRenderedGrepEntry[], omitted: number, maxChars: number): string { const kept = [...rendered] let omittedCount = omitted let marker = grepOmissionMarker(omittedCount, Math.max(0, maxChars - 2)) - while (kept.length > 1 && joinedLength(kept, marker) > maxChars) { - kept.pop() - omittedCount += 1 + while (kept.length > 1 && joinedLength(kept.map((entry) => entry.text), marker) > maxChars) { + omittedCount += kept.pop()!.matchCount marker = grepOmissionMarker(omittedCount, Math.max(0, maxChars - 2)) } if (!marker) { - return formatMcpGrepLineForBudget(kept[0]!, maxChars) + return formatMcpGrepLineForBudget(kept[0]!.text, maxChars) + } + + const keptText = kept.map((entry) => entry.text) + if (joinedLength(keptText, marker) <= maxChars) { + return withGrepOmissionMarker(keptText, marker) } - if (joinedLength(kept, marker) <= maxChars) { - return withGrepOmissionMarker(kept, marker) + const firstEntry = kept[0]! + if (firstEntry.matchCount > 1) { + return formatMcpGrepEntryForBudget({ hits: firstEntry.hits, text: firstEntry.text }, omittedCount, maxChars) } - const line = formatMcpGrepLineForBudget(kept[0]!, maxCharsForFirstHit(maxChars, marker)) + const line = formatMcpGrepLineForBudget(firstEntry.text, maxCharsForFirstHit(maxChars, marker)) return marker ? withGrepOmissionMarker([line], marker) : line } diff --git a/packages/mcp/test/mcp.test.ts b/packages/mcp/test/mcp.test.ts index 76756f8..e371816 100644 --- a/packages/mcp/test/mcp.test.ts +++ b/packages/mcp/test/mcp.test.ts @@ -840,6 +840,125 @@ describe('formatMcpGrepHits keeps its single-line ↩-joined format', () => { expect(output).toBe('src/auth.ts:10-12:3 | function verify() { ↩ return true ↩ }') }) + it('merges overlapping same-file context windows while keeping every match locator', () => { + const output = formatMcpGrepHits([ + makeGrepHit({ + range: { startLine: 2, endLine: 2 }, + line: 2, + column: 16, + match: 'NEEDLE', + snippet: "const before = true\nconst first = 'NEEDLE'\nconst shared = true", + snippetRange: { startLine: 1, endLine: 3 } + }), + makeGrepHit({ + range: { startLine: 4, endLine: 4 }, + line: 4, + column: 17, + match: 'NEEDLE', + snippet: "const shared = true\nconst second = 'NEEDLE'\nconst after = true", + snippetRange: { startLine: 3, endLine: 5 } + }) + ]) + + expect(output).toBe([ + 'src/auth.ts:2:16', + 'src/auth.ts:4:17', + '1 | const before = true', + "2 | const first = 'NEEDLE'", + '3 | const shared = true', + "4 | const second = 'NEEDLE'", + '5 | const after = true' + ].join('\n')) + expect(output.match(/3 \| const shared = true/g)).toHaveLength(1) + }) + + it('keeps no-context adjacent hits in the existing compact format', () => { + const output = formatMcpGrepHits([ + makeGrepHit({ + range: { startLine: 2, endLine: 2 }, + line: 2, + snippet: "const first = 'NEEDLE'", + snippetRange: { startLine: 2, endLine: 2 } + }), + makeGrepHit({ + range: { startLine: 3, endLine: 3 }, + line: 3, + snippet: "const second = 'NEEDLE'", + snippetRange: { startLine: 3, endLine: 3 } + }) + ]) + + expect(output).toBe([ + "src/auth.ts:2:3 | const first = 'NEEDLE'", + "src/auth.ts:3:3 | const second = 'NEEDLE'" + ].join('\n')) + }) + + it('counts omitted grouped grep entries by matches rather than groups', () => { + const output = formatMcpGrepHits([ + makeGrepHit({ + range: { startLine: 2, endLine: 2 }, + line: 2, + snippet: "a\nconst first = 'NEEDLE'\nb", + snippetRange: { startLine: 1, endLine: 3 } + }), + makeGrepHit({ + range: { startLine: 4, endLine: 4 }, + line: 4, + snippet: "b\nconst second = 'NEEDLE'\nc", + snippetRange: { startLine: 3, endLine: 5 } + }), + makeGrepHit({ + range: { startLine: 20, endLine: 20 }, + line: 20, + snippet: "x\nconst third = 'NEEDLE'\ny", + snippetRange: { startLine: 19, endLine: 21 } + }), + makeGrepHit({ + range: { startLine: 22, endLine: 22 }, + line: 22, + snippet: "y\nconst fourth = 'NEEDLE'\nz", + snippetRange: { startLine: 21, endLine: 23 } + }) + ], { maxTokens: 45 }) + + expect(output).toContain('src/auth.ts:2:3') + expect(output).toContain('src/auth.ts:4:3') + expect(output).not.toContain('src/auth.ts:20:3') + expect(output).not.toContain('src/auth.ts:22:3') + expect(output).toContain('matches_omitted=2') + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(45) + }) + + it('keeps all locators from the first merged group when its snippet is over budget', () => { + const output = formatMcpGrepHits([ + makeGrepHit({ + range: { startLine: 2, endLine: 2 }, + line: 2, + snippet: `${'a'.repeat(80)}\nconst first = 'NEEDLE'\n${'b'.repeat(80)}`, + snippetRange: { startLine: 1, endLine: 3 } + }), + makeGrepHit({ + range: { startLine: 4, endLine: 4 }, + line: 4, + snippet: `${'b'.repeat(80)}\nconst second = 'NEEDLE'\n${'c'.repeat(80)}`, + snippetRange: { startLine: 3, endLine: 5 } + }), + makeGrepHit({ + range: { startLine: 20, endLine: 20 }, + line: 20, + snippet: "const third = 'NEEDLE'", + snippetRange: { startLine: 20, endLine: 20 } + }) + ], { maxTokens: 35 }) + + expect(output).toContain('src/auth.ts:2:3') + expect(output).toContain('src/auth.ts:4:3') + expect(output).not.toContain('src/auth.ts:20:3') + expect(output).toContain('matches_omitted=1') + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(35) + }) + it('preserves first hits and appends a compact omission marker when maxTokens truncates matches', () => { const hits = Array.from({ length: 6 }, (_, index) => makeGrepHit({ @@ -913,6 +1032,34 @@ describe('grep_code MCP output budgeting', () => { expect(output).toContain('matches_omitted=') expect(output).not.toContain('NEEDLE_19') }) + + it('merges nearby grep_code context windows through callMcpTool', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-mcp-grep-context-')) + temporaryDirectories.push(repoRoot) + + await mkdir(join(repoRoot, 'src'), { recursive: true }) + await writeFile( + join(repoRoot, 'src', 'cluster.ts'), + [ + 'const before = true', + "const first = 'NEEDLE'", + 'const shared = true', + "const second = 'NEEDLE'", + 'const after = true' + ].join('\n'), + 'utf8' + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + + const output = await callMcpTool(repo, 'grep_code', { pattern: 'NEEDLE', path_glob: 'src/**', context_lines: 1 }) + expect(output).toContain('src/cluster.ts:2:16') + expect(output).toContain('src/cluster.ts:4:17') + expect(output.match(/3 \| const shared = true/g)).toHaveLength(1) + expect(output).toContain("2 | const first = 'NEEDLE'") + expect(output).toContain("4 | const second = 'NEEDLE'") + }) }) describe('read_chunk MCP output budgeting', () => { From 723042f5a3dfc253fe1555c3fe219f138396fbc4 Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 04:24:59 +0530 Subject: [PATCH 10/13] Budget find_symbol eval follow-ups --- packages/eval/src/index.ts | 10 +++---- packages/eval/test/eval.test.ts | 46 +++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 7 deletions(-) diff --git a/packages/eval/src/index.ts b/packages/eval/src/index.ts index cb7d45c..9e993b1 100644 --- a/packages/eval/src/index.ts +++ b/packages/eval/src/index.ts @@ -851,7 +851,7 @@ async function runCodesiftPolicy(repo: Awaited>, que rangesOverlap(matchingDef.range, query.expectedLineRange) if (!bodyResolves && matchingDef) { - const followUp = await readRangeSafely(repo, matchingDef.file, matchingDef.range) + const followUp = await readSymbolChunkSafely(repo, matchingDef) if (followUp !== null) { tokensToResolution += estimateTokenCount(followUp) callsToResolution = 2 @@ -1324,12 +1324,8 @@ async function readChunkSafely(repo: Awaited>, id: s } } -async function readRangeSafely(repo: Awaited>, file: string, range: Range): Promise { - try { - return await repo.readRange(file, range.startLine, range.endLine) - } catch { - return null - } +async function readSymbolChunkSafely(repo: Awaited>, definition: SymbolDefinition): Promise { + return readChunkSafely(repo, `${definition.file}:${definition.range.startLine}-${definition.range.endLine}`) } function normalizePath(value: string): string { diff --git a/packages/eval/test/eval.test.ts b/packages/eval/test/eval.test.ts index 4509ef9..ef70600 100644 --- a/packages/eval/test/eval.test.ts +++ b/packages/eval/test/eval.test.ts @@ -4,6 +4,8 @@ import { join } from 'node:path' import { afterEach, describe, expect, it } from 'vitest' +import { DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS, DEFAULT_MCP_READ_CHUNK_MAX_TOKENS } from '@codesift/mcp' + import { DEFAULT_MANIFEST_PATH, candidateFromRipgrep, candidateMatchesExpected, createEmptyManifest, diffLossBudgets, evaluateManifest, formatSummary, isRerankEvalEnabled, loadManifest, proveRoutingPolicy, runRipgrep, summarizeEmptyRun, type GoldenQuery, type GoldenQueryType, type RipgrepHit } from '../src/index.js' const temporaryDirectories: string[] = [] @@ -137,6 +139,50 @@ describe('@codesift/eval', () => { expect(codesiftRun?.tokensToResolution).toBeLessThanOrEqual(730) }, 30_000) + it('accounts body-less symbol follow-ups through the read_chunk MCP budget', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-eval-symbol-followup-budget-')) + temporaryDirectories.push(repoRoot) + + await mkdir(join(repoRoot, 'src'), { recursive: true }) + + const makeFunction = (returnValue: string) => [ + 'export function inflateTarget() {', + ...Array.from({ length: 140 }, (_, index) => ` const noisy${index} = '${returnValue}_${index}_${'x'.repeat(180)}'`), + ` return '${returnValue}'`, + '}', + '' + ].join('\n') + + await Promise.all([ + writeFile(join(repoRoot, 'src', 'a.ts'), makeFunction('target'), 'utf8'), + writeFile(join(repoRoot, 'src', 'b.ts'), makeFunction('otherB'), 'utf8'), + writeFile(join(repoRoot, 'src', 'c.ts'), makeFunction('otherC'), 'utf8'), + writeFile(join(repoRoot, 'src', 'd.ts'), makeFunction('otherD'), 'utf8') + ]) + + const summary = await evaluateManifest({ + repos: [{ id: 'symbol-followup-budget', language: 'typescript', repoPath: repoRoot }], + queries: [ + { + id: 'symbol-followup-budget', + repoId: 'symbol-followup-budget', + queryType: 'symbol-def', + query: 'inflateTarget', + expected: [{ file: 'src/a.ts', symbol: 'inflateTarget' }], + expectedLineRange: { startLine: 1, endLine: 143 } + } + ] + }, { resultLimit: 10, inspectionLimit: 5, latencyToleranceMs: Number.POSITIVE_INFINITY }) + + const codesiftRun = summary.runs.find((run) => run.tool === 'codesift') + + expect(codesiftRun?.taskSuccess).toBe(true) + expect(codesiftRun?.callsToResolution).toBe(2) + expect(codesiftRun?.tokensToResolution).toBeLessThanOrEqual( + DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS + DEFAULT_MCP_READ_CHUNK_MAX_TOKENS + ) + }, 30_000) + it('measures end-to-end calls-to-resolution truthfully across both tools', async () => { const manifest = await loadManifest(DEFAULT_MANIFEST_PATH) const m3Manifest = { From 43c642c78f273570ddc00351a392f58cdd643d21 Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 14:05:50 +0530 Subject: [PATCH 11/13] chore: pin Node 22 and ensure native SQLite install hooks run - .nvmrc 20 -> 22 so better-sqlite3 uses published macOS arm64 prebuilds - add .npmrc (ignore-scripts=false) so the native SQLite deps run install hooks - smoke-pack-install: force ignore-scripts=false for the packed-install check - README: document the Node 22 + .npmrc requirement Pre-existing local environment setup; committed alongside the graph work because the build/test gate depends on Node 22. Co-Authored-By: Claude Opus 4.8 (1M context) --- .npmrc | 1 + .nvmrc | 2 +- README.md | 3 +++ scripts/smoke-pack-install.mjs | 7 ++++++- 4 files changed, 11 insertions(+), 2 deletions(-) create mode 100644 .npmrc diff --git a/.npmrc b/.npmrc new file mode 100644 index 0000000..e60ae71 --- /dev/null +++ b/.npmrc @@ -0,0 +1 @@ +ignore-scripts=false diff --git a/.nvmrc b/.nvmrc index 209e3ef..2bd5a0a 100644 --- a/.nvmrc +++ b/.nvmrc @@ -1 +1 @@ -20 +22 diff --git a/README.md b/README.md index 33d18d0..e402cf9 100644 --- a/README.md +++ b/README.md @@ -84,6 +84,9 @@ packages/ ## Quickstart +Use Node 22 locally (`.nvmrc`) on macOS arm64 so `better-sqlite3` can use its published prebuilds. +If you normally run with `ignore-scripts=true`, this repo overrides it via `.npmrc` because the native SQLite deps must run their install hooks. + ```bash pnpm install pnpm build diff --git a/scripts/smoke-pack-install.mjs b/scripts/smoke-pack-install.mjs index 57dcb04..8b93207 100644 --- a/scripts/smoke-pack-install.mjs +++ b/scripts/smoke-pack-install.mjs @@ -15,6 +15,11 @@ function run(command, args, options = {}) { cwd: repoRoot, encoding: 'utf8', stdio: 'pipe', + env: { + ...process.env, + npm_config_ignore_scripts: 'false', + ...options.env + }, ...options }) @@ -101,7 +106,7 @@ async function main() { run(npmBin, ['init', '-y'], { cwd: installDirectory }) const installLog = run( npmBin, - ['install', '--foreground-scripts', '--loglevel', 'verbose', coreTarball, mcpTarball, cliTarball], + ['install', '--ignore-scripts=false', '--foreground-scripts', '--loglevel', 'verbose', coreTarball, mcpTarball, cliTarball], { cwd: installDirectory } ) From 37809e04d35849472c5a0e945f26c56779310061 Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 14:05:50 +0530 Subject: [PATCH 12/13] feat(core,mcp): persistent relational code graph (callers/refs/imports/impl/impact) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a persistent, incrementally-maintained `edges` graph so codesift can answer relational queries grep structurally cannot — who calls/references X, what imports a file, who implements an interface, and bounded impact ("what breaks if I change X") — all index-time, zero-egress, and honestly labeled import-resolved vs name-only. Phases (per GRAPH_PLAN.md): - 0: edges table + indexes, per-file plumbing in applySyncChanges, SCHEMA_VERSION 7->8, Edge/EdgeKind/EdgeResolution types. - 1: index-time edge extraction for TS/JS + Python riding the cached chunking AST; caller attribution by enclosing range; with_usages now reads the index instead of re-parsing files at query time. - 2: core findCallers/findReferences/findImporters + MCP tools find_callers / find_refs / find_importers. - 3: TS implements/extends heritage edges; name-only call edges for Go/Java/Ruby/Rust; findImplementers + who_implements. - 4: opt-in find_symbol relations bundle (with_callers) and bounded impact(name, depth) BFS (hard node/depth caps + cycle guard) + impact tool. Reviewed and fixed: default-export resolution, name-only leakage into disambiguated queries, shadow/member-access mislabeling, case-insensitive edge matching, and stale-inbound-edge cleanup on file removal. Honesty: never type-resolved; name-only surfaced as approx. Additive — no ranking/fusion/search behavior changed. Gate green on Node 22 (pnpm build && typecheck && test: 11 files, 177 tests). Co-Authored-By: Claude Opus 4.8 (1M context) --- GRAPH_PLAN.md | 188 ++ packages/core/src/chunking.ts | 9 +- packages/core/src/index.ts | 8 + packages/core/src/repo.ts | 1912 +++++++++++++++-- packages/core/src/types.ts | 124 +- packages/core/test/core.test.ts | 788 ++++++- .../eval/fixtures/heritage-ts/src/contract.ts | 9 + .../eval/fixtures/heritage-ts/src/impl.ts | 9 + packages/mcp/src/index.ts | 678 +++++- packages/mcp/test/mcp.test.ts | 426 +++- 10 files changed, 3883 insertions(+), 268 deletions(-) create mode 100644 GRAPH_PLAN.md create mode 100644 packages/eval/fixtures/heritage-ts/src/contract.ts create mode 100644 packages/eval/fixtures/heritage-ts/src/impl.ts diff --git a/GRAPH_PLAN.md b/GRAPH_PLAN.md new file mode 100644 index 0000000..f9e5403 --- /dev/null +++ b/GRAPH_PLAN.md @@ -0,0 +1,188 @@ +# codesift Graph Plan — the relational moat (callers / refs / impact) + +> North goal: make codesift answer the questions `rg` *structurally cannot* — "who calls X", "what uses X", "what imports this", "who implements this", "what breaks if I change X" — in **one warm, fresh, zero-egress call**. This is the durable moat: an index can hold relationships; a stateless line-matcher never can. + +Status: proposal (2026-06). All file:line anchors below are verified against current `main`. + +--- + +## The one-paragraph thesis + +We already compute the hard part and throw it away. `findTypeScriptUsages` (`repo.ts:2415-2524`) does real **import resolution** — it follows `import` declarations, resolves the module specifier back to the defining file, tracks the local binding name (named / default / namespace), then AST-walks for genuine use sites while excluding the definition and the import statements themselves. `findPythonUsages` (`repo.ts:2559+`) does the regex equivalent. But this runs **at query time, re-parsing files from disk on every call** (`repo.ts:2423-2436`), **only for the top hit** (`attachUsagesToTopDefinitionHit`, `repo.ts:2350`), **capped at 5** (`repo.ts:2424`), **opt-in/off by default**, and is **discarded after each query**. The `symbols` table stores **definitions only** (`repo.ts:1198-1208`); there is no edges table anywhere in the schema (`repo.ts:1155-1221`). + +The plan: **relocate that resolution from query time to index time, persist it as an `edges` table, and serve relational queries as cheap indexed lookups.** This is a moat move *and* a runtime win — the per-query AST parse disappears, `with_usages` becomes free + always-on + uncapped, and `find_callers` / `find_refs` / `find_importers` / `who_implements` come essentially for free on top of the persisted edges. + +--- + +## The data model + +One new table, derived per **source** file (the file that contains the reference), so it maintains incrementally through the path that already exists. + +```sql +create table if not exists edges( + id integer primary key autoincrement, + src_file text not null references files(path) on delete cascade, + src_line integer not null, + src_symbol text, -- enclosing caller symbol (nullable; resolved by containing range) + dst_name text not null, -- referenced identifier + dst_file text, -- resolved defining file (import-resolved); null for name-only + edge_kind text not null, -- 'ref' | 'call' | 'import' | 'implements' | 'extends' + resolution text not null, -- 'import-resolved' | 'name-only' (NEVER a fake 'type-resolved') + language text +); +create index if not exists idx_edges_dst on edges(dst_name, dst_file); +create index if not exists idx_edges_src on edges(src_file); +create index if not exists idx_edges_kind on edges(edge_kind); +``` + +The relational queries then become: + +- **who calls / uses X** → `where dst_name = ? and (dst_file = ? or dst_file is null) and edge_kind in ('call','ref')`. The `dst_file` match is what disambiguates the three different `validate` defs in `fixtures/collision-ts/` — the exact collision case `rg` can't resolve. +- **what imports this file** → `where dst_file = ? and edge_kind = 'import'`. +- **who implements / extends Y** → `where dst_name = ? and edge_kind in ('implements','extends')`. + +### Two design rules that keep this honest and fresh + +1. **Resolution honesty (don't repeat the SCIP mistake).** `MOAT_PLAN.md` already killed "a SCIP graph the scanners can't back." Every edge carries `resolution`. `import-resolved` = we followed the binding (high confidence — what TS/Python already do). `name-only` = identifier matched without resolving the import (the fallback for the Go/Java/Ruby/Rust regex scanners), and it must be surfaced to the agent as *"may include unrelated same-named symbols."* We **never** emit a `type-resolved` label — we don't have a type checker. +2. **Resolve `dst_symbol` lazily by name+file at query time** (via the existing `symbols` table containing-range lookup), rather than storing a hard `dst_symbol_id`. That keeps edges robust when a definition just moves lines — only the *source* file's edges need re-extraction, never every inbound edge. + +### Schema-version decision + +Adding `edges` is a schema change. Existing indexes won't have edges until rebuilt. Recommendation: **bump `SCHEMA_VERSION`** so the existing compat path (`getIndexCompatibility`, `repo.ts:1269`; `schema_version_mismatch` → guided `--rebuild`) triggers a one-time rebuild that backfills edges. Clean and correct; the one-time reindex cost is acceptable and already a supported flow. (Alternative — add the table without a bump and lazy-backfill on next per-file sync — leaves the graph incomplete until every file is touched; rejected.) + +--- + +## Phasing + +### Phase 0 — Foundation (schema + per-file plumbing) + +Lay the table and the maintenance plumbing **before** any extraction logic, so extraction and the read-side tools can be built against a real (if initially empty) edges store in parallel. + +- Add the `edges` table + indexes to the schema block (`repo.ts:1155-1221`). +- Add `delete from edges` to `clearIndex` (`repo.ts:1228-1237`). +- In `applySyncChanges` (`repo.ts:1327`), add a per-file `deleteEdgesBySrcFile` prepared statement and call it in the replacement loop alongside `deleteSymbolsByFile`/`deleteChunksByFile` (`repo.ts:1391-1398`), plus an `insertEdge` statement and an **extraction hook point** in the changed-files insert loop (`repo.ts:1411`) — initially a no-op `extractEdges(file) -> []`. +- Bump `SCHEMA_VERSION`. +- Add `Edge` / `EdgeKind` / `EdgeResolution` types to `types.ts`; widen `SymbolUsage.resolution` (`types.ts:34`) from the lone `'import-resolved'` literal to `'import-resolved' | 'name-only'`. + +This phase ships dark (no behavior change — extractor returns `[]`), so it's safe to land first. + +### Phase 1 — Index-time extraction for TS/JS + Python (the core move) + +Promote the existing query-time resolvers to index-time extractors that emit `Edge[]` for a single file, and repoint `with_usages` to read persisted edges. + +- Refactor `findTypeScriptUsages` (`repo.ts:2415-2524`) into a pure `extractTypeScriptEdges(file, content, sourceFile)` that returns all import-resolved `ref`/`call`/`import` edges for the file (drop the 5-cap; the cap was a query-time budget, not a correctness limit). **Free-ride the existing parse**: TS/JS files are already AST-parsed during chunking (`buildTypeScriptChunks`, `chunking.ts:70-170`) — extract edges in that same pass so marginal index cost is ~zero. +- Same for `findPythonUsages` → `extractPythonEdges`. +- **Caller attribution:** for each use site, resolve `src_symbol` by looking up the `symbols` row whose `[start_line, end_line]` contains `src_line` (cheap range query; the index `idx_chunks_file_range` / symbols table already supports it). A use site inside a function body → that function is the caller; mark the edge `call` when the reference is in call position, else `ref`. +- Rewrite `attachUsagesToTopDefinitionHit` (`repo.ts:2344-2360`) to **read from the `edges` table** instead of parsing — an indexed `select` keyed by the definition's name+file. `with_usages` can now default-on safely (it's a bounded indexed read, no longer an O(files) AST walk — this removes the never-slower hazard `MOAT_NEXT.md:66` flagged). + +End state: `with_usages` is free, always-fresh, and complete; the graph is populated for the two languages where resolution is real. + +### Phase 2 — Relational query API + MCP tools + +Expose the graph. Read-side only; depends on Phase 0's schema (can be built in parallel with Phase 1 against seeded rows). + +- Core: `findCallers(name, opts)`, `findReferences(name, opts)`, `findImporters(file)` on the `Repo` class — each a thin indexed `select` over `edges` joined to `symbols` for the def-site, with the same token-budgeting treatment as existing hits. +- MCP: register `find_callers`, `find_refs`, `find_importers` — add to `MCP_TOOL_NAMES` (`mcp/src/index.ts:49-55`), zod input schemas (alongside `:176-219`), `getToolDefinitions` (`:256-324`), the dispatch `switch (name)` (`:431-452`), and `registerTool` (`:521-533`), with formatters mirroring `formatMcpSymbols` (`:1208`). Update `MCP_SERVER_INSTRUCTIONS` (`:167-171`) routing. + +### Phase 3 — `implements`/`extends` + multi-language name-only edges + +- TS `implements`/`extends` edges from the AST (`heritageClauses` — already in the tree we parse). +- `name-only` edges for Go / Java / Ruby / Rust by piggybacking on their existing regex scanners (`chunking.ts:172-503`), masking strings/comments via the existing `maskCStyleSyntax`. Labeled `name-only`, surfaced honestly. +- `who_implements` MCP tool. + +### Phase 4 — One-call relational answer + impact analysis + +- Upgrade `find_symbol` to optionally return **definition body + top callers + usages + same-file neighbors** in a single response (answers the open question in `MOAT_NEXT.md:148`; collapses the agent's `find_symbol → grep → read → grep` refactor loop into one call). +- `impact(name, depth)` — transitive callers up to depth N (bounded), the "what breaks if I change this" query, with a hard node/time budget. + +--- + +## Parallelization + +### Dependency graph + +``` +PHASE 0 (foundation — lands first, ships dark) + schema + clearIndex + applySyncChanges hook + types + │ + ├─────────────────────────┬───────────────────────────┐ + ▼ ▼ ▼ +PHASE 1 (extraction) PHASE 2 (tools/read API) (fixtures, seeded rows) + TS/JS + Python edges build against schema with usages-ts + collision-ts + + with_usages repoint hand-seeded edge rows, already exist — reuse + │ wire to real edges when 1 lands + └─────────────┬───────────┘ + ▼ + PHASE 3 (implements/extends + name-only langs) + ▼ + PHASE 4 (one-call answer + impact) +``` + +The unlock: **Phase 1 (extraction) and Phase 2 (read-side tools) are independent once Phase 0's schema exists.** Phase 2 can be developed and tested against hand-seeded `edges` rows, then flipped to live data the moment Phase 1 lands. That's two tracks running concurrently after a small, fast Phase 0. + +### Tracks + +- **Track A — extraction (the novel work).** Phase 1 → Phase 3 extraction. Owns the edge-extraction algorithms, caller attribution, and resolution-confidence semantics. +- **Track B — read-side surface (mechanical).** Phase 2 query methods + MCP tools, developed against seeded rows. Then Phase 4 wiring. +- **Track C — foundation (mechanical).** Phase 0; lands first, unblocks A and B. + +### File-conflict hotspot + +`packages/core/src/repo.ts` is the shared file. Conflicts are avoidable because the work lands in **different functions**, but two pairs need ordering discipline: + +| Work | Functions in `repo.ts` | Note | +|------|------------------------|------| +| Phase 0 | schema block (`1155`), `clearIndex` (`1228`), `applySyncChanges` (`1327`) | lands first | +| Phase 1 | `extractTypeScriptEdges`/`extractPythonEdges` (was `2415`/`2559`), hook in `applySyncChanges` insert loop (`1411`) | **touches `applySyncChanges` too → land after Phase 0, same owner for that function** | +| Phase 2 | new `findCallers`/`findReferences`/`findImporters` (new functions, disjoint) | safe in parallel | + +`packages/mcp/src/index.ts` is touched only by Phase 2/4 (new tools) — disjoint from the `repo.ts` extraction work, so Track B's MCP edits never collide with Track A. `types.ts` is touched by Phase 0 (and read by all) — land the type additions in Phase 0 so the rest compile against them. + +**Recommendation:** give each track its own branch/worktree. The only function needing a coordination conversation is `applySyncChanges` (Phase 0 lands the hook point; Phase 1 fills the extractor). + +--- + +## Delegation: Claude (creative) vs pi (mechanical) + +Same split principle the moat work used (`MOAT_NEXT.md:99`): **Claude owns novel algorithms, resolution/confidence semantics, caller-attribution logic, and anything that changes what an edge *means* or how a tool *answers*; pi owns deterministic schema/plumbing, MCP wiring, formatters, and fixtures under a written spec** with exact file:line targets, the acceptance test, and a guardrail. + +| Task | Owner | Why this split | +|------|-------|----------------| +| `edges` schema + index design, `resolution` taxonomy, schema-version decision | **Claude** | Data-model + honesty semantics | +| Phase 0 wiring — table in schema block, `clearIndex` line, `applySyncChanges` delete/insert/hook, `SCHEMA_VERSION` bump, `types.ts` additions | **pi** | Pure mechanical plumbing once schema is fixed | +| Phase 1 extraction algorithm — refactor resolvers to pure per-file extractors, drop the 5-cap, edge-kind (`call` vs `ref`) classification | **Claude** | Novel: changes what we extract and how we classify it | +| Phase 1 **caller attribution** — containing-range `src_symbol` resolution design | **Claude** | Judgment: how to attribute a use site to its enclosing symbol | +| Phase 1 `attachUsagesToTopDefinitionHit` repoint to indexed read + default-on flip | **pi** | Mechanical once the read shape is specified | +| Phase 2 tool **semantics** — what each tool returns, output shape, routing-instruction wording, collision/ambiguity handling | **Claude** | Changes the agent-facing contract | +| Phase 2 MCP **plumbing** — `MCP_TOOL_NAMES`, zod schemas, `getToolDefinitions`, dispatch `switch`, `registerTool`, formatters | **pi** | Deterministic wiring against shipped patterns | +| Phase 2 core read methods (`findCallers`/`findReferences`/`findImporters`) | **pi** | Thin indexed selects once the SQL shape is specified | +| Phase 3 `implements`/`extends` extraction + name-only multi-language design + honest surfacing | **Claude** | Confidence labeling + cross-language judgment | +| Phase 3 regex-scanner edge plumbing per language | **pi** | Mechanical, mirrors the chunking scanners | +| Phase 4 one-call answer composition + `impact` traversal + budgets | **Claude** | Novel composition + bounded-traversal algorithm | +| Correctness fixtures/tests (reuse `fixtures/usages-ts`, `fixtures/collision-ts`) | **pi** | Mechanical fixture/test authoring to a spec | + +**Net:** Claude carries the extraction algorithms (1, 3, 4 extraction), caller attribution, tool/answer semantics, and the data model. pi carries Phase 0 foundation, all MCP/SQL plumbing, the with_usages repoint, the regex-scanner edge wiring, and fixtures. **Wave 1 = pi (Phase 0) lands fast → then Claude (Phase 1 extraction) ∥ pi (Phase 2 plumbing on seeded rows) run concurrently.** + +### How to hand a task to pi + +Use the `pi-delegate` skill, **one task per delegation**, each with: target files + line anchors, the exact change, the acceptance test (`pnpm test` or a named vitest), and the guardrail *"do not modify extraction/resolution/ranking logic — plumbing only."* + +> **Local gotcha:** `pi-delegate` crashes on startup in this repo via the broken agent-view extension — invoke it with `-ne` to disable extensions. Build/test under Node 22 (`.nvmrc`), since `better-sqlite3`'s native binary needs it. + +Good first pi delegations (independent, fully specified): **Phase 0 foundation** and **the Phase 2 MCP plumbing against seeded edge rows**. + +--- + +## What NOT to do (guardrails) + +- **Don't claim type resolution.** No `type-resolved` label, ever — we have no type checker. Label `import-resolved` vs `name-only` and surface `name-only`'s ambiguity to the agent. +- **Don't make extraction unbounded on the default path.** Edge extraction rides the existing per-file parse (index time, once per file change) — never an O(files) walk at query time. That is the whole point and the thing that keeps "never slower than rg" intact. +- **Don't break zero-egress / offline.** The graph is pure local AST/regex work — no network, no model. Keep it inside `pnpm run test:offline`. +- **Don't store hard `dst_symbol_id` pointers** — resolve by name+file at query time so a moved definition doesn't require rewriting every inbound edge. +- **Don't over-promise multi-language.** Go/Java/Ruby/Rust ship as `name-only` first; upgrade per language to import-resolved only when a real lightweight import parser exists for it. + +--- + +## Proof (lightweight, product-correctness — not a new benchmark) + +Reuse the fixtures that already exist: `fixtures/usages-ts/` (a `parseToken` def with two real call sites in `src/api.ts` + `src/worker.ts`) proves caller/usage extraction; `fixtures/collision-ts/` (three `validate` defs) proves `dst_file` disambiguation. Add core tests asserting: (1) edges survive incremental re-sync of a changed source file, (2) `find_callers` resolves the right `validate` by file, (3) `with_usages` returns identical results read-from-index as the old parse path. These are correctness gates on the feature, not a competitive eval. diff --git a/packages/core/src/chunking.ts b/packages/core/src/chunking.ts index 38b50ff..42000a0 100644 --- a/packages/core/src/chunking.ts +++ b/packages/core/src/chunking.ts @@ -43,6 +43,12 @@ const MAX_STRUCTURAL_CHUNK_TOKENS = 1_200 const SPLIT_WINDOW_LINES = 90 const SPLIT_OVERLAP_LINES = 12 +const typeScriptSourceFileCache = new WeakMap() + +export function getCachedTypeScriptSourceFile(file: ScannedFile): ts.SourceFile | undefined { + return typeScriptSourceFileCache.get(file) +} + export function buildChunks(file: ScannedFile): ChunkRecord[] { let chunks: ChunkRecord[] = [] @@ -76,6 +82,7 @@ function buildTypeScriptChunks(file: ScannedFile): ChunkRecord[] { true, scriptKindFromPath(file.relativePath) ) + typeScriptSourceFileCache.set(file, sourceFile) const chunks: ChunkRecord[] = [] @@ -1195,7 +1202,7 @@ function countRubyClosers(line: string): number { return /^end\b/.test(line.trim()) ? 1 : 0 } -function maskCStyleSyntax(lines: string[]): string[] { +export function maskCStyleSyntax(lines: string[]): string[] { const maskedLines: string[] = [] let inBlockComment = false let stringQuote: '"' | "'" | '`' | null = null diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index ecd26cd..f783f18 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -1,8 +1,14 @@ export type { + EdgeResult, EmbeddingBatchOptions, EmbeddingProvider, EmbeddingRole, + FindEdgeOptions, + FindImportersOptions, FindSymbolOptions, + ImpactNode, + ImpactOptions, + ImpactResult, GrepHit, GrepOptions, IndexCompatibilitySnapshot, @@ -27,6 +33,8 @@ export type { StopWatching, SymbolDefinition, SymbolKind, + SymbolNeighbor, + SymbolRelations, SymbolUsage, SyncOptions, SyncProgressEvent, diff --git a/packages/core/src/repo.ts b/packages/core/src/repo.ts index cd327b6..c42867e 100644 --- a/packages/core/src/repo.ts +++ b/packages/core/src/repo.ts @@ -8,14 +8,50 @@ import Database from 'better-sqlite3' import { minimatch } from 'minimatch' import * as sqliteVec from 'sqlite-vec' -import { buildChunks, type ChunkRecord } from './chunking.js' +import { buildChunks, getCachedTypeScriptSourceFile, maskCStyleSyntax, type ChunkRecord } from './chunking.js' import { readConfig } from './config.js' import { DEFAULT_EMBEDDING_PROVIDER_ID, expandTermToOrGroup, getEmbeddingProvider, isCloudEmbeddingProvider, isLearnedEmbeddingProvider } from './embedding.js' -import { isCodeLanguage, isDocumentationLanguage, isPythonLike, isTypeScriptLike } from './languages.js' +import { isCodeLanguage, isDocumentationLanguage, isGoLike, isJavaLike, isPythonLike, isRubyLike, isRustLike, isTypeScriptLike } from './languages.js' import { resolveReranker } from './reranker.js' import { scanRepository, scanRepositoryManifest, type ScannedFile } from './scan.js' import { prepareForCloud } from './secret-scan.js' -import { DEFAULT_SEARCH_K, type EmbeddingProvider, type FindSymbolOptions, type GrepHit, type GrepOptions, type IndexCompatibilitySnapshot, type IndexCompatibilityStatus, type ReadChunkOptions, type ReadRangeOptions, type RerankResult, type Repo, type RepoOptions, type RepoStaleReason, type RepoStatus, type RepoSyncStatus, type SearchHit, type SearchOptions, type SearchReasonTag, type StopWatching, type SymbolDefinition, type SymbolKind, type SymbolUsage, type SyncOptions, type SyncResult, type VectorSearchStatus, type WatchOptions } from './types.js' +import { + DEFAULT_SEARCH_K, + type Edge, + type EdgeResult, + type EmbeddingProvider, + type FindEdgeOptions, + type FindImportersOptions, + type FindSymbolOptions, + type GrepHit, + type GrepOptions, + type ImpactNode, + type ImpactOptions, + type ImpactResult, + type IndexCompatibilitySnapshot, + type IndexCompatibilityStatus, + type ReadChunkOptions, + type ReadRangeOptions, + type Repo, + type RepoOptions, + type RepoStaleReason, + type RepoStatus, + type RepoSyncStatus, + type RerankResult, + type SearchHit, + type SearchOptions, + type SearchReasonTag, + type StopWatching, + type SymbolDefinition, + type SymbolKind, + type SymbolNeighbor, + type SymbolRelations, + type SymbolUsage, + type SyncOptions, + type SyncResult, + type VectorSearchStatus, + type WatchOptions +} from './types.js' interface ChunkRow { id: string @@ -70,6 +106,54 @@ interface RankedChunkRow { reasons: Set } +interface SourceSymbolRange { + startLine: number + endLine: number + symbol?: string + kind?: SymbolKind | null +} + +interface EdgeUsageRow { + id: number + src_file: string + src_line: number + language: string | null + resolution: Edge['resolution'] +} + +interface EdgeResultRow { + id: number + src_file: string + src_line: number + src_symbol: string | null + edge_kind: Edge['edgeKind'] + resolution: Edge['resolution'] + language: string | null +} + +interface ImpactNodeRow extends EdgeResultRow { + name: string + depth: number +} + +interface EdgeBinding { + dstName: string + dstFile?: string + resolution: Edge['resolution'] +} + +interface DefinitionEdgeTarget { + names: string[] + resolutionMode: Edge['resolution'] + file?: string +} + +interface ReadRowsResult { + items: T[] + tokenTruncated: boolean + readFailures: number +} + interface IndexedFileRow { path: string language: string @@ -128,7 +212,7 @@ interface SyncApplyContext { completedAt: string } -const SCHEMA_VERSION = '7' +const SCHEMA_VERSION = '8' const DEFAULT_RRF_K = 60 const DEFAULT_VECTOR_LIMIT = 50 const DEFAULT_PATH_FILTERED_LIMIT = 200 @@ -155,6 +239,11 @@ const MIN_RELAXATION_ROWS = 3 // Above this many exact rows the identifier collides across the repo, so picking // one body to inline would be misleading — keep every row compact instead. const FIND_SYMBOL_INLINE_MAX_EXACT_ROWS = 3 +const FIND_SYMBOL_RELATION_MAX_SITES = 5 +const FIND_SYMBOL_RELATION_MAX_NEIGHBORS = 4 +const DEFAULT_IMPACT_DEPTH = 2 +const DEFAULT_IMPACT_MAX_NODES = 50 +const HARD_MAX_IMPACT_BOUND = 50 // When an identifier-shaped query collides across ≥2 definitions, single_best does // NOT collapse to one (that would hide the collision and silently pick a winner); // instead it returns up to this many candidates with an "ambiguous: N defs" hint so @@ -653,56 +742,62 @@ export class SqliteRepo implements Repo { this.ensureIndexCompatibleForQueries(db) const kinds = normalizeKinds(options?.kind) + const exactWhere = ['lower(name) = lower(?)'] + const exactParams: string[] = [name] + if (kinds.length > 0) { + exactWhere.push(`kind in (${kinds.map(() => '?').join(', ')})`) + exactParams.push(...kinds) + } + if (options?.pathGlob) { + exactWhere.push('codesift_minimatch(file_path, ?) = 1') + exactParams.push(options.pathGlob) + } - const exactWhere = ['lower(name) = lower(?)'] - const exactParams: string[] = [name] - if (kinds.length > 0) { - exactWhere.push(`kind in (${kinds.map(() => '?').join(', ')})`) - exactParams.push(...kinds) - } - if (options?.pathGlob) { - exactWhere.push('codesift_minimatch(file_path, ?) = 1') - exactParams.push(options.pathGlob) - } + const exactRows = db + .prepare( + ` + select id, name, file_path, start_line, end_line, kind, signature, parent, language + from symbols + where ${exactWhere.join(' and ')} + order by file_path asc, start_line asc + ` + ) + .all(...exactParams) - const exactRows = db - .prepare( - ` - select id, name, file_path, start_line, end_line, kind, signature, parent, language - from symbols - where ${exactWhere.join(' and ')} - order by file_path asc, start_line asc - ` - ) - .all(...exactParams) + const partialRows = + exactRows.length > 0 + ? [] + : selectPartialSymbolRows(db, name, kinds, options?.pathGlob) - const partialRows = - exactRows.length > 0 - ? [] - : selectPartialSymbolRows(db, name, kinds, options?.pathGlob) + const rows = [...exactRows, ...partialRows] - const rows = [...exactRows, ...partialRows] + const canEnrichTopExactRow = exactRows.length > 0 && exactRows.length <= FIND_SYMBOL_INLINE_MAX_EXACT_ROWS - // One-call moat: inline the verbatim enclosing-symbol body for the top exact - // match when the lookup is unambiguous, so an identifier query resolves - // without a mandatory follow-up read. Disk-fresh, capped, never throws. - let topBody: string | undefined - if ( - options?.withBody !== false && - exactRows.length > 0 && - exactRows.length <= FIND_SYMBOL_INLINE_MAX_EXACT_ROWS - ) { - const top = exactRows[0]! - try { - const source = await this.readRange(top.file_path, top.start_line, top.end_line) - const { body } = capInlineBody(source, chunkLocator(top.file_path, top.start_line, top.end_line)) - if (body) { - topBody = body + // One-call moat: inline the verbatim enclosing-symbol body for the top exact + // match when the lookup is unambiguous, so an identifier query resolves + // without a mandatory follow-up read. Disk-fresh, capped, never throws. + let topBody: string | undefined + if (options?.withBody !== false && canEnrichTopExactRow) { + const top = exactRows[0]! + try { + const source = await this.readRange(top.file_path, top.start_line, top.end_line) + const { body } = capInlineBody(source, chunkLocator(top.file_path, top.start_line, top.end_line)) + if (body) { + topBody = body + } + } catch { + // A disk failure or out-of-range read leaves every row compact. + } + } + + let topRelations: SymbolRelations | undefined + if (options?.withCallers === true && canEnrichTopExactRow) { + try { + topRelations = await this.readFindSymbolRelations(db, exactRows[0]!) + } catch { + // Relation bundling is best-effort and must never fail the base lookup. } - } catch { - // A disk failure or out-of-range read leaves every row compact. } - } return rows.map((row, index) => { const definition: SymbolDefinition = { @@ -728,11 +823,16 @@ export class SqliteRepo implements Repo { definition.language = row.language } - // Only the top exact row (rows[0] when exactRows is non-empty) carries a body. + // Only the top exact row (rows[0] when exactRows is non-empty) carries + // the optional single-call enrichments. if (index === 0 && topBody !== undefined) { definition.body = topBody } + if (index === 0 && topRelations !== undefined) { + definition.relations = topRelations + } + return definition }) } finally { @@ -740,6 +840,192 @@ export class SqliteRepo implements Repo { } } + async findCallers(name: string, options?: FindEdgeOptions): Promise { + return this.findDefinitionEdges(name, options, ['call'], true) + } + + async findReferences(name: string, options?: FindEdgeOptions): Promise { + return this.findDefinitionEdges(name, options, ['call', 'ref'], false) + } + + async findImplementers(name: string, options?: FindEdgeOptions): Promise { + return this.findDefinitionEdges(name, options, ['implements', 'extends'], false) + } + + async findImporters(file: string, options?: FindImportersOptions): Promise { + if (!file.trim() || !existsSync(this.indexPath)) { + return [] + } + + const releaseDatabase = await this.enterDatabaseUser() + let rows: EdgeResultRow[] = [] + try { + const db = this.openDatabase() + this.ensureIndexCompatibleForQueries(db) + rows = selectImporterEdgeRows(db, file) + } finally { + releaseDatabase() + } + + return (await readEdgeResultsFromRows(this.root, rows, options?.maxTokens)).items + } + + async impact(name: string, options?: ImpactOptions): Promise { + const depthLimit = normalizeImpactDepth(options?.depth) + const maxNodes = normalizeImpactMaxNodes(options?.maxNodes) + if (!name.trim() || !existsSync(this.indexPath)) { + return { nodes: [], depthLimit, maxNodes } + } + + const releaseDatabase = await this.enterDatabaseUser() + let rows: ImpactNodeRow[] = [] + let depthCapped = false + let nodesCapped = false + try { + const db = this.openDatabase() + this.ensureIndexCompatibleForQueries(db) + + const initialDefinitions = selectExactDefinitionRows(db, name, normalizeKinds(options?.kind), options?.pathGlob) + if (initialDefinitions.length === 0) { + return { nodes: [], depthLimit, maxNodes } + } + + const defaultExportCache = new Map>>() + const queue: Array<{ name: string; targetFiles: string[]; depth: number }> = [ + { name, targetFiles: [...new Set(initialDefinitions.map((definition) => definition.file_path))], depth: 0 } + ] + const visitedTargets = new Set(initialDefinitions.map((definition) => `${definition.name.toLowerCase()}\u0000${definition.file_path}`)) + + while (queue.length > 0 && rows.length < maxNodes) { + const current = queue.shift()! + const remainingNodes = maxNodes - rows.length + const currentTargets = await resolveDefinitionEdgeTargetsForNameAndFiles( + db, + this.root, + current.name, + current.targetFiles, + defaultExportCache + ) + if (currentTargets.length === 0) { + continue + } + + const edgeRows = selectDefinitionEdgeRows(db, currentTargets, ['call'], true, remainingNodes + 1) + + if (edgeRows.length > remainingNodes) { + nodesCapped = true + } + + for (const edgeRow of edgeRows.slice(0, remainingNodes)) { + const nodeName = edgeRow.src_symbol?.trim() || 'top-level' + rows.push({ ...edgeRow, name: nodeName, depth: current.depth }) + + if (current.depth >= depthLimit) { + if (edgeRow.src_symbol) { + depthCapped = true + } + continue + } + + if (!edgeRow.src_symbol) { + continue + } + + const nextKey = `${edgeRow.src_symbol.toLowerCase()}\u0000${edgeRow.src_file}` + if (visitedTargets.has(nextKey)) { + continue + } + + visitedTargets.add(nextKey) + queue.push({ + name: edgeRow.src_symbol, + targetFiles: [edgeRow.src_file], + depth: current.depth + 1 + }) + } + + if (nodesCapped) { + break + } + } + + if (queue.length > 0) { + nodesCapped = true + } + } finally { + releaseDatabase() + } + + const { items: nodes, tokenTruncated } = await readImpactNodesFromRows(this.root, rows, options?.maxTokens) + return { + nodes, + impactTruncated: nodesCapped || tokenTruncated, + depthCapped, + nodesCapped, + depthLimit, + maxNodes + } + } + + private async readFindSymbolRelations(db: Database.Database, definition: SymbolRow): Promise { + const targets = await resolveDefinitionEdgeTargets(this.root, [definition]) + const siteRows = selectDefinitionEdgeRows(db, targets, ['call', 'ref'], true, FIND_SYMBOL_RELATION_MAX_SITES) + const { items: sites } = await readEdgeResultsFromRows(this.root, siteRows, undefined) + const totalSites = countDefinitionEdgeRows(db, targets, ['call', 'ref']) + + const neighborRows = selectSameFileNeighborRows( + db, + definition.file_path, + definition.id, + definition.start_line, + FIND_SYMBOL_RELATION_MAX_NEIGHBORS + ) + const neighbors = neighborRows.map((row) => buildSymbolNeighbor(row)) + const totalNeighbors = countSameFileNeighborRows(db, definition.file_path, definition.id) + const omitted = Math.max(0, totalSites - siteRows.length) + Math.max(0, totalNeighbors - neighbors.length) + + if (sites.length === 0 && neighbors.length === 0 && omitted === 0) { + return undefined + } + + const relations: SymbolRelations = { sites, neighbors } + if (omitted > 0) { + relations.omitted = omitted + } + + return relations + } + + private async findDefinitionEdges( + name: string, + options: FindEdgeOptions | undefined, + edgeKinds: ReadonlyArray, + preferCallsFirst: boolean + ): Promise { + if (!name.trim() || !existsSync(this.indexPath)) { + return [] + } + + const releaseDatabase = await this.enterDatabaseUser() + let rows: EdgeResultRow[] = [] + try { + const db = this.openDatabase() + this.ensureIndexCompatibleForQueries(db) + + const definitions = selectExactDefinitionRows(db, name, normalizeKinds(options?.kind), options?.pathGlob) + if (definitions.length === 0) { + return [] + } + + const targets = await resolveDefinitionEdgeTargets(this.root, definitions) + rows = selectDefinitionEdgeRows(db, targets, edgeKinds, preferCallsFirst) + } finally { + releaseDatabase() + } + + return (await readEdgeResultsFromRows(this.root, rows, options?.maxTokens)).items + } + async readChunk(id: string, options?: ReadChunkOptions): Promise { const releaseDatabase = await this.enterDatabaseUser() let parsedChunkId: { file: string; startLine: number; endLine: number } | null @@ -1210,6 +1496,22 @@ export class SqliteRepo implements Repo { create index if not exists idx_symbols_name on symbols(name); create index if not exists idx_symbols_kind on symbols(kind); + create table if not exists edges( + id integer primary key autoincrement, + src_file text not null references files(path) on delete cascade, + src_line integer not null, + src_symbol text, + dst_name text not null, + dst_file text, + edge_kind text not null, + resolution text not null, + language text + ); + + create index if not exists idx_edges_dst on edges(dst_name, dst_file); + create index if not exists idx_edges_src on edges(src_file); + create index if not exists idx_edges_kind on edges(edge_kind); + create table if not exists embedding_cache( provider_id text not null, provider_dims integer not null, @@ -1228,6 +1530,7 @@ export class SqliteRepo implements Repo { private clearIndex(db: Database.Database): void { db.exec(` delete from symbols; + delete from edges; delete from chunks_fts; delete from chunks; delete from files; @@ -1372,6 +1675,12 @@ function applySyncChanges(db: Database.Database, context: SyncApplyContext): voi values (?, ?, ?, ?, ?, ?, ?, ?) ` ) + const insertEdge = db.prepare( + ` + insert into edges(src_file, src_line, src_symbol, dst_name, dst_file, edge_kind, resolution, language) + values (?, ?, ?, ?, ?, ?, ?, ?) + ` + ) const upsertEmbeddingCache = db.prepare( ` insert into embedding_cache(provider_id, provider_dims, model_version, content_hash, embedding, updated_at) @@ -1389,9 +1698,12 @@ function applySyncChanges(db: Database.Database, context: SyncApplyContext): voi const selectChunkIdsByFile = db.prepare<[string], { id: string }>('select id from chunks where file_path = ?') const deleteChunkFts = db.prepare<[string]>('delete from chunks_fts where chunk_id = ?') const deleteSymbolsByFile = db.prepare<[string]>('delete from symbols where file_path = ?') + const deleteEdgesBySrcFile = db.prepare<[string]>('delete from edges where src_file = ?') + const deleteEdgesByDstFile = db.prepare<[string]>('delete from edges where dst_file = ?') const deleteChunksByFile = db.prepare<[string]>('delete from chunks where file_path = ?') const deleteFile = db.prepare<[string]>('delete from files where path = ?') const replacementPaths = [...context.diff.removedPaths, ...context.diff.changedFiles.map((file) => file.relativePath)] + const symbolRangesByFile = groupSymbolRangesByFile(context.embeddedChunks) db.transaction(() => { for (const filePath of replacementPaths) { @@ -1400,16 +1712,34 @@ function applySyncChanges(db: Database.Database, context: SyncApplyContext): voi } deleteSymbolsByFile.run(filePath) + deleteEdgesBySrcFile.run(filePath) deleteChunksByFile.run(filePath) deleteFile.run(filePath) } + for (const removedPath of context.diff.removedPaths) { + deleteEdgesByDstFile.run(removedPath) + } + for (const file of context.diff.touchedFiles) { updateFileManifest.run(fileToManifestRow(file)) } for (const file of context.diff.changedFiles) { insertFile.run(fileToManifestRow(file)) + + for (const edge of extractEdges(file, symbolRangesByFile.get(file.relativePath) ?? [])) { + insertEdge.run( + edge.srcFile, + edge.srcLine, + edge.srcSymbol ?? null, + edge.dstName, + edge.dstFile ?? null, + edge.edgeKind, + edge.resolution, + edge.language ?? null + ) + } } for (const { row, embedding } of context.embeddedChunks) { @@ -1474,6 +1804,43 @@ function applySyncChanges(db: Database.Database, context: SyncApplyContext): voi })() } +function extractEdges(file: ScannedFile, symbolRanges: SourceSymbolRange[]): Edge[] { + if (isTypeScriptLike(file.language)) { + return extractTypeScriptEdges(file, symbolRanges) + } + + if (isPythonLike(file.language)) { + return extractPythonEdges(file, symbolRanges) + } + + if (isGoLike(file.language) || isJavaLike(file.language) || isRubyLike(file.language) || isRustLike(file.language)) { + return extractNameOnlyCallEdges(file, symbolRanges) + } + + return [] +} + +function groupSymbolRangesByFile(records: EmbeddedChunkRecord[]): Map { + const rangesByFile = new Map() + + for (const { row } of records) { + if (!row.symbol || !row.kind || row.kind === 'file') { + continue + } + + const ranges = rangesByFile.get(row.file) ?? [] + ranges.push({ + startLine: row.startLine, + endLine: row.endLine, + symbol: row.symbol, + kind: row.kind + }) + rangesByFile.set(row.file, ranges) + } + + return rangesByFile +} + function fileToManifestRow(file: ScannedFile): IndexedFileRow { return { path: file.relativePath, @@ -2149,6 +2516,22 @@ function normalizeSearchTokenBudget(value: number | undefined): number | undefin return Math.floor(value) } +function normalizeImpactDepth(value: number | undefined): number { + if (value === undefined || !Number.isFinite(value) || value < 0) { + return DEFAULT_IMPACT_DEPTH + } + + return Math.min(Math.floor(value), HARD_MAX_IMPACT_BOUND) +} + +function normalizeImpactMaxNodes(value: number | undefined): number { + if (value === undefined || !Number.isFinite(value) || value <= 0) { + return DEFAULT_IMPACT_MAX_NODES + } + + return Math.min(Math.floor(value), HARD_MAX_IMPACT_BOUND) +} + function estimateSearchHitTokens(row: ChunkRow, snippet: string): number { const header = `${row.file_path}:${row.start_line}-${row.end_line} ${row.symbol ?? ''} ${row.kind ?? ''}` // The compact renderer emits up to COMPACT_SNIPPET_MAX_LINES lines, each with a @@ -2359,171 +2742,334 @@ async function attachUsagesToTopDefinitionHit( } } +const MAX_ATTACHED_USAGES = 5 + async function findImportResolvedUsages( db: Database.Database, root: string, definition: SearchHit, options?: SearchOptions ): Promise { - if (!definition.symbol || !definition.language) { + if ( + !definition.symbol || + !definition.language || + (!isTypeScriptLike(definition.language) && !isPythonLike(definition.language)) + ) { return [] } - const candidateFiles = selectUsageCandidateFiles(db, definition.language, options?.pathGlob) - if (candidateFiles.length === 0) { + const targets = await resolveDefinitionEdgeTargets(root, [ + { + name: definition.symbol, + file_path: definition.file, + language: definition.language + } + ]) + if (targets.length === 0) { return [] } - if (isTypeScriptLike(definition.language)) { - return findTypeScriptUsages(root, definition, candidateFiles) - } - - if (isPythonLike(definition.language)) { - return findPythonUsages(root, definition, candidateFiles) - } + const rows = selectDefinitionUsageRows(db, targets, { + excludedSrcFile: definition.file, + excludedStartLine: definition.range.startLine, + excludedEndLine: definition.range.endLine, + ...(options?.pathGlob ? { pathGlob: options.pathGlob } : {}), + limit: MAX_ATTACHED_USAGES + }) - return [] + return readUsagesFromEdgeRows(root, rows) } -function selectUsageCandidateFiles(db: Database.Database, language: string, pathGlob?: string): Array<{ path: string; language: string }> { - const family = isTypeScriptLike(language) - ? ['typescript', 'tsx', 'javascript', 'jsx'] - : isPythonLike(language) - ? ['python'] - : [language] +async function readUsagesFromEdgeRows(root: string, rows: EdgeUsageRow[]): Promise { + const linesByFile = new Map() + const usages: SymbolUsage[] = [] + + for (const row of rows) { + let lines = linesByFile.get(row.src_file) + if (!lines) { + try { + lines = splitLines(await readFile(resolve(root, row.src_file), 'utf8')) + linesByFile.set(row.src_file, lines) + } catch { + continue + } + } - const whereClauses = [`language in (${family.map(() => '?').join(', ')})`] - const params: string[] = [...family] + const usage: SymbolUsage = { + file: row.src_file, + range: { startLine: row.src_line, endLine: row.src_line }, + line: row.src_line, + snippet: (lines[row.src_line - 1] ?? '').trimEnd(), + resolution: row.resolution + } - if (pathGlob) { - whereClauses.push('codesift_minimatch(path, ?) = 1') - params.push(pathGlob) + if (row.language) { + usage.language = row.language + } + + usages.push(usage) } - return db - .prepare( - ` - select path, language - from files - where ${whereClauses.join(' and ')} - order by path asc - ` - ) - .all(...params) + return usages } -async function findTypeScriptUsages( - root: string, - definition: SearchHit, - candidateFiles: Array<{ path: string; language: string }> -): Promise { - const usages: SymbolUsage[] = [] - const seen = new Set() +function extractTypeScriptEdges(file: ScannedFile, symbolRanges: SourceSymbolRange[]): Edge[] { + const sourceFile = + getCachedTypeScriptSourceFile(file) ?? + ts.createSourceFile(file.relativePath, file.content, ts.ScriptTarget.Latest, true, scriptKindFromPath(file.relativePath)) + const directBindings = new Map() + const namespaceBindings = new Map() + const sameFileSymbols = collectSameFileSymbols(symbolRanges) + const scopeDeclarations = collectTypeScriptUsageScopeDeclarations(sourceFile) + const scopeStack: Array> = [] + const edgesByKey = new Map() - for (const candidate of candidateFiles) { - if (usages.length >= 5) { - break + for (const statement of sourceFile.statements) { + if (!ts.isImportDeclaration(statement)) { + continue } - const absolutePath = resolve(root, candidate.path) - let content: string - try { - content = await readFile(absolutePath, 'utf8') - } catch { + const moduleSpecifier = ts.isStringLiteral(statement.moduleSpecifier) ? statement.moduleSpecifier.text : null + const dstFile = moduleSpecifier ? resolveTypeScriptModuleTarget(file, moduleSpecifier) : null + const clause = statement.importClause + if (!clause) { continue } - const sourceFile = ts.createSourceFile(candidate.path, content, ts.ScriptTarget.Latest, true, scriptKindFromPath(candidate.path)) - const directNames = new Set() - const namespaceNames = new Set() - - if (candidate.path === definition.file) { - directNames.add(definition.symbol!) + if (clause.name && dstFile) { + directBindings.set(clause.name.text, { dstName: 'default', dstFile, resolution: 'import-resolved' }) + pushEdge( + edgesByKey, + buildEdge({ + srcFile: file.relativePath, + srcLine: lineNumberForOffsetInSourceFile(sourceFile, clause.name.getStart(sourceFile)), + srcSymbol: selectEnclosingSourceSymbol(symbolRanges, lineNumberForOffsetInSourceFile(sourceFile, clause.name.getStart(sourceFile))), + dstName: 'default', + dstFile, + edgeKind: 'import', + resolution: 'import-resolved', + language: file.language + }) + ) } - for (const statement of sourceFile.statements) { - if (!ts.isImportDeclaration(statement)) { - continue - } - - const moduleSpecifier = ts.isStringLiteral(statement.moduleSpecifier) ? statement.moduleSpecifier.text : null - if (!moduleSpecifier || !moduleSpecifierCouldResolveToFile(candidate.path, moduleSpecifier, definition.file)) { - continue - } - - const clause = statement.importClause - if (!clause) { - continue - } + const namedBindings = clause.namedBindings + if (!namedBindings) { + continue + } - if (clause.name && definition.symbol === 'default') { - directNames.add(clause.name.text) - } + if (ts.isNamedImports(namedBindings)) { + for (const element of namedBindings.elements) { + const importedName = (element.propertyName ?? element.name).text + if (!dstFile) { + continue + } - const namedBindings = clause.namedBindings - if (!namedBindings) { - continue + directBindings.set(element.name.text, { + dstName: importedName, + dstFile, + resolution: 'import-resolved' + }) + const srcLine = lineNumberForOffsetInSourceFile(sourceFile, element.name.getStart(sourceFile)) + pushEdge( + edgesByKey, + buildEdge({ + srcFile: file.relativePath, + srcLine, + srcSymbol: selectEnclosingSourceSymbol(symbolRanges, srcLine), + dstName: importedName, + dstFile, + edgeKind: 'import', + resolution: 'import-resolved', + language: file.language + }) + ) } + continue + } - if (ts.isNamedImports(namedBindings)) { - for (const element of namedBindings.elements) { - const importedName = (element.propertyName ?? element.name).text - if (importedName === definition.symbol) { - directNames.add(element.name.text) - } - } - } else if (ts.isNamespaceImport(namedBindings)) { - namespaceNames.add(namedBindings.name.text) - } + if (ts.isNamespaceImport(namedBindings) && dstFile) { + namespaceBindings.set(namedBindings.name.text, { + dstName: '*', + dstFile, + resolution: 'import-resolved' + }) + const srcLine = lineNumberForOffsetInSourceFile(sourceFile, namedBindings.name.getStart(sourceFile)) + pushEdge( + edgesByKey, + buildEdge({ + srcFile: file.relativePath, + srcLine, + srcSymbol: selectEnclosingSourceSymbol(symbolRanges, srcLine), + dstName: '*', + dstFile, + edgeKind: 'import', + resolution: 'import-resolved', + language: file.language + }) + ) } + } - if (directNames.size === 0 && namespaceNames.size === 0) { + for (const statement of sourceFile.statements) { + if (!ts.isClassDeclaration(statement) && !ts.isInterfaceDeclaration(statement)) { continue } - const lines = splitLines(content) - const lineStarts = buildLineStarts(content) - const recordUsage = (offset: number, lineNumberOverride?: number) => { - const lineNumber = lineNumberOverride ?? lineNumberForOffset(lineStarts, offset) - if (candidate.path === definition.file && lineNumber >= definition.range.startLine && lineNumber <= definition.range.endLine) { - return + for (const clause of statement.heritageClauses ?? []) { + const edgeKind = clause.token === ts.SyntaxKind.ImplementsKeyword + ? 'implements' + : clause.token === ts.SyntaxKind.ExtendsKeyword + ? 'extends' + : null + if (!edgeKind) { + continue } - const key = `${candidate.path}:${lineNumber}` - if (seen.has(key)) { - return + for (const typeNode of clause.types) { + const binding = resolveTypeScriptHeritageTarget(typeNode.expression, directBindings, namespaceBindings, sameFileSymbols, file) + if (!binding) { + continue + } + + const srcLine = lineNumberForOffsetInSourceFile(sourceFile, typeNode.expression.getStart(sourceFile)) + pushEdge( + edgesByKey, + buildEdge({ + srcFile: file.relativePath, + srcLine, + srcSymbol: selectEnclosingSourceSymbol(symbolRanges, srcLine), + dstName: binding.dstName, + ...(binding.dstFile ? { dstFile: binding.dstFile } : {}), + edgeKind, + resolution: binding.resolution, + language: file.language + }) + ) } - seen.add(key) + } + } - usages.push({ - file: candidate.path, - range: { startLine: lineNumber, endLine: lineNumber }, - line: lineNumber, - snippet: (lines[lineNumber - 1] ?? '').trimEnd(), - language: candidate.language, - resolution: 'import-resolved' - }) + const visit = (node: ts.Node): void => { + const declarations = scopeDeclarations.get(node) + if (declarations) { + scopeStack.push(declarations) } - const visit = (node: ts.Node): void => { - if (ts.isPropertyAccessExpression(node) && ts.isIdentifier(node.expression) && namespaceNames.has(node.expression.text) && node.name.text === definition.symbol) { - recordUsage(node.name.getStart(sourceFile)) + try { + if ((ts.isPropertyAccessExpression(node) || ts.isPropertyAccessChain(node)) && ts.isIdentifier(node.expression)) { + const binding = namespaceBindings.get(node.expression.text) + if (binding?.dstFile) { + const srcLine = lineNumberForOffsetInSourceFile(sourceFile, node.name.getStart(sourceFile)) + pushEdge( + edgesByKey, + buildEdge({ + srcFile: file.relativePath, + srcLine, + srcSymbol: selectEnclosingSourceSymbol(symbolRanges, srcLine), + dstName: node.name.text, + dstFile: binding.dstFile, + edgeKind: isTypeScriptCallTarget(node) ? 'call' : 'ref', + resolution: binding.resolution, + language: file.language + }) + ) + } } - if (ts.isIdentifier(node) && directNames.has(node.text) && isUsageIdentifier(node)) { - recordUsage(node.getStart(sourceFile)) + if (ts.isIdentifier(node) && isUsageIdentifier(node)) { + const binding = directBindings.get(node.text) + const srcLine = lineNumberForOffsetInSourceFile(sourceFile, node.getStart(sourceFile)) + if (binding?.dstFile) { + pushEdge( + edgesByKey, + buildEdge({ + srcFile: file.relativePath, + srcLine, + srcSymbol: selectEnclosingSourceSymbol(symbolRanges, srcLine), + dstName: binding.dstName, + dstFile: binding.dstFile, + edgeKind: isTypeScriptCallTarget(node) ? 'call' : 'ref', + resolution: binding.resolution, + language: file.language + }) + ) + } else if (sameFileSymbols.has(node.text) && !isTypeScriptSameFileReferenceShadowed(node.text, scopeStack)) { + pushEdge( + edgesByKey, + buildEdge({ + srcFile: file.relativePath, + srcLine, + srcSymbol: selectEnclosingSourceSymbol(symbolRanges, srcLine), + dstName: node.text, + dstFile: file.relativePath, + edgeKind: isTypeScriptCallTarget(node) ? 'call' : 'ref', + resolution: 'import-resolved', + language: file.language + }) + ) + } } ts.forEachChild(node, visit) + } finally { + if (declarations) { + scopeStack.pop() + } } - - visit(sourceFile) } - return usages.slice(0, 5) + visit(sourceFile) + return [...edgesByKey.values()] } -function isUsageIdentifier(node: ts.Identifier): boolean { +function resolveTypeScriptHeritageTarget( + expression: ts.Expression, + directBindings: Map, + namespaceBindings: Map, + sameFileSymbols: Set, + file: ScannedFile +): EdgeBinding | undefined { + if (ts.isIdentifier(expression)) { + const binding = directBindings.get(expression.text) + if (binding?.dstFile) { + return binding + } + + if (sameFileSymbols.has(expression.text)) { + return { + dstName: expression.text, + dstFile: file.relativePath, + resolution: 'import-resolved' + } + } + + return { + dstName: expression.text, + resolution: 'name-only' + } + } + + if ((ts.isPropertyAccessExpression(expression) || ts.isPropertyAccessChain(expression)) && ts.isIdentifier(expression.expression)) { + const binding = namespaceBindings.get(expression.expression.text) + if (binding?.dstFile) { + return { + dstName: expression.name.text, + dstFile: binding.dstFile, + resolution: binding.resolution + } + } + + return { + dstName: expression.name.text, + resolution: 'name-only' + } + } + + return undefined +} + +function isUsageIdentifier(node: ts.Identifier): boolean { const parent = node.parent if (!parent) { return true @@ -2541,7 +3087,7 @@ function isUsageIdentifier(node: ts.Identifier): boolean { return false } - if (ts.isPropertyAccessExpression(parent) && parent.name === node) { + if ((ts.isPropertyAccessExpression(parent) || ts.isPropertyAccessChain(parent)) && parent.name === node) { return false } @@ -2556,115 +3102,515 @@ function isUsageIdentifier(node: ts.Identifier): boolean { return true } -async function findPythonUsages( - root: string, - definition: SearchHit, - candidateFiles: Array<{ path: string; language: string }> -): Promise { - const usages: SymbolUsage[] = [] - const seen = new Set() +function isTypeScriptCallTarget(node: ts.Node): boolean { + const parent = node.parent + if (!parent) { + return false + } - for (const candidate of candidateFiles) { - if (usages.length >= 5) { - break - } + return ( + ((ts.isCallExpression(parent) || ts.isCallChain(parent)) && parent.expression === node) || + (ts.isNewExpression(parent) && parent.expression === node) || + (ts.isTaggedTemplateExpression(parent) && parent.tag === node) + ) +} - const absolutePath = resolve(root, candidate.path) - let content: string - try { - content = await readFile(absolutePath, 'utf8') - } catch { +function extractPythonEdges(file: ScannedFile, symbolRanges: SourceSymbolRange[]): Edge[] { + const lines = splitLines(file.content) + const maskedLines = maskPythonLines(lines) + const edgesByKey = new Map() + const directBindings = new Map() + const moduleBindings = new Map() + + for (const symbol of collectSameFileSymbols(symbolRanges)) { + directBindings.set(symbol, { + dstName: symbol, + dstFile: file.relativePath, + resolution: 'import-resolved' + }) + } + + for (let index = 0; index < maskedLines.length; index += 1) { + const lineNumber = index + 1 + const trimmed = maskedLines[index]?.trim() ?? '' + if (!trimmed) { continue } - const lines = splitLines(content) - const maskedLines = maskPythonLines(lines) - const directNames = new Set() - const moduleAliases = new Set() + const fromImport = trimmed.match(/^from\s+([.A-Za-z0-9_]+)\s+import\s+(.+)$/) + if (fromImport) { + const dstFile = resolvePythonModuleTarget(file, fromImport[1]!) + for (const part of fromImport[2]!.split(',')) { + const match = part.trim().match(/^([A-Za-z_][A-Za-z0-9_]*)(?:\s+as\s+([A-Za-z_][A-Za-z0-9_]*))?$/) + const importedName = match?.[1] + if (!importedName) { + continue + } - if (candidate.path === definition.file) { - directNames.add(definition.symbol!) + directBindings.set(match?.[2] ?? importedName, { + dstName: importedName, + ...(dstFile ? { dstFile } : {}), + resolution: dstFile ? 'import-resolved' : 'name-only' + }) + } + continue } - for (let index = 0; index < lines.length; index += 1) { - const rawLine = lines[index] ?? '' - const trimmed = rawLine.trim() - if (!trimmed) { - continue + const importMatch = trimmed.match(/^import\s+([.A-Za-z0-9_]+)(?:\s+as\s+([A-Za-z_][A-Za-z0-9_]*))?$/) + if (importMatch) { + const dstFile = resolvePythonModuleTarget(file, importMatch[1]!) + // Known limitation: `import a.b.c` still keys the top-level alias only; deeper dotted-module binding stays approximate noise. + const localName = importMatch[2] ?? importMatch[1]!.split('.')[0] ?? '' + if (localName) { + moduleBindings.set(localName, { + dstName: '*', + ...(dstFile ? { dstFile } : {}), + resolution: dstFile ? 'import-resolved' : 'name-only' + }) } + } + } - const fromImport = trimmed.match(/^from\s+([.A-Za-z0-9_]+)\s+import\s+(.+)$/) - if (fromImport && pythonModuleCouldResolveToFile(candidate.path, fromImport[1]!, definition.file)) { - for (const part of fromImport[2]!.split(',')) { - const match = part.trim().match(/^([A-Za-z_][A-Za-z0-9_]*)(?:\s+as\s+([A-Za-z_][A-Za-z0-9_]*))?$/) - const importedName = match?.[1] - if (importedName && importedName === definition.symbol) { - directNames.add(match?.[2] ?? importedName) - } + for (let index = 0; index < maskedLines.length; index += 1) { + const lineNumber = index + 1 + const masked = maskedLines[index] ?? '' + const trimmed = masked.trim() + if (!trimmed || /^(from\s+.+\s+import\s+.+|import\s+.+|class\s+.+|def\s+.+)$/.test(trimmed)) { + continue + } + + for (const [alias, binding] of moduleBindings) { + const pattern = new RegExp(`\\b${escapeRegExp(alias)}\\.([A-Za-z_][A-Za-z0-9_]*)\\b`, 'g') + for (const match of masked.matchAll(pattern)) { + const propertyName = match[1] + const matchText = match[0] + const matchIndex = match.index ?? -1 + if (!propertyName || matchIndex < 0) { + continue } + + const edgeKind = /^\s*\(/.test(masked.slice(matchIndex + matchText.length)) ? 'call' : 'ref' + pushEdge( + edgesByKey, + buildEdge({ + srcFile: file.relativePath, + srcLine: lineNumber, + srcSymbol: selectEnclosingSourceSymbol(symbolRanges, lineNumber), + dstName: propertyName, + ...(binding.dstFile ? { dstFile: binding.dstFile } : {}), + edgeKind, + resolution: binding.resolution, + language: file.language + }) + ) } + } + + for (const [localName, binding] of directBindings) { + const pattern = new RegExp(`\\b${escapeRegExp(localName)}\\b`, 'g') + for (const match of masked.matchAll(pattern)) { + const matchText = match[0] + const matchIndex = match.index ?? -1 + if (!matchText || matchIndex < 0) { + continue + } + + if (matchIndex > 0 && masked[matchIndex - 1] === '.') { + continue + } - const importAlias = trimmed.match(/^import\s+([.A-Za-z0-9_]+)\s+as\s+([A-Za-z_][A-Za-z0-9_]*)$/) - if (importAlias && pythonModuleCouldResolveToFile(candidate.path, importAlias[1]!, definition.file)) { - moduleAliases.add(importAlias[2]!) + const edgeKind = /^\s*\(/.test(masked.slice(matchIndex + matchText.length)) ? 'call' : 'ref' + pushEdge( + edgesByKey, + buildEdge({ + srcFile: file.relativePath, + srcLine: lineNumber, + srcSymbol: selectEnclosingSourceSymbol(symbolRanges, lineNumber), + dstName: binding.dstName, + ...(binding.dstFile ? { dstFile: binding.dstFile } : {}), + edgeKind, + resolution: binding.resolution, + language: file.language + }) + ) } } + } + + return [...edgesByKey.values()] +} + +function extractNameOnlyCallEdges(file: ScannedFile, symbolRanges: SourceSymbolRange[]): Edge[] { + const lines = splitLines(file.content) + const maskedLines = isRubyLike(file.language) ? maskRubySyntaxLines(lines) : maskCStyleSyntax(lines) + const callPattern = callPatternForNameOnlyLanguage(file.language) + const edgesByKey = new Map() + + for (let index = 0; index < maskedLines.length; index += 1) { + const lineNumber = index + 1 + const masked = maskedLines[index] ?? '' + const trimmed = masked.trim() + if (!trimmed || shouldSkipNameOnlyCallLine(file.language, trimmed)) { + continue + } - if (directNames.size === 0 && moduleAliases.size === 0) { + const enclosing = selectEnclosingSourceSymbolRange(symbolRanges, lineNumber) + if (isGoLike(file.language) && enclosing?.kind && !['function', 'method', 'variable', 'constant'].includes(enclosing.kind)) { continue } - for (let index = 0; index < maskedLines.length; index += 1) { - const lineNumber = index + 1 - if (candidate.path === definition.file && lineNumber >= definition.range.startLine && lineNumber <= definition.range.endLine) { + for (const match of masked.matchAll(callPattern)) { + const callee = match[1] + const matchIndex = match.index ?? -1 + if (!callee || matchIndex < 0 || isNameOnlyCallKeyword(file.language, callee)) { continue } - const rawLine = lines[index]?.trim() ?? '' - if (/^(from\s+.+\s+import\s+.+|import\s+.+)$/.test(rawLine)) { + if ( + enclosing?.symbol === callee && + enclosing.startLine === lineNumber && + (enclosing.kind === 'function' || enclosing.kind === 'method') + ) { continue } - const masked = maskedLines[index] ?? '' - const directMatch = [...directNames].some((name) => new RegExp(`\\b${escapeRegExp(name)}\\b`).test(masked)) - const moduleMatch = [...moduleAliases].some((alias) => new RegExp(`\\b${escapeRegExp(alias)}\\.${escapeRegExp(definition.symbol!)}\\b`).test(masked)) - if (!directMatch && !moduleMatch) { + pushEdge( + edgesByKey, + buildEdge({ + srcFile: file.relativePath, + srcLine: lineNumber, + srcSymbol: enclosing?.symbol, + dstName: callee, + edgeKind: 'call', + resolution: 'name-only', + language: file.language + }) + ) + } + } + + return [...edgesByKey.values()] +} + +function callPatternForNameOnlyLanguage(language: string): RegExp { + if (isJavaLike(language)) { + return /([A-Za-z_$][A-Za-z0-9_$]*)\s*\(/g + } + + if (isRubyLike(language)) { + return /([A-Za-z_][A-Za-z0-9_!?=]*)\s*\(/g + } + + return /([A-Za-z_][A-Za-z0-9_]*)\s*\(/g +} + +function shouldSkipNameOnlyCallLine(language: string, trimmed: string): boolean { + if (isGoLike(language)) { + return /^(?:package|import)\b/.test(trimmed) + } + + if (isJavaLike(language)) { + return /^(?:package|import)\b/.test(trimmed) || trimmed.startsWith('@') + } + + if (isRubyLike(language)) { + return /^(?:class|module|def)\b/.test(trimmed) + } + + if (isRustLike(language)) { + return /^(?:use|mod)\b/.test(trimmed) || trimmed.startsWith('#[') + } + + return false +} + +function isNameOnlyCallKeyword(language: string, callee: string): boolean { + if (isGoLike(language)) { + return GO_NAME_ONLY_CALL_KEYWORDS.has(callee) + } + + if (isJavaLike(language)) { + return JAVA_NAME_ONLY_CALL_KEYWORDS.has(callee) + } + + if (isRubyLike(language)) { + return RUBY_NAME_ONLY_CALL_KEYWORDS.has(callee) + } + + if (isRustLike(language)) { + return RUST_NAME_ONLY_CALL_KEYWORDS.has(callee) + } + + return false +} + +const GO_NAME_ONLY_CALL_KEYWORDS = new Set(['if', 'for', 'switch', 'select']) +const JAVA_NAME_ONLY_CALL_KEYWORDS = new Set(['if', 'for', 'while', 'switch', 'catch', 'try', 'synchronized', 'do']) +const RUBY_NAME_ONLY_CALL_KEYWORDS = new Set(['if', 'unless', 'while', 'until', 'for', 'case']) +const RUST_NAME_ONLY_CALL_KEYWORDS = new Set(['if', 'for', 'while', 'match', 'loop']) + +// Known limitation: Ruby heredocs, `%w`, and `=begin` blocks are only partially masked here; downstream Ruby edges stay name-only/approximate. +function maskRubySyntaxLines(lines: string[]): string[] { + const maskedLines: string[] = [] + + for (const line of lines) { + let masked = '' + let stringQuote: '"' | "'" | null = null + let escaped = false + + for (let index = 0; index < line.length; index += 1) { + const char = line[index]! + + if (stringQuote) { + if (escaped) { + escaped = false + masked += ' ' + continue + } + + if (char === '\\') { + escaped = true + masked += ' ' + continue + } + + if (char === stringQuote) { + stringQuote = null + } + + masked += ' ' continue } - const key = `${candidate.path}:${lineNumber}` - if (seen.has(key)) { + if (char === '#') { + masked += ' '.repeat(line.length - index) + break + } + + if (char === '"' || char === "'") { + stringQuote = char + masked += ' ' continue } - seen.add(key) - usages.push({ - file: candidate.path, - range: { startLine: lineNumber, endLine: lineNumber }, - line: lineNumber, - snippet: (lines[index] ?? '').trimEnd(), - language: candidate.language, - resolution: 'import-resolved' - }) + masked += char + } - if (usages.length >= 5) { - break + maskedLines.push(masked) + } + + return maskedLines +} + +function collectSameFileSymbols(symbolRanges: SourceSymbolRange[]): Set { + return new Set(symbolRanges.flatMap((range) => (range.symbol ? [range.symbol] : []))) +} + +function collectTypeScriptUsageScopeDeclarations(sourceFile: ts.SourceFile): Map> { + const declarationsByScope = new Map>() + + const addToScope = (scope: ts.Node | undefined, name: string | undefined): void => { + if (!scope || !name || ts.isSourceFile(scope)) { + return + } + + const declarations = declarationsByScope.get(scope) ?? new Set() + declarations.add(name) + declarationsByScope.set(scope, declarations) + } + + const visit = (node: ts.Node): void => { + if (ts.isParameter(node)) { + for (const name of collectTypeScriptBindingNames(node.name)) { + addToScope(findNearestTypeScriptFunctionScope(node), name) + } + } else if (ts.isVariableDeclaration(node)) { + const declarationList = ts.isVariableDeclarationList(node.parent) ? node.parent : undefined + const scope = declarationList && (declarationList.flags & ts.NodeFlags.BlockScoped) + ? findNearestTypeScriptBlockScope(node) + : findNearestTypeScriptFunctionScope(node) + for (const name of collectTypeScriptBindingNames(node.name)) { + addToScope(scope, name) } + } else if (ts.isFunctionDeclaration(node) && node.name) { + addToScope(findNearestTypeScriptBlockScope(node), node.name.text) + } else if (ts.isClassDeclaration(node) && node.name) { + addToScope(findNearestTypeScriptBlockScope(node), node.name.text) + } else if (ts.isInterfaceDeclaration(node) || ts.isTypeAliasDeclaration(node) || ts.isEnumDeclaration(node)) { + addToScope(findNearestTypeScriptBlockScope(node), node.name.text) + } else if (ts.isFunctionExpression(node) && node.name) { + addToScope(node, node.name.text) + } + + ts.forEachChild(node, visit) + } + + visit(sourceFile) + return declarationsByScope +} + +function collectTypeScriptBindingNames(name: ts.BindingName): string[] { + if (ts.isIdentifier(name)) { + return [name.text] + } + + const names: string[] = [] + for (const element of name.elements) { + if (ts.isOmittedExpression(element)) { + continue + } + + names.push(...collectTypeScriptBindingNames(element.name)) + } + + return names +} + +function findNearestTypeScriptFunctionScope(node: ts.Node): ts.Node | undefined { + let current: ts.Node | undefined = node.parent + while (current) { + if (ts.isFunctionLike(current) || ts.isSourceFile(current)) { + return current + } + current = current.parent + } + + return undefined +} + +function findNearestTypeScriptBlockScope(node: ts.Node): ts.Node | undefined { + let current: ts.Node | undefined = node.parent + while (current) { + if ( + ts.isBlock(current) || + ts.isModuleBlock(current) || + ts.isCaseBlock(current) || + ts.isForStatement(current) || + ts.isForInStatement(current) || + ts.isForOfStatement(current) || + ts.isCatchClause(current) || + ts.isFunctionLike(current) || + ts.isSourceFile(current) + ) { + return current + } + current = current.parent + } + + return undefined +} + +function isTypeScriptSameFileReferenceShadowed(name: string, scopeStack: ReadonlyArray>): boolean { + return scopeStack.some((scope) => scope.has(name)) +} + +function selectEnclosingSourceSymbolRange(symbolRanges: SourceSymbolRange[], lineNumber: number): SourceSymbolRange | undefined { + let best: SourceSymbolRange | undefined + + for (const range of symbolRanges) { + if (!range.symbol || lineNumber < range.startLine || lineNumber > range.endLine) { + continue + } + + if (!best) { + best = range + continue + } + + const span = range.endLine - range.startLine + const bestSpan = best.endLine - best.startLine + if (span < bestSpan || (span === bestSpan && range.startLine >= best.startLine)) { + best = range + } + } + + return best +} + +function selectEnclosingSourceSymbol(symbolRanges: SourceSymbolRange[], lineNumber: number): string | undefined { + return selectEnclosingSourceSymbolRange(symbolRanges, lineNumber)?.symbol +} + +function buildEdge(params: { + srcFile: string + srcLine: number + srcSymbol: string | undefined + dstName: string + dstFile?: string + edgeKind: Edge['edgeKind'] + resolution: Edge['resolution'] + language?: string +}): Edge { + const edge: Edge = { + srcFile: params.srcFile, + srcLine: params.srcLine, + dstName: params.dstName, + edgeKind: params.edgeKind, + resolution: params.resolution + } + + if (params.srcSymbol) { + edge.srcSymbol = params.srcSymbol + } + + if (params.dstFile) { + edge.dstFile = params.dstFile + } + + if (params.language) { + edge.language = params.language + } + + return edge +} + +function pushEdge(edgesByKey: Map, edge: Edge): void { + const bucket = edge.edgeKind === 'import' + ? 'import' + : edge.edgeKind === 'call' || edge.edgeKind === 'ref' + ? 'usage' + : edge.edgeKind + const key = [bucket, edge.srcFile, edge.srcLine, edge.dstName, edge.dstFile ?? ''].join('\0') + const existing = edgesByKey.get(key) + if (!existing) { + edgesByKey.set(key, edge) + return + } + + if (existing.edgeKind === 'ref' && edge.edgeKind === 'call') { + existing.edgeKind = 'call' + } + + if (existing.resolution === 'name-only' && edge.resolution === 'import-resolved') { + existing.resolution = 'import-resolved' + if (edge.dstFile) { + existing.dstFile = edge.dstFile } } +} - return usages.slice(0, 5) +function lineNumberForOffsetInSourceFile(sourceFile: ts.SourceFile, offset: number): number { + return sourceFile.getLineAndCharacterOfPosition(offset).line + 1 } -function moduleSpecifierCouldResolveToFile(importerFile: string, moduleSpecifier: string, targetFile: string): boolean { +function resolveTypeScriptModuleTarget(file: ScannedFile, moduleSpecifier: string): string | null { if (!moduleSpecifier.startsWith('.')) { - return false + return null } + for (const candidate of buildTypeScriptModuleCandidates(file.relativePath, moduleSpecifier)) { + if (existsSync(resolve(scannedFileRoot(file), candidate))) { + return candidate + } + } + + return null +} + +function buildTypeScriptModuleCandidates(importerFile: string, moduleSpecifier: string): string[] { const importerDirectory = dirname(importerFile) - const base = normalizeRepoPath(resolve('/', importerDirectory, moduleSpecifier)) - const target = normalizeRepoPath(resolve('/', targetFile)) - const candidates = new Set([ + const base = normalizeRelativeRepoPath(resolve('/', importerDirectory, moduleSpecifier)) + return [ base, `${base}.ts`, `${base}.tsx`, @@ -2678,14 +3624,21 @@ function moduleSpecifierCouldResolveToFile(importerFile: string, moduleSpecifier `${base}/index.jsx`, `${base}/index.mts`, `${base}/index.cts` - ]) + ] +} + +function resolvePythonModuleTarget(file: ScannedFile, moduleSpecifier: string): string | null { + for (const candidate of buildPythonModuleCandidates(file.relativePath, moduleSpecifier)) { + if (existsSync(resolve(scannedFileRoot(file), candidate))) { + return candidate + } + } - return candidates.has(target) + return null } -function pythonModuleCouldResolveToFile(importerFile: string, moduleSpecifier: string, targetFile: string): boolean { +function buildPythonModuleCandidates(importerFile: string, moduleSpecifier: string): string[] { const importerDirectory = dirname(importerFile) - const target = normalizeRepoPath(resolve('/', targetFile)) const relativeMatch = moduleSpecifier.match(/^(\.+)(.*)$/) let base: string @@ -2693,19 +3646,34 @@ function pythonModuleCouldResolveToFile(importerFile: string, moduleSpecifier: s const dots = relativeMatch[1]!.length const rest = relativeMatch[2] ?? '' const parentDirectory = dots > 1 ? resolve('/', importerDirectory, ...Array.from({ length: dots - 1 }, () => '..')) : resolve('/', importerDirectory) - base = normalizeRepoPath(resolve(parentDirectory, rest.replace(/\./g, '/'))) + base = normalizeRelativeRepoPath(resolve(parentDirectory, rest.replace(/\./g, '/'))) } else { - base = normalizeRepoPath(resolve('/', moduleSpecifier.replace(/\./g, '/'))) + base = normalizeRelativeRepoPath(resolve('/', moduleSpecifier.replace(/\./g, '/'))) + } + + return [base, `${base}.py`, `${base}/__init__.py`] +} + +function scannedFileRoot(file: ScannedFile): string { + const normalizedAbsolute = normalizePath(file.absolutePath) + const normalizedRelative = normalizeRelativeRepoPath(file.relativePath) + const suffix = normalizedRelative ? `/${normalizedRelative}` : '' + + if (suffix && normalizedAbsolute.endsWith(suffix)) { + return normalizedAbsolute.slice(0, -suffix.length) || '/' } - const candidates = new Set([base, `${base}.py`, `${base}/__init__.py`]) - return candidates.has(target) + return dirname(normalizedAbsolute) } function normalizeRepoPath(value: string): string { return value.replace(/\\/g, '/').replace(/\/+/g, '/').replace(/\/$/, '') } +function normalizeRelativeRepoPath(value: string): string { + return normalizeRepoPath(value).replace(/^\/+/, '') +} + function scriptKindFromPath(filePath: string): ts.ScriptKind { if (filePath.endsWith('.ts')) { return ts.ScriptKind.TS @@ -2927,6 +3895,480 @@ function selectPartialSymbolRows(db: Database.Database, name: string, kinds: Sym .all(...params, name) } +function selectExactDefinitionRows( + db: Database.Database, + name: string, + kinds: SymbolKind[], + pathGlob?: string +): SymbolRow[] { + const whereClauses = ['lower(name) = lower(?)'] + const params: string[] = [name] + + if (kinds.length > 0) { + whereClauses.push(`kind in (${kinds.map(() => '?').join(', ')})`) + params.push(...kinds) + } + + if (pathGlob) { + whereClauses.push('codesift_minimatch(file_path, ?) = 1') + params.push(pathGlob) + } + + return db + .prepare( + ` + select id, name, file_path, start_line, end_line, kind, signature, parent, language + from symbols + where ${whereClauses.join(' and ')} + order by file_path asc, start_line asc, id asc + ` + ) + .all(...params) +} + +function selectDefinitionRowsByNameAndFiles( + db: Database.Database, + name: string, + targetFiles: ReadonlyArray +): SymbolRow[] { + if (targetFiles.length === 0) { + return [] + } + + const targetFilePlaceholders = targetFiles.map(() => '?').join(', ') + return db + .prepare( + ` + select id, name, file_path, start_line, end_line, kind, signature, parent, language + from symbols + where lower(name) = lower(?) and file_path in (${targetFilePlaceholders}) + order by file_path asc, start_line asc, id asc + ` + ) + .all(name, ...targetFiles) +} + +async function resolveDefinitionEdgeTargetsForNameAndFiles( + db: Database.Database, + root: string, + name: string, + targetFiles: ReadonlyArray, + defaultExportCache?: Map>> +): Promise { + const definitions = selectDefinitionRowsByNameAndFiles(db, name, targetFiles) + return resolveDefinitionEdgeTargets(root, definitions, defaultExportCache) +} + +async function resolveDefinitionEdgeTargets( + root: string, + definitions: ReadonlyArray>, + defaultExportCache: Map>> = new Map() +): Promise { + const targetsByKey = new Map() + + for (const definition of definitions) { + const resolutionMode: Edge['resolution'] = isImportResolvedDefinitionLanguage(definition.language) + ? 'import-resolved' + : 'name-only' + const names = new Set([definition.name]) + + if (resolutionMode === 'import-resolved' && definition.language && isTypeScriptLike(definition.language)) { + const defaultExportNames = await getTypeScriptDefaultExportNames(root, definition.file_path, defaultExportCache) + if (defaultExportNames.has(definition.name) || definition.name === 'default') { + names.add('default') + } + } + + const key = resolutionMode === 'import-resolved' + ? `import-resolved\u0000${definition.file_path}` + : `name-only\u0000${definition.name.toLowerCase()}` + const existing = targetsByKey.get(key) + if (existing) { + for (const name of names) { + if (!existing.names.includes(name)) { + existing.names.push(name) + } + } + continue + } + + const target: DefinitionEdgeTarget = { + names: [...names], + resolutionMode + } + if (resolutionMode === 'import-resolved') { + target.file = definition.file_path + } + targetsByKey.set(key, target) + } + + return [...targetsByKey.values()] +} + +function isImportResolvedDefinitionLanguage(language: string | null | undefined): boolean { + return Boolean(language && (isTypeScriptLike(language) || isPythonLike(language))) +} + +async function getTypeScriptDefaultExportNames( + root: string, + file: string, + cache: Map>> +): Promise> { + const cached = cache.get(file) + if (cached) { + return cached + } + + const load = readFile(resolve(root, file), 'utf8') + .then((content) => { + const sourceFile = ts.createSourceFile(file, content, ts.ScriptTarget.Latest, true, scriptKindFromPath(file)) + return collectTypeScriptDefaultExportNames(sourceFile) + }) + .catch(() => new Set()) + cache.set(file, load) + return load +} + +function collectTypeScriptDefaultExportNames(sourceFile: ts.SourceFile): Set { + const names = new Set() + + for (const statement of sourceFile.statements) { + if ((ts.isFunctionDeclaration(statement) || ts.isClassDeclaration(statement)) && hasTypeScriptDefaultExportModifier(statement)) { + names.add(statement.name?.text ?? 'default') + continue + } + + if (ts.isExportAssignment(statement) && !statement.isExportEquals) { + if (ts.isIdentifier(statement.expression)) { + names.add(statement.expression.text) + } else if (ts.isFunctionExpression(statement.expression) || ts.isClassExpression(statement.expression)) { + names.add(statement.expression.name?.text ?? 'default') + } else { + names.add('default') + } + continue + } + + if (ts.isExportDeclaration(statement) && statement.exportClause && ts.isNamedExports(statement.exportClause)) { + for (const element of statement.exportClause.elements) { + if (element.name.text === 'default') { + names.add((element.propertyName ?? element.name).text) + } + } + } + } + + return names +} + +function hasTypeScriptDefaultExportModifier(node: ts.Node): boolean { + const modifiers = ts.canHaveModifiers(node) ? (ts.getModifiers(node) ?? []) : [] + return modifiers.some((modifier) => modifier.kind === ts.SyntaxKind.ExportKeyword) && + modifiers.some((modifier) => modifier.kind === ts.SyntaxKind.DefaultKeyword) +} + +function selectDefinitionEdgeRows( + db: Database.Database, + targets: ReadonlyArray, + edgeKinds: ReadonlyArray, + preferCallsFirst: boolean, + limit?: number +): EdgeResultRow[] { + if (targets.length === 0 || edgeKinds.length === 0) { + return [] + } + + const edgeKindPlaceholders = edgeKinds.map(() => '?').join(', ') + const rowsById = new Map() + + for (const target of targets) { + if (target.names.length === 0) { + continue + } + + const nameClauses = target.names.map(() => 'lower(dst_name) = lower(?)').join(' or ') + const params: string[] = [...target.names] + let whereClause = `(${nameClauses}) and dst_file is null and edge_kind in (${edgeKindPlaceholders})` + if (target.resolutionMode === 'import-resolved' && target.file) { + whereClause = `(${nameClauses}) and dst_file = ? and edge_kind in (${edgeKindPlaceholders})` + params.push(target.file) + } + params.push(...edgeKinds) + + const targetRows = db + .prepare( + ` + select id, src_file, src_line, src_symbol, edge_kind, resolution, language + from edges + where ${whereClause} + ` + ) + .all(...params) + + for (const row of targetRows) { + rowsById.set(row.id, row) + } + } + + const rows = [...rowsById.values()].sort((left, right) => compareEdgeResultRows(left, right, preferCallsFirst)) + return limit === undefined ? rows : rows.slice(0, limit) +} + +function compareEdgeResultRows(left: EdgeResultRow, right: EdgeResultRow, preferCallsFirst: boolean): number { + if (preferCallsFirst && left.edge_kind !== right.edge_kind) { + if (left.edge_kind === 'call') { + return -1 + } + if (right.edge_kind === 'call') { + return 1 + } + } + + return ( + left.src_file.localeCompare(right.src_file) || + left.src_line - right.src_line || + left.id - right.id + ) +} + +function selectDefinitionUsageRows( + db: Database.Database, + targets: ReadonlyArray, + options: { + excludedSrcFile: string + excludedStartLine: number + excludedEndLine: number + pathGlob?: string + limit?: number + } +): EdgeUsageRow[] { + const rows = selectDefinitionEdgeRows(db, targets, ['call', 'ref'], false) + .filter( + (row) => + !(row.src_file === options.excludedSrcFile && row.src_line >= options.excludedStartLine && row.src_line <= options.excludedEndLine) + ) + .filter((row) => !options.pathGlob || minimatch(row.src_file, options.pathGlob)) + .map((row) => ({ + id: row.id, + src_file: row.src_file, + src_line: row.src_line, + language: row.language, + resolution: row.resolution + })) + + return options.limit === undefined ? rows : rows.slice(0, options.limit) +} + +function countDefinitionEdgeRows( + db: Database.Database, + targets: ReadonlyArray, + edgeKinds: ReadonlyArray +): number { + return selectDefinitionEdgeRows(db, targets, edgeKinds, false).length +} + +function selectImporterEdgeRows(db: Database.Database, file: string): EdgeResultRow[] { + return db + .prepare( + ` + select id, src_file, src_line, src_symbol, edge_kind, resolution, language + from edges + where dst_file = ? and edge_kind = 'import' + order by src_file asc, src_line asc, id asc + ` + ) + .all(file) +} + +function selectSameFileNeighborRows( + db: Database.Database, + file: string, + excludedId: number, + anchorStartLine: number, + limit: number +): SymbolRow[] { + return db + .prepare( + ` + select id, name, file_path, start_line, end_line, kind, signature, parent, language + from symbols + where file_path = ? and id <> ? + order by abs(start_line - ?) asc, start_line asc, id asc + limit ${limit} + ` + ) + .all(file, excludedId, anchorStartLine) +} + +function countSameFileNeighborRows(db: Database.Database, file: string, excludedId: number): number { + return db + .prepare( + ` + select count(*) as count + from symbols + where file_path = ? and id <> ? + ` + ) + .get(file, excludedId)?.count ?? 0 +} + +function buildSymbolNeighbor(row: SymbolRow): SymbolNeighbor { + const neighbor: SymbolNeighbor = { + name: row.name, + file: row.file_path, + range: { startLine: row.start_line, endLine: row.end_line }, + kind: row.kind + } + + if (row.parent) { + neighbor.parent = row.parent + } + + if (row.language) { + neighbor.language = row.language + } + + return neighbor +} + +async function readEdgeResultsFromRows( + root: string, + rows: EdgeResultRow[], + maxTokens: number | undefined +): Promise> { + const tokenBudget = normalizeSearchTokenBudget(maxTokens) + const linesByFile = new Map() + const results: EdgeResult[] = [] + let tokensUsed = 0 + let tokenTruncated = false + let readFailures = 0 + + for (let index = 0; index < rows.length; index += 1) { + const row = rows[index]! + let lines = linesByFile.get(row.src_file) + if (!lines) { + try { + lines = splitLines(await readFile(resolve(root, row.src_file), 'utf8')) + linesByFile.set(row.src_file, lines) + } catch { + readFailures += 1 + continue + } + } + + const result: EdgeResult = { + file: row.src_file, + range: { startLine: row.src_line, endLine: row.src_line }, + line: row.src_line, + snippet: (lines[row.src_line - 1] ?? '').trimEnd(), + edgeKind: row.edge_kind, + resolution: row.resolution + } + + if (row.src_symbol) { + result.srcSymbol = row.src_symbol + } + + if (row.language) { + result.language = row.language + } + + const estimatedTokens = estimateEdgeResultTokens(result) + if (tokenBudget !== undefined && tokensUsed + estimatedTokens > tokenBudget) { + if (results.length === 0) { + results.push({ + ...result, + snippet: truncateToTokenBudget(result.snippet, Math.max(1, Math.max(0, tokenBudget - 8))) + }) + tokenTruncated = index < rows.length - 1 + } else { + tokenTruncated = true + } + break + } + + results.push(result) + tokensUsed += estimatedTokens + } + + return { items: results, tokenTruncated, readFailures } +} + +function estimateEdgeResultTokens(result: EdgeResult): number { + const header = `${result.file}:${result.line} ${result.srcSymbol ?? ''} ${result.edgeKind} ${result.resolution}` + return SEARCH_HIT_TOKEN_OVERHEAD + estimateTokenCount(header) + estimateTokenCount(result.snippet) +} + +async function readImpactNodesFromRows( + root: string, + rows: ImpactNodeRow[], + maxTokens: number | undefined +): Promise> { + const tokenBudget = normalizeSearchTokenBudget(maxTokens) + const linesByFile = new Map() + const nodes: ImpactNode[] = [] + let tokensUsed = 0 + let tokenTruncated = false + let readFailures = 0 + + for (let index = 0; index < rows.length; index += 1) { + const row = rows[index]! + let lines = linesByFile.get(row.src_file) + if (!lines) { + try { + lines = splitLines(await readFile(resolve(root, row.src_file), 'utf8')) + linesByFile.set(row.src_file, lines) + } catch { + readFailures += 1 + continue + } + } + + const node: ImpactNode = { + name: row.name, + file: row.src_file, + range: { startLine: row.src_line, endLine: row.src_line }, + line: row.src_line, + snippet: (lines[row.src_line - 1] ?? '').trimEnd(), + depth: row.depth, + edgeKind: row.edge_kind, + resolution: row.resolution + } + + if (row.src_symbol) { + node.srcSymbol = row.src_symbol + } + + if (row.language) { + node.language = row.language + } + + const estimatedTokens = estimateImpactNodeTokens(node) + if (tokenBudget !== undefined && tokensUsed + estimatedTokens > tokenBudget) { + if (nodes.length === 0) { + nodes.push({ + ...node, + snippet: truncateToTokenBudget(node.snippet, Math.max(1, Math.max(0, tokenBudget - 10))) + }) + tokenTruncated = index < rows.length - 1 + } else { + tokenTruncated = true + } + break + } + + nodes.push(node) + tokensUsed += estimatedTokens + } + + return { items: nodes, tokenTruncated, readFailures } +} + +function estimateImpactNodeTokens(node: ImpactNode): number { + const header = `${node.file}:${node.line} ${node.name} d${node.depth} ${node.edgeKind} ${node.resolution}` + return SEARCH_HIT_TOKEN_OVERHEAD + estimateTokenCount(header) + estimateTokenCount(node.snippet) +} + function selectGrepCandidateFiles(db: Database.Database, options?: GrepOptions): Array<{ path: string; language: string }> { const whereClauses: string[] = [] const params: string[] = [] diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index c9e2ace..fcade26 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -22,6 +22,22 @@ export interface Range { export type SearchReasonTag = '=' | '~' | '+' +export type EdgeKind = 'ref' | 'call' | 'import' | 'implements' | 'extends' + +export type EdgeResolution = 'import-resolved' | 'name-only' + +export interface Edge { + id?: number + srcFile: string + srcLine: number + srcSymbol?: string + dstName: string + dstFile?: string + edgeKind: EdgeKind + resolution: EdgeResolution + language?: string +} + export interface SymbolUsage { file: string range: Range @@ -29,9 +45,34 @@ export interface SymbolUsage { snippet: string language?: string /** - * Honest provenance marker: usages are import-resolved/local, not type-resolved. + * Honest provenance marker: usages are import-resolved or name-only, never type-resolved. + */ + resolution: EdgeResolution +} + +export interface SymbolNeighbor { + name: string + file: string + range: Range + kind: SymbolKind + parent?: string + language?: string +} + +export interface SymbolRelations { + /** + * Top caller/reference sites for the resolved definition. These reuse the + * persisted edge rows and stay honest about resolution confidence. */ - resolution: 'import-resolved' + sites: EdgeResult[] + /** + * Nearby same-file symbols around the resolved definition for quick local context. + */ + neighbors: SymbolNeighbor[] + /** + * Additional relation items omitted by the bounded read-side bundle. + */ + omitted?: number } export interface SearchHit { @@ -99,6 +140,12 @@ export interface SymbolDefinition { * on compact/ambiguous rows. */ body?: string + /** + * OPT-IN bounded relational context for the top exact match: top caller/ref + * sites plus same-file neighbors. Present only when requested and only on the + * top exact row of an unambiguous lookup. + */ + relations?: SymbolRelations } export interface SearchOptions { @@ -153,6 +200,74 @@ export interface FindSymbolOptions { * compact. Capped like search bodies. Set false for a compact name→location list. */ withBody?: boolean + /** + * OPT-IN bounded relational addendum for the top exact match: top caller/ref + * sites plus same-file neighbors. Default false. + */ + withCallers?: boolean +} + +export interface FindEdgeOptions { + /** + * Target definition kind filter used while resolving the destination symbol. + */ + kind?: SymbolKind | SymbolKind[] + /** + * Target definition path glob used for collision disambiguation, e.g. `src/schema/**`. + */ + pathGlob?: string + /** + * Approx output token budget for the returned edge rows. + */ + maxTokens?: number +} + +export interface FindImportersOptions { + /** + * Approx output token budget for the returned importer rows. + */ + maxTokens?: number +} + +export interface ImpactOptions extends FindEdgeOptions { + /** + * Caller depth to traverse: 0=direct callers only, 1=callers of callers, etc. + * Default 2. + */ + depth?: number + /** + * Hard cap on returned graph nodes. Default 50. + */ + maxNodes?: number +} + +export interface EdgeResult { + file: string + range: Range + line: number + snippet: string + srcSymbol?: string + edgeKind: EdgeKind + resolution: EdgeResolution + language?: string +} + +export interface ImpactNode extends EdgeResult { + /** + * Symbol name traversed at this node. Falls back to `top-level` when the edge + * originates outside an enclosing symbol. + */ + name: string + depth: number +} + +export interface ImpactResult { + nodes: ImpactNode[] + impactTruncated?: boolean + depthCapped?: boolean + nodesCapped?: boolean + depthLimit: number + maxNodes: number } export interface SyncProgressEvent { @@ -280,6 +395,11 @@ export interface Repo { search(query: string, options?: SearchOptions): Promise grep(pattern: string, options?: GrepOptions): Promise findSymbol(name: string, options?: FindSymbolOptions): Promise + findCallers(name: string, options?: FindEdgeOptions): Promise + impact(name: string, options?: ImpactOptions): Promise + findReferences(name: string, options?: FindEdgeOptions): Promise + findImplementers(name: string, options?: FindEdgeOptions): Promise + findImporters(file: string, options?: FindImportersOptions): Promise readChunk(id: string, options?: ReadChunkOptions): Promise readRange(file: string, startLine: number, endLine: number, options?: ReadRangeOptions): Promise status(): Promise diff --git a/packages/core/test/core.test.ts b/packages/core/test/core.test.ts index 794d53a..5d9e1fc 100644 --- a/packages/core/test/core.test.ts +++ b/packages/core/test/core.test.ts @@ -1,5 +1,5 @@ import Database from 'better-sqlite3' -import { mkdtemp, mkdir, rm, symlink, unlink, writeFile } from 'node:fs/promises' +import { cp, mkdtemp, mkdir, rm, symlink, unlink, writeFile } from 'node:fs/promises' import { tmpdir } from 'node:os' import { join } from 'node:path' @@ -107,6 +107,14 @@ TokenVerifier is the main entry point described in docs. TokenVerifier appears h return repoRoot } +async function copyFixtureRepository(fixtureName: string, prefix: string): Promise { + const parentDirectory = await mkdtemp(join(tmpdir(), prefix)) + const repoRoot = join(parentDirectory, 'repo') + temporaryDirectories.push(parentDirectory) + await cp(join(process.cwd(), 'packages', 'eval', 'fixtures', fixtureName), repoRoot, { recursive: true }) + return repoRoot +} + async function sleep(ms: number): Promise { await new Promise((resolve) => setTimeout(resolve, ms)) } @@ -740,7 +748,197 @@ describe('@codesift/core find_symbol body inlining', () => { }) }) +describe('@codesift/core relational find_symbol + impact', () => { + it('bundles caller/ref sites and same-file neighbors for the top exact definition when withCallers:true', async () => { + const repoRoot = await copyFixtureRepository('usages-ts', 'codesift-findsymbol-relations-') + const repo = await openRepo(repoRoot) + + await repo.sync() + + const [definition] = await repo.findSymbol('parseToken', { withCallers: true }) + + expect(definition?.file).toBe('src/token.ts') + expect(definition?.body).toContain('export function parseToken') + expect(definition?.relations?.sites.map((site) => `${site.file}:${site.line}:${site.srcSymbol}:${site.edgeKind}:${site.resolution}`)).toEqual([ + 'src/api.ts:4:readSubject:call:import-resolved', + 'src/worker.ts:4:enqueueToken:call:import-resolved' + ]) + expect(definition?.relations?.neighbors).toMatchObject([ + { + name: 'ParsedToken', + file: 'src/token.ts', + kind: 'interface', + range: { startLine: 1, endLine: 4 } + } + ]) + }) + + it('walks bounded transitive callers and reports depth/node caps without exploding', async () => { + const repoRoot = await copyFixtureRepository('usages-ts', 'codesift-impact-') + await writeFile( + join(repoRoot, 'src', 'service.ts'), + `import { readSubject } from './api' + +export function handleToken(header: string): string { + return readSubject(header) +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'src', 'app.ts'), + `import { handleToken } from './service' + +export function runApp(header: string): string { + return handleToken(header) +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'src', 'cli.ts'), + `import { runApp } from './app' + +export function main(header: string): string { + return runApp(header) +} +`, + 'utf8' + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + + const impact = await repo.impact('parseToken', { depth: 2 }) + expect(impact.nodes.map((node) => `${node.depth}:${node.file}:${node.line}:${node.srcSymbol}:${node.edgeKind}:${node.resolution}`)).toEqual([ + '0:src/api.ts:4:readSubject:call:import-resolved', + '0:src/worker.ts:4:enqueueToken:call:import-resolved', + '1:src/service.ts:4:handleToken:call:import-resolved', + '2:src/app.ts:4:runApp:call:import-resolved' + ]) + expect(impact.depthCapped).toBe(true) + expect(impact.nodesCapped).not.toBe(true) + expect(impact.impactTruncated).not.toBe(true) + + const capped = await repo.impact('parseToken', { depth: 3, maxNodes: 2 }) + expect(capped.nodes).toHaveLength(2) + expect(capped.nodesCapped).toBe(true) + expect(capped.impactTruncated).toBe(true) + }) + + it('guards impact traversal against caller cycles', async () => { + const repoRoot = await copyFixtureRepository('usages-ts', 'codesift-impact-cycle-') + await writeFile( + join(repoRoot, 'src', 'a.ts'), + `import { parseToken } from './token' +import { c } from './c' + +export function a(raw: string): string { + return raw.length > 0 ? parseToken(raw).subject : c(raw) +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'src', 'b.ts'), + `import { a } from './a' + +export function b(raw: string): string { + return a(raw) +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'src', 'c.ts'), + `import { b } from './b' + +export function c(raw: string): string { + return b(raw) +} +`, + 'utf8' + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + + const impact = await repo.impact('parseToken', { depth: 10, maxNodes: 50 }) + expect(impact.nodes.map((node) => `${node.depth}:${node.file}:${node.line}:${node.srcSymbol}:${node.edgeKind}`)).toEqual([ + '0:src/a.ts:5:a:call', + '0:src/api.ts:4:readSubject:call', + '0:src/worker.ts:4:enqueueToken:call', + '1:src/b.ts:4:b:call', + '2:src/c.ts:4:c:call', + '3:src/a.ts:5:a:call' + ]) + expect(impact.depthCapped).not.toBe(true) + expect(impact.nodesCapped).not.toBe(true) + expect(impact.impactTruncated).not.toBe(true) + expect(impact.maxNodes).toBe(50) + + const hardCapped = await repo.impact('parseToken', { depth: 999, maxNodes: 1_000_000 }) + expect(hardCapped.depthLimit).toBe(50) + expect(hardCapped.maxNodes).toBe(50) + }) +}) + describe('@codesift/core import-resolved usages', () => { + it('writes TS edges at index time and serves fixture usages from the persisted index', async () => { + const repoRoot = await copyFixtureRepository('usages-ts', 'codesift-fixture-usages-ts-') + const repo = await openRepo(repoRoot) + + await repo.sync() + + const db = new Database(join(repoRoot, '.codesift', 'index.db')) + const rows = db + .prepare< + [string], + { + src_file: string + src_line: number + src_symbol: string | null + dst_file: string | null + edge_kind: string + resolution: string + } + >( + ` + select src_file, src_line, src_symbol, dst_file, edge_kind, resolution + from edges + where dst_name = ? + order by src_file asc, src_line asc, edge_kind asc + ` + ) + .all('parseToken') + db.close() + + expect(rows).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + src_file: 'src/api.ts', + src_line: 4, + src_symbol: 'readSubject', + dst_file: 'src/token.ts', + edge_kind: 'call', + resolution: 'import-resolved' + }), + expect.objectContaining({ + src_file: 'src/worker.ts', + src_line: 4, + src_symbol: 'enqueueToken', + dst_file: 'src/token.ts', + edge_kind: 'call', + resolution: 'import-resolved' + }) + ]) + ) + + const hits = await repo.search('parseToken', { k: 1, withUsages: true }) + expect(hits[0]?.symbol).toBe('parseToken') + expect(hits[0]?.usages?.map((usage) => `${usage.file}:${usage.line}`)).toEqual(['src/api.ts:4', 'src/worker.ts:4']) + }) + it('bundles TS/JS import-resolved and same-file usages for the top definition hit', async () => { const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-usages-ts-')) temporaryDirectories.push(repoRoot) @@ -791,6 +989,449 @@ export const ok = auth.verifyJwtToken('eyJ-demo') expect(hits[0]?.usages?.every((usage) => usage.resolution === 'import-resolved')).toBe(true) }) + it('keys persisted usages by definition file for colliding TS definitions', async () => { + const repoRoot = await copyFixtureRepository('collision-ts', 'codesift-fixture-collision-ts-') + const repo = await openRepo(repoRoot) + + await repo.sync() + + const hits = await repo.search('name must be a non-empty string', { k: 1, withUsages: true }) + expect(hits[0]?.file).toBe('src/schema/validator.ts') + expect(hits[0]?.symbol).toBe('validate') + expect(hits[0]?.usages?.map((usage) => `${usage.file}:${usage.line}`)).toEqual(['src/api/handler.ts:6']) + + const db = new Database(join(repoRoot, '.codesift', 'index.db')) + const rows = db + .prepare<[string, string], { src_file: string; src_line: number; dst_file: string | null; edge_kind: string }>( + ` + select src_file, src_line, dst_file, edge_kind + from edges + where dst_name = ? and dst_file = ? and edge_kind in ('call', 'ref') + order by src_file asc, src_line asc + ` + ) + .all('validate', 'src/schema/validator.ts') + db.close() + + expect(rows).toEqual([ + { + src_file: 'src/api/handler.ts', + src_line: 6, + dst_file: 'src/schema/validator.ts', + edge_kind: 'call' + } + ]) + }) + + it('finds callers, references, and importers from persisted edges', async () => { + const repoRoot = await copyFixtureRepository('usages-ts', 'codesift-fixture-graph-usages-ts-') + const repo = await openRepo(repoRoot) + + await repo.sync() + + const callers = await repo.findCallers('parseToken') + const references = await repo.findReferences('parseToken') + const importers = await repo.findImporters('src/token.ts') + + expect(callers).toMatchObject([ + { + file: 'src/api.ts', + line: 4, + range: { startLine: 4, endLine: 4 }, + srcSymbol: 'readSubject', + edgeKind: 'call', + resolution: 'import-resolved', + snippet: " const claims = parseToken(header.replace('Bearer ', ''))", + language: 'typescript' + }, + { + file: 'src/worker.ts', + line: 4, + range: { startLine: 4, endLine: 4 }, + srcSymbol: 'enqueueToken', + edgeKind: 'call', + resolution: 'import-resolved', + snippet: ' const claims = parseToken(raw)', + language: 'typescript' + } + ]) + + expect(references).toMatchObject([ + { + file: 'src/api.ts', + line: 4, + srcSymbol: 'readSubject', + edgeKind: 'call', + resolution: 'import-resolved' + }, + { + file: 'src/worker.ts', + line: 4, + srcSymbol: 'enqueueToken', + edgeKind: 'call', + resolution: 'import-resolved' + } + ]) + + expect(importers).toMatchObject([ + { + file: 'src/api.ts', + line: 1, + range: { startLine: 1, endLine: 1 }, + edgeKind: 'import', + resolution: 'import-resolved', + snippet: "import { parseToken } from './token'", + language: 'typescript' + }, + { + file: 'src/worker.ts', + line: 1, + range: { startLine: 1, endLine: 1 }, + edgeKind: 'import', + resolution: 'import-resolved', + snippet: "import { parseToken } from './token'", + language: 'typescript' + } + ]) + }) + + it('distinguishes callers from non-call references', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-graph-call-vs-ref-')) + temporaryDirectories.push(repoRoot) + await mkdir(join(repoRoot, 'src'), { recursive: true }) + + await writeFile( + join(repoRoot, 'src', 'token.ts'), + `export function parseToken(token: string): string { + return token.trim() +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'src', 'api.ts'), + `import { parseToken } from './token' + +export function readSubject(header: string): string { + return parseToken(header) +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'src', 'ref.ts'), + `import { parseToken } from './token' + +export function keepParser(token: string): string { + const parser = parseToken + return parser(token) +} +`, + 'utf8' + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + + expect((await repo.findCallers('parseToken')).map((result) => `${result.file}:${result.line}:${result.srcSymbol}:${result.edgeKind}`)).toEqual([ + 'src/api.ts:4:readSubject:call' + ]) + expect((await repo.findReferences('parseToken')).map((result) => `${result.file}:${result.line}:${result.srcSymbol}:${result.edgeKind}`)).toEqual([ + 'src/api.ts:4:readSubject:call', + 'src/ref.ts:4:keepParser:ref' + ]) + }) + + it('resolves default-import callers, refs, impact, and usages to the real default-exported symbol name', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-default-import-')) + temporaryDirectories.push(repoRoot) + await mkdir(join(repoRoot, 'src'), { recursive: true }) + + await writeFile( + join(repoRoot, 'src', 'token.ts'), + `export default function parseToken(token: string): string { + return token.trim() +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'src', 'api.ts'), + `import parseToken from './token' + +export function readSubject(header: string): string { + return parseToken(header) +} +`, + 'utf8' + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + + expect((await repo.findCallers('parseToken')).map((result) => `${result.file}:${result.line}:${result.srcSymbol}:${result.edgeKind}`)).toEqual([ + 'src/api.ts:4:readSubject:call' + ]) + expect((await repo.findReferences('parseToken')).map((result) => `${result.file}:${result.line}:${result.srcSymbol}:${result.edgeKind}`)).toEqual([ + 'src/api.ts:4:readSubject:call' + ]) + expect((await repo.impact('parseToken', { depth: 1 })).nodes.map((node) => `${node.depth}:${node.file}:${node.line}:${node.srcSymbol}:${node.edgeKind}`)).toEqual([ + '0:src/api.ts:4:readSubject:call' + ]) + expect((await repo.search('parseToken', { k: 1, withUsages: true }))[0]?.usages?.map((usage) => `${usage.file}:${usage.line}`)).toEqual([ + 'src/api.ts:4' + ]) + }) + + it('disambiguates callers and references by destination file for colliding definitions', async () => { + const repoRoot = await copyFixtureRepository('collision-ts', 'codesift-fixture-graph-collision-ts-') + const repo = await openRepo(repoRoot) + + await repo.sync() + + const callers = await repo.findCallers('validate', { kind: 'function', pathGlob: 'src/schema/**' }) + const references = await repo.findReferences('validate', { kind: 'function', pathGlob: 'src/schema/**' }) + + expect(callers).toMatchObject([ + { + file: 'src/api/handler.ts', + line: 6, + srcSymbol: 'handleRequest', + edgeKind: 'call', + resolution: 'import-resolved' + } + ]) + expect(references).toMatchObject([ + { + file: 'src/api/handler.ts', + line: 6, + srcSymbol: 'handleRequest', + edgeKind: 'call', + resolution: 'import-resolved' + } + ]) + }) + + it('writes TS implements/extends edges and finds implementers from persisted heritage clauses', async () => { + const repoRoot = await copyFixtureRepository('heritage-ts', 'codesift-fixture-heritage-ts-') + const repo = await openRepo(repoRoot) + + await repo.sync() + + const db = new Database(join(repoRoot, '.codesift', 'index.db')) + const rows = db + .prepare< + [string], + { + src_file: string + src_line: number + src_symbol: string | null + dst_file: string | null + edge_kind: string + resolution: string + } + >( + ` + select src_file, src_line, src_symbol, dst_file, edge_kind, resolution + from edges + where dst_name = ? and edge_kind in ('implements', 'extends') + order by src_file asc, src_line asc, edge_kind asc + ` + ) + .all('AuthStrategy') + db.close() + + expect(rows).toEqual([ + { + src_file: 'src/impl.ts', + src_line: 3, + src_symbol: 'JwtVerifier', + dst_file: 'src/contract.ts', + edge_kind: 'implements', + resolution: 'import-resolved' + }, + { + src_file: 'src/impl.ts', + src_line: 9, + src_symbol: 'StrictStrategy', + dst_file: 'src/contract.ts', + edge_kind: 'extends', + resolution: 'import-resolved' + } + ]) + + const authImplementers = await repo.findImplementers('AuthStrategy') + const baseImplementers = await repo.findImplementers('BaseVerifier') + + expect(authImplementers).toMatchObject([ + { + file: 'src/impl.ts', + line: 3, + srcSymbol: 'JwtVerifier', + edgeKind: 'implements', + resolution: 'import-resolved' + }, + { + file: 'src/impl.ts', + line: 9, + srcSymbol: 'StrictStrategy', + edgeKind: 'extends', + resolution: 'import-resolved' + } + ]) + expect(baseImplementers).toMatchObject([ + { + file: 'src/impl.ts', + line: 3, + srcSymbol: 'JwtVerifier', + edgeKind: 'extends', + resolution: 'import-resolved' + } + ]) + }) + + it('matches caller, ref, implementer, and impact lookups case-insensitively', async () => { + const usagesRepoRoot = await copyFixtureRepository('usages-ts', 'codesift-graph-case-usages-') + const usagesRepo = await openRepo(usagesRepoRoot) + await usagesRepo.sync() + + expect((await usagesRepo.findCallers('ParseToken')).map((result) => `${result.file}:${result.line}:${result.srcSymbol}:${result.edgeKind}`)).toEqual([ + 'src/api.ts:4:readSubject:call', + 'src/worker.ts:4:enqueueToken:call' + ]) + expect((await usagesRepo.findReferences('ParseToken')).map((result) => `${result.file}:${result.line}:${result.srcSymbol}:${result.edgeKind}`)).toEqual([ + 'src/api.ts:4:readSubject:call', + 'src/worker.ts:4:enqueueToken:call' + ]) + expect((await usagesRepo.impact('ParseToken', { depth: 0 })).nodes.map((node) => `${node.depth}:${node.file}:${node.line}:${node.srcSymbol}:${node.edgeKind}`)).toEqual([ + '0:src/api.ts:4:readSubject:call', + '0:src/worker.ts:4:enqueueToken:call' + ]) + + const heritageRepoRoot = await copyFixtureRepository('heritage-ts', 'codesift-graph-case-heritage-') + const heritageRepo = await openRepo(heritageRepoRoot) + await heritageRepo.sync() + + expect((await heritageRepo.findImplementers('authstrategy')).map((result) => `${result.file}:${result.line}:${result.srcSymbol}:${result.edgeKind}`)).toEqual([ + 'src/impl.ts:3:JwtVerifier:implements', + 'src/impl.ts:9:StrictStrategy:extends' + ]) + }) + + it('does not mix name-only callers into a disambiguated import-resolved query', async () => { + const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-graph-name-only-disambiguation-')) + temporaryDirectories.push(repoRoot) + await mkdir(join(repoRoot, 'src'), { recursive: true }) + await mkdir(join(repoRoot, 'auth'), { recursive: true }) + + await writeFile( + join(repoRoot, 'src', 'token.ts'), + `export function parseToken(token: string): string { + return token.trim() +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'src', 'api.ts'), + `import { parseToken } from './token' + +export function readSubject(header: string): string { + return parseToken(header) +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'auth', 'token.go'), + `package auth + +func parseToken(token string) string { + return token +} + +func handleToken(token string) string { + return parseToken(token) +} +`, + 'utf8' + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + + expect((await repo.findCallers('parseToken', { kind: 'function', pathGlob: 'src/**' })).map((result) => `${result.file}:${result.line}:${result.srcSymbol}:${result.resolution}`)).toEqual([ + 'src/api.ts:4:readSubject:import-resolved' + ]) + expect((await repo.findCallers('parseToken', { kind: 'function', pathGlob: 'auth/**' })).map((result) => `${result.file}:${result.line}:${result.srcSymbol}:${result.resolution}`)).toEqual([ + 'auth/token.go:8:handleToken:name-only' + ]) + }) + + it('writes name-only Go call edges and serves callers from persisted edges', async () => { + const repoRoot = await copyFixtureRepository('m3-go', 'codesift-fixture-graph-m3-go-') + await writeFile( + join(repoRoot, 'auth', 'consumer.go'), + `package auth + +func ValidateBearer(token string) bool { + verifier := NewTokenVerifier() + return verifier.VerifyToken(token) +} +`, + 'utf8' + ) + const repo = await openRepo(repoRoot) + + await repo.sync() + + const db = new Database(join(repoRoot, '.codesift', 'index.db')) + const rows = db + .prepare< + [string], + { + src_file: string + src_line: number + src_symbol: string | null + dst_file: string | null + edge_kind: string + resolution: string + } + >( + ` + select src_file, src_line, src_symbol, dst_file, edge_kind, resolution + from edges + where dst_name = ? and edge_kind = 'call' + order by src_file asc, src_line asc + ` + ) + .all('VerifyToken') + db.close() + + expect(rows).toEqual([ + { + src_file: 'auth/consumer.go', + src_line: 5, + src_symbol: 'ValidateBearer', + dst_file: null, + edge_kind: 'call', + resolution: 'name-only' + } + ]) + + const callers = await repo.findCallers('VerifyToken', { kind: 'method', pathGlob: 'auth/**' }) + expect(callers).toMatchObject([ + { + file: 'auth/consumer.go', + line: 5, + srcSymbol: 'ValidateBearer', + edgeKind: 'call', + resolution: 'name-only', + language: 'go' + } + ]) + }) + it('bundles Python import-resolved usages for the top definition hit', async () => { const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-usages-py-')) temporaryDirectories.push(repoRoot) @@ -817,11 +1458,154 @@ def handle(token: str) -> bool: const repo = await openRepo(repoRoot) await repo.sync() - const hits = await repo.search('verify_token', { k: 1, withUsages: true, pathGlob: 'pkg/**' }) + const db = new Database(join(repoRoot, '.codesift', 'index.db')) + const rows = db + .prepare< + [string], + { + src_file: string + src_line: number + src_symbol: string | null + dst_file: string | null + edge_kind: string + resolution: string + } + >( + ` + select src_file, src_line, src_symbol, dst_file, edge_kind, resolution + from edges + where dst_name = ? and edge_kind in ('call', 'ref') + order by src_file asc, src_line asc + ` + ) + .all('verify_token') + db.close() + + expect(rows).toEqual([ + { + src_file: 'pkg/consumer.py', + src_line: 5, + src_symbol: 'handle', + dst_file: 'pkg/token.py', + edge_kind: 'call', + resolution: 'import-resolved' + } + ]) + + const hits = await repo.search('verify_token', { k: 1, withUsages: true, pathGlob: 'pkg/**' }) expect(hits[0]?.symbol).toBe('verify_token') expect(hits[0]?.usages?.map((usage) => `${usage.file}:${usage.line}`)).toEqual(['pkg/consumer.py:5']) }) + + it('skips shadowed TS same-file matches and Python member-access matches under import-resolved labels', async () => { + const tsRepoRoot = await mkdtemp(join(tmpdir(), 'codesift-graph-shadowed-ts-')) + temporaryDirectories.push(tsRepoRoot) + await mkdir(join(tsRepoRoot, 'src'), { recursive: true }) + await writeFile( + join(tsRepoRoot, 'src', 'token.ts'), + `export function parseToken(token: string): string { + return token.trim() +} + +export function wrapToken(raw: string): string { + const parseToken = (value: string): string => value.toUpperCase() + return parseToken(raw) +} +`, + 'utf8' + ) + + const tsRepo = await openRepo(tsRepoRoot) + await tsRepo.sync() + expect(await tsRepo.findReferences('parseToken')).toEqual([]) + + const pyRepoRoot = await mkdtemp(join(tmpdir(), 'codesift-graph-member-access-py-')) + temporaryDirectories.push(pyRepoRoot) + await mkdir(join(pyRepoRoot, 'pkg'), { recursive: true }) + await writeFile(join(pyRepoRoot, 'pkg', '__init__.py'), '', 'utf8') + await writeFile( + join(pyRepoRoot, 'pkg', 'token.py'), + `def validate(token: str) -> str: + return token + +class Service: + def validate(self, token: str) -> str: + return token + + def run(self, token: str) -> str: + return self.validate(token) +`, + 'utf8' + ) + + const pyRepo = await openRepo(pyRepoRoot) + await pyRepo.sync() + expect(await pyRepo.findReferences('validate', { kind: 'function', pathGlob: 'pkg/token.py' })).toEqual([]) + }) + + it('replaces persisted edges on incremental re-sync when a source file changes', async () => { + const repoRoot = await copyFixtureRepository('usages-ts', 'codesift-fixture-usages-ts-resync-') + const repo = await openRepo(repoRoot) + + await repo.sync() + await writeFile( + join(repoRoot, 'src', 'api.ts'), + `export function readSubject(header: string): string { + return header.replace('Bearer ', '') +} +`, + 'utf8' + ) + await repo.sync() + + const db = new Database(join(repoRoot, '.codesift', 'index.db')) + const apiEdgeCount = db + .prepare<[string, string], { value: number }>('select count(*) as value from edges where src_file = ? and dst_name = ?') + .get('src/api.ts', 'parseToken') + const workerRows = db + .prepare<[string, string], { src_file: string; src_line: number; edge_kind: string }>( + ` + select src_file, src_line, edge_kind + from edges + where src_file = ? and dst_name = ? and edge_kind in ('call', 'ref') + order by src_line asc + ` + ) + .all('src/worker.ts', 'parseToken') + db.close() + + expect(apiEdgeCount?.value).toBe(0) + expect(workerRows).toEqual([{ src_file: 'src/worker.ts', src_line: 4, edge_kind: 'call' }]) + }) + + it('clears both outbound and inbound edges when a file is removed', async () => { + const repoRoot = await copyFixtureRepository('usages-ts', 'codesift-fixture-usages-ts-remove-') + await writeFile( + join(repoRoot, 'src', 'service.ts'), + `import { readSubject } from './api' + +export function handleToken(header: string): string { + return readSubject(header) +} +`, + 'utf8' + ) + const repo = await openRepo(repoRoot) + + await repo.sync() + await unlink(join(repoRoot, 'src', 'api.ts')) + await repo.sync() + + const db = new Database(join(repoRoot, '.codesift', 'index.db')) + const outboundCount = db.prepare<[string], { value: number }>('select count(*) as value from edges where src_file = ?').get('src/api.ts') + const inboundCount = db.prepare<[string], { value: number }>('select count(*) as value from edges where dst_file = ?').get('src/api.ts') + db.close() + + expect(outboundCount?.value).toBe(0) + expect(inboundCount?.value).toBe(0) + expect(await repo.findImporters('src/api.ts')).toEqual([]) + }) }) describe('buildFtsQuery synonym OR-expansion', () => { diff --git a/packages/eval/fixtures/heritage-ts/src/contract.ts b/packages/eval/fixtures/heritage-ts/src/contract.ts new file mode 100644 index 0000000..682bd70 --- /dev/null +++ b/packages/eval/fixtures/heritage-ts/src/contract.ts @@ -0,0 +1,9 @@ +export interface AuthStrategy { + verify(token: string): boolean +} + +export class BaseVerifier { + verify(token: string): boolean { + return token.length > 0 + } +} diff --git a/packages/eval/fixtures/heritage-ts/src/impl.ts b/packages/eval/fixtures/heritage-ts/src/impl.ts new file mode 100644 index 0000000..0077138 --- /dev/null +++ b/packages/eval/fixtures/heritage-ts/src/impl.ts @@ -0,0 +1,9 @@ +import { BaseVerifier, type AuthStrategy } from './contract' + +export class JwtVerifier extends BaseVerifier implements AuthStrategy { + override verify(token: string): boolean { + return super.verify(token) + } +} + +export interface StrictStrategy extends AuthStrategy {} diff --git a/packages/mcp/src/index.ts b/packages/mcp/src/index.ts index 1d420d4..ae4802c 100644 --- a/packages/mcp/src/index.ts +++ b/packages/mcp/src/index.ts @@ -6,9 +6,13 @@ import { DEFAULT_SEARCH_K as CORE_DEFAULT_SEARCH_K, getDefaultEmbeddingProvider, isLearnedEmbeddingProvider, + type EdgeResult, + type FindEdgeOptions, type FindSymbolOptions, type GrepHit, type GrepOptions, + type ImpactOptions, + type ImpactResult, type Repo, type RepoStatus, type SearchHit, @@ -26,6 +30,7 @@ export const DEFAULT_SEARCH_K = CORE_DEFAULT_SEARCH_K export const DEFAULT_MCP_SEARCH_MAX_TOKENS = 700 export const DEFAULT_MCP_GREP_MAX_TOKENS = 700 export const DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS = 700 +export const DEFAULT_MCP_RELATION_MAX_TOKENS = 700 export const DEFAULT_MCP_READ_CHUNK_MAX_TOKENS = 1000 export const DEFAULT_MCP_INDEX_STATUS_MAX_TOKENS = 200 export const MIN_MCP_READ_CHUNK_MAX_TOKENS = 6 @@ -49,6 +54,11 @@ const SYMBOL_KINDS = [ export const MCP_TOOL_NAMES = [ 'search_code', 'find_symbol', + 'find_callers', + 'find_refs', + 'find_importers', + 'who_implements', + 'impact', 'grep_code', 'read_chunk', 'index_status' @@ -73,6 +83,39 @@ export interface FindSymbolArgs { kind?: FindSymbolOptions['kind'] | undefined path_glob?: string | undefined with_body?: boolean | undefined + with_callers?: boolean | undefined + max_tokens?: number | undefined +} + +export interface FindCallersArgs { + name: string + kind?: FindEdgeOptions['kind'] | undefined + path_glob?: string | undefined + max_tokens?: number | undefined +} + +export interface FindRefsArgs { + name: string + kind?: FindEdgeOptions['kind'] | undefined + path_glob?: string | undefined + max_tokens?: number | undefined +} + +export interface FindImportersArgs { + file: string + max_tokens?: number | undefined +} + +export interface FindImplementersArgs { + name: string + max_tokens?: number | undefined +} + +export interface ImpactArgs { + name: string + depth?: number | undefined + kind?: ImpactOptions['kind'] | undefined + path_glob?: string | undefined max_tokens?: number | undefined } @@ -107,6 +150,14 @@ export interface FormatMcpSymbolsOptions { maxTokens?: number | undefined } +export interface FormatMcpEdgeResultsOptions { + maxTokens?: number | undefined +} + +export interface FormatMcpImpactOptions { + maxTokens?: number | undefined +} + export interface ReadChunkArgs { id: string context_lines?: number | undefined @@ -148,6 +199,11 @@ export interface McpServerHandle { export interface McpRouter { searchCode(args: SearchCodeArgs): Promise findSymbol(args: FindSymbolArgs): Promise + findCallers(args: FindCallersArgs): Promise + findReferences(args: FindRefsArgs): Promise + findImporters(args: FindImportersArgs): Promise + findImplementers(args: FindImplementersArgs): Promise + impact(args: ImpactArgs): Promise grepCode(args: GrepCodeArgs): Promise readChunk(args: ReadChunkArgs): Promise indexStatus(): Promise @@ -165,9 +221,9 @@ export type McpJsonRpcResponse = | { jsonrpc: '2.0'; id: string | number | null; error: { code: number; message: string; data?: unknown } } export const MCP_SERVER_INSTRUCTIONS = [ - 'Use codesift before host grep/read. Route identifiers/definitions->find_symbol; literals/env/errors/operators/regex->grep_code; concepts/unknown names->search_code.', - 'Top search_code body is inline; read_chunk only for non-top hits/wider context. Broad search_code: k=5-8. Keep grep_code context small.', - 'Check index_status; if missing/stale/running/failed/aborted, warn and suggest codesift index/sync.' + 'codesift;identifiers->find_symbol;callers/uses->find_callers/find_refs;breakage/transitive callers->impact;importers->find_importers;implements/extends->who_implements;G/J/Rb/Rs approx:name-only;literals/regex/errors/env/operators->grep_code;concepts/fuzzy names->search_code', + 'Top search_code body inline;read_chunk only for non-top/wider context;Broad search k=5-8.', + 'Check index_status;warn+sync if missing/stale/running/failed.' ].join('\n') const symbolKindSchema = z.enum(SYMBOL_KINDS) @@ -190,6 +246,39 @@ const findSymbolInputSchema = { kind: kindFilterSchema.optional().describe('Symbol kind filter.'), path_glob: z.string().min(1).optional().describe('Repo glob, e.g. "src/auth/**".'), with_body: z.boolean().optional().describe('Inline top exact body. Default true.'), + with_callers: z.boolean().optional().describe('Append bounded caller/ref + same-file relations for the top exact hit.'), + max_tokens: z.number().int().positive().max(4000).optional().describe('Approx output tokens. Default 700.') +} + +const findCallersInputSchema = { + name: z.string().min(1).describe('Exact definition name to resolve through the symbol index.'), + kind: kindFilterSchema.optional().describe('Target definition kind filter for disambiguation.'), + path_glob: z.string().min(1).optional().describe('Target definition path glob for disambiguation, e.g. "src/schema/**".'), + max_tokens: z.number().int().positive().max(4000).optional().describe('Approx output tokens. Default 700.') +} + +const findRefsInputSchema = { + name: z.string().min(1).describe('Exact definition name to resolve through the symbol index.'), + kind: kindFilterSchema.optional().describe('Target definition kind filter for disambiguation.'), + path_glob: z.string().min(1).optional().describe('Target definition path glob for disambiguation, e.g. "src/schema/**".'), + max_tokens: z.number().int().positive().max(4000).optional().describe('Approx output tokens. Default 700.') +} + +const findImportersInputSchema = { + file: z.string().min(1).describe('Repo-relative file path to inspect for importing edges, e.g. "src/token.ts".'), + max_tokens: z.number().int().positive().max(4000).optional().describe('Approx output tokens. Default 700.') +} + +const findImplementersInputSchema = { + name: z.string().min(1).describe('Exact interface/class name to resolve through the symbol index.'), + max_tokens: z.number().int().positive().max(4000).optional().describe('Approx output tokens. Default 700.') +} + +const impactInputSchema = { + name: z.string().min(1).describe('Exact definition name to walk through transitive callers.'), + depth: z.number().int().min(0).max(8).optional().describe('Caller depth: 0=direct callers only. Default 2.'), + kind: kindFilterSchema.optional().describe('Target definition kind filter for root disambiguation.'), + path_glob: z.string().min(1).optional().describe('Target definition path glob for root disambiguation, e.g. "src/schema/**".'), max_tokens: z.number().int().positive().max(4000).optional().describe('Approx output tokens. Default 700.') } @@ -220,6 +309,11 @@ const indexStatusInputSchema = { const searchCodeArgsSchema = z.object(searchCodeInputSchema).strict() const findSymbolArgsSchema = z.object(findSymbolInputSchema).strict() +const findCallersArgsSchema = z.object(findCallersInputSchema).strict() +const findRefsArgsSchema = z.object(findRefsInputSchema).strict() +const findImportersArgsSchema = z.object(findImportersInputSchema).strict() +const findImplementersArgsSchema = z.object(findImplementersInputSchema).strict() +const impactArgsSchema = z.object(impactInputSchema).strict() const grepCodeArgsSchema = z.object(grepCodeInputSchema).strict() const readChunkArgsSchema = z.object(readChunkInputSchema).strict() const indexStatusArgsSchema = z.object(indexStatusInputSchema).strict() @@ -256,8 +350,8 @@ class StdioMcpServerHandle implements McpServerHandle { export function getToolDefinitions(): readonly McpToolDefinition[] { const provider = getDefaultEmbeddingProvider() const searchDescription = isLearnedEmbeddingProvider(provider) - ? 'Hybrid lexical + semantic repo search for concepts/unknown names. Top body inline; read_chunk for other hits. Use grep_code for literals/regex, find_symbol for definitions.' - : 'Lexical repo search for concepts/unknown names. Top body inline; read_chunk for other hits. Use grep_code for literals/regex, find_symbol for definitions.' + ? 'Hybrid lexical + semantic search for concepts/unknown names. Top body inline; grep_code literals/regex, find_symbol defs.' + : 'Lexical search for concepts/unknown names. Top body inline; grep_code literals/regex, find_symbol defs.' return [ { @@ -277,18 +371,66 @@ export function getToolDefinitions(): readonly McpToolDefinition[] { }, { name: 'find_symbol', - description: 'Exact identifier/definition lookup. Top unambiguous body inline; use before search_code for names.', + description: 'Exact identifier lookup. Top body inline; optional bounded relations.', inputSchema: jsonSchema(['name'], { name: { type: 'string' }, kind: kindJsonSchema(), path_glob: { type: 'string' }, with_body: { type: 'boolean', default: true }, + with_callers: { type: 'boolean' }, max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS } }) }, + { + name: 'find_callers', + description: 'Callers (Go/Java/Ruby/Rust approx:name-only).', + inputSchema: jsonSchema(['name'], { + name: { type: 'string' }, + kind: kindJsonSchema(), + path_glob: { type: 'string' }, + max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: DEFAULT_MCP_RELATION_MAX_TOKENS } + }) + }, + { + name: 'find_refs', + description: 'Refs (Go/Java/Ruby/Rust approx:name-only).', + inputSchema: jsonSchema(['name'], { + name: { type: 'string' }, + kind: kindJsonSchema(), + path_glob: { type: 'string' }, + max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: DEFAULT_MCP_RELATION_MAX_TOKENS } + }) + }, + { + name: 'find_importers', + description: 'Indexed importers for a repo file.', + inputSchema: jsonSchema(['file'], { + file: { type: 'string' }, + max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: DEFAULT_MCP_RELATION_MAX_TOKENS } + }) + }, + { + name: 'who_implements', + description: 'Implementers/extends (Go/Java/Ruby/Rust approx:name-only).', + inputSchema: jsonSchema(['name'], { + name: { type: 'string' }, + max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: DEFAULT_MCP_RELATION_MAX_TOKENS } + }) + }, + { + name: 'impact', + description: 'Transitive callers/impact (Go/Java/Ruby/Rust approx:name-only).', + inputSchema: jsonSchema(['name'], { + name: { type: 'string' }, + depth: { type: 'integer', minimum: 0, maximum: 8, default: 2 }, + kind: kindJsonSchema(), + path_glob: { type: 'string' }, + max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: DEFAULT_MCP_RELATION_MAX_TOKENS } + }) + }, { name: 'grep_code', - description: 'Literal/regex search for env vars, errors, operators, exact strings.', + description: 'Literal/regex search for env vars, errors, operators, exact text.', inputSchema: jsonSchema(['pattern'], { pattern: { type: 'string' }, regex: { type: 'boolean', default: false }, @@ -306,7 +448,7 @@ export function getToolDefinitions(): readonly McpToolDefinition[] { }, { name: 'read_chunk', - description: 'Read non-top hit/wider context by id. Not needed for top search_code/find_symbol hits returned inline.', + description: 'Read non-top/wider context by id. Not needed for top search_code/find_symbol hits returned inline.', inputSchema: jsonSchema(['id'], { id: { type: 'string' }, context_lines: { type: 'integer', minimum: 0, maximum: 50 }, @@ -370,8 +512,63 @@ export function createRouter(repo: Repo): McpRouter { options.withBody = args.with_body } + if (args.with_callers !== undefined) { + options.withCallers = args.with_callers + } + return repo.findSymbol(args.name, options) }, + async findCallers(args) { + const options: FindEdgeOptions = {} + + if (args.kind) { + options.kind = args.kind + } + + if (args.path_glob) { + options.pathGlob = args.path_glob + } + + return repo.findCallers(args.name, options) + }, + async findReferences(args) { + const options: FindEdgeOptions = {} + + if (args.kind) { + options.kind = args.kind + } + + if (args.path_glob) { + options.pathGlob = args.path_glob + } + + return repo.findReferences(args.name, options) + }, + async findImporters(args) { + return repo.findImporters(args.file) + }, + async findImplementers(args) { + return repo.findImplementers(args.name) + }, + async impact(args) { + const options: ImpactOptions = { + maxTokens: args.max_tokens ?? DEFAULT_MCP_RELATION_MAX_TOKENS + } + + if (args.depth !== undefined) { + options.depth = args.depth + } + + if (args.kind) { + options.kind = args.kind + } + + if (args.path_glob) { + options.pathGlob = args.path_glob + } + + return repo.impact(args.name, options) + }, async grepCode(args) { const options: GrepOptions = {} @@ -439,6 +636,31 @@ export async function callMcpTool(repo: Repo, name: McpToolName, args: unknown): const parsed = findSymbolArgsSchema.parse(args) return formatMcpSymbols(await router.findSymbol(parsed), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS }) } + case 'find_callers': + { + const parsed = findCallersArgsSchema.parse(args) + return formatMcpCallers(await router.findCallers(parsed), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_RELATION_MAX_TOKENS }) + } + case 'find_refs': + { + const parsed = findRefsArgsSchema.parse(args) + return formatMcpReferences(await router.findReferences(parsed), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_RELATION_MAX_TOKENS }) + } + case 'find_importers': + { + const parsed = findImportersArgsSchema.parse(args) + return formatMcpImporters(await router.findImporters(parsed), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_RELATION_MAX_TOKENS }) + } + case 'who_implements': + { + const parsed = findImplementersArgsSchema.parse(args) + return formatMcpImplementers(await router.findImplementers(parsed), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_RELATION_MAX_TOKENS }) + } + case 'impact': + { + const parsed = impactArgsSchema.parse(args) + return formatMcpImpact(await router.impact(parsed), { maxTokens: parsed.max_tokens ?? DEFAULT_MCP_RELATION_MAX_TOKENS }) + } case 'grep_code': { const parsed = grepCodeArgsSchema.parse(args) @@ -524,6 +746,21 @@ export function createSdkServer(repo: Repo): McpServer { server.registerTool('find_symbol', { description: toolDescription('find_symbol'), inputSchema: findSymbolInputSchema }, async (args) => textResult(formatMcpSymbols(await router.findSymbol(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS })) ) + server.registerTool('find_callers', { description: toolDescription('find_callers'), inputSchema: findCallersInputSchema }, async (args) => + textResult(formatMcpCallers(await router.findCallers(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_RELATION_MAX_TOKENS })) + ) + server.registerTool('find_refs', { description: toolDescription('find_refs'), inputSchema: findRefsInputSchema }, async (args) => + textResult(formatMcpReferences(await router.findReferences(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_RELATION_MAX_TOKENS })) + ) + server.registerTool('find_importers', { description: toolDescription('find_importers'), inputSchema: findImportersInputSchema }, async (args) => + textResult(formatMcpImporters(await router.findImporters(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_RELATION_MAX_TOKENS })) + ) + server.registerTool('who_implements', { description: toolDescription('who_implements'), inputSchema: findImplementersInputSchema }, async (args) => + textResult(formatMcpImplementers(await router.findImplementers(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_RELATION_MAX_TOKENS })) + ) + server.registerTool('impact', { description: toolDescription('impact'), inputSchema: impactInputSchema }, async (args) => + textResult(formatMcpImpact(await router.impact(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_RELATION_MAX_TOKENS })) + ) server.registerTool('grep_code', { description: toolDescription('grep_code'), inputSchema: grepCodeInputSchema }, async (args) => textResult(formatMcpGrepHits(await router.grepCode(args), { maxTokens: args.max_tokens ?? DEFAULT_MCP_GREP_MAX_TOKENS })) ) @@ -1221,24 +1458,73 @@ export function formatMcpSymbols(definitions: SymbolDefinition[], options: Forma } function formatMcpSymbolDefinition(definition: SymbolDefinition, index: number): string { - const header = formatMcpSymbolHeader(definition, index) + const sections = [formatMcpSymbolHeader(definition, index)] // The top exact match carries a paste-ready body so the identifier resolves // in a single call; render it with the same line-numbered block as search. if (definition.body !== undefined) { - return `${header}\n${formatBodyBlock(definition.body, definition.range.startLine)}` + sections.push(formatBodyBlock(definition.body, definition.range.startLine)) + } + + const relations = formatMcpSymbolRelations(definition.relations) + if (relations) { + sections.push(relations) } - return header + + return sections.join('\n') } function formatMcpSymbolHeader(definition: SymbolDefinition, index: number): string { return `#${index + 1} ${definition.kind} ${definition.name} ${definition.file}:${formatRange(definition.range.startLine, definition.range.endLine)}` } +function formatMcpSymbolRelations(relations: SymbolDefinition['relations']): string { + if (!relations) { + return '' + } + + const itemLines = buildSymbolRelationLines(relations) + const sections = ['relations:', ...itemLines] + if (relations.omitted && relations.omitted > 0) { + sections.push(symbolRelationOmissionMarker(relations.omitted)) + } + + return itemLines.length > 0 || (relations.omitted ?? 0) > 0 ? joinSections(sections) : '' +} + +function buildSymbolRelationLines(relations: NonNullable): string[] { + return [ + ...relations.sites.map((site) => formatMcpSymbolRelationSite(site)), + ...relations.neighbors.map((neighbor) => formatMcpSymbolNeighborLine(neighbor)) + ] +} + +function formatMcpSymbolRelationSite(site: NonNullable['sites'][number]): string { + const context = site.srcSymbol ?? 'top-level' + const resolution = site.resolution === 'name-only' ? 'approx:name-only' : site.resolution + return `- ${site.edgeKind} ${site.file}:${site.line} ${context} ${resolution}` +} + +function formatMcpSymbolNeighborLine(neighbor: NonNullable['neighbors'][number]): string { + const symbolPath = [neighbor.parent, neighbor.name].filter(Boolean).join(' > ') + return `- neighbor ${neighbor.kind} ${symbolPath} ${formatRange(neighbor.range.startLine, neighbor.range.endLine)}` +} + +function countSymbolRelationItems(relations: SymbolDefinition['relations']): number { + if (!relations) { + return 0 + } + + return relations.sites.length + relations.neighbors.length + (relations.omitted ?? 0) +} + function fitSymbolOutputForBudget(definitions: SymbolDefinition[], rendered: string[], maxChars: number): string { const omittedAfterFirst = definitions.length - 1 + const firstHasRelations = countSymbolRelationItems(definitions[0]?.relations) > 0 const firstOnlyMarker = omittedAfterFirst > 0 ? symbolOmissionMarker(omittedAfterFirst, Math.max(0, maxChars - 2)) - : symbolBodyTruncationMarker(Math.max(0, maxChars - 2)) + : firstHasRelations + ? undefined + : symbolBodyTruncationMarker(Math.max(0, maxChars - 2)) const firstOnlyBudget = maxCharsForFirstHit(maxChars, firstOnlyMarker) if (rendered[0]!.length > firstOnlyBudget) { const first = formatMcpSymbolDefinitionForBudget(definitions[0]!, 0, firstOnlyBudget) @@ -1268,13 +1554,97 @@ function fitSymbolOutputForBudget(definitions: SymbolDefinition[], rendered: str function formatMcpSymbolDefinitionForBudget(definition: SymbolDefinition, index: number, maxChars: number): string { const header = formatMcpSymbolHeader(definition, index) - if (definition.body === undefined || maxChars <= header.length + 1) { + if (maxChars <= header.length) { return truncateWithEllipsis(header, maxChars) } - const bodyBudget = maxChars - header.length - 1 - const body = truncateSymbolBodyBlock(formatBodyBlock(definition.body, definition.range.startLine), bodyBudget) - return body ? `${header}\n${body}` : truncateWithEllipsis(header, maxChars) + const sections = [header] + const relationItemCount = countSymbolRelationItems(definition.relations) + if (definition.body !== undefined) { + const relationReserve = relationItemCount > 0 ? minimumSymbolRelationBudget(definition.relations) + 1 : 0 + const bodyBudget = Math.max(0, maxChars - header.length - 1 - relationReserve) + const body = truncateSymbolBodyBlock(formatBodyBlock(definition.body, definition.range.startLine), bodyBudget) + if (body) { + sections.push(body) + } + } + + if (relationItemCount > 0 && definition.relations) { + const remainingForRelations = Math.max(0, maxChars - joinSections(sections).length - 1) + const relations = formatMcpSymbolRelationsForBudget(definition.relations, remainingForRelations) + if (relations) { + sections.push(relations) + } + } + + const output = joinSections(sections) + return output.length <= maxChars ? output : truncateWithEllipsis(output, maxChars) +} + +function minimumSymbolRelationBudget(relations: SymbolDefinition['relations']): number { + const count = countSymbolRelationItems(relations) + return count > 0 ? symbolRelationOmissionMarker(count).length : 0 +} + +function formatMcpSymbolRelationsForBudget(relations: NonNullable, maxChars: number): string { + if (maxChars <= 0) { + return '' + } + + const itemLines = buildSymbolRelationLines(relations) + const baseOmitted = relations.omitted ?? 0 + const full = joinSections([ + 'relations:', + ...itemLines, + baseOmitted > 0 ? symbolRelationOmissionMarker(baseOmitted) : '' + ]) + if (full && full.length <= maxChars) { + return full + } + + if (itemLines.length === 0) { + return bestEffortOmissionMarker(symbolRelationOmissionMarkerVariants(baseOmitted), maxChars) + } + + const header = 'relations:' + if (header.length > maxChars) { + return bestEffortOmissionMarker(symbolRelationOmissionMarkerVariants(baseOmitted + itemLines.length), maxChars) + } + + const kept = [header] + let keptItems = 0 + for (let index = 0; index < itemLines.length; index += 1) { + const line = itemLines[index]! + const omittedAfter = baseOmitted + itemLines.length - (keptItems + 1) + const marker = omittedAfter > 0 + ? symbolRelationOmissionMarker(omittedAfter, remainingCharsForTrailingLine([...kept, line], maxChars)) + : '' + const candidate = joinSections([...kept, line, marker]) + if (candidate.length <= maxChars) { + kept.push(line) + keptItems += 1 + continue + } + break + } + + const omitted = baseOmitted + itemLines.length - keptItems + if (omitted <= 0) { + return joinSections(kept) + } + + let marker = symbolRelationOmissionMarker(omitted, remainingCharsForTrailingLine(kept, maxChars)) + while (!marker && kept.length > 1) { + kept.pop() + keptItems -= 1 + marker = symbolRelationOmissionMarker(baseOmitted + itemLines.length - keptItems, remainingCharsForTrailingLine(kept, maxChars)) + } + + if (!marker) { + return bestEffortOmissionMarker(symbolRelationOmissionMarkerVariants(baseOmitted + itemLines.length), maxChars) + } + + return joinSections([...kept, marker]) } function truncateSymbolBodyBlock(body: string, maxChars: number): string { @@ -1303,6 +1673,17 @@ function symbolOmissionMarker(omitted: number, maxChars?: number): string { return bare.length <= maxChars ? bare : '' } +function symbolRelationOmissionMarker(omitted: number, maxChars?: number): string { + return firstMarkerThatFits(symbolRelationOmissionMarkerVariants(omitted), maxChars) +} + +function symbolRelationOmissionMarkerVariants(omitted: number): string[] { + return [ + `relations_omitted=${omitted}; raise max_tokens`, + `relations_omitted=${omitted}` + ] +} + function symbolBodyTruncationMarker(maxChars?: number): string { const full = 'symbol_body_truncated=true; refine name/kind/path_glob or raise max_tokens.' if (maxChars === undefined || full.length <= maxChars) { @@ -1318,6 +1699,253 @@ function symbolBodyTruncationMarker(maxChars?: number): string { return bare.length <= maxChars ? bare : '' } +export function formatMcpCallers(results: EdgeResult[], options: FormatMcpEdgeResultsOptions = {}): string { + return formatMcpEdgeResults(results, 'no_callers', 'callers', options) +} + +export function formatMcpReferences(results: EdgeResult[], options: FormatMcpEdgeResultsOptions = {}): string { + return formatMcpEdgeResults(results, 'no_refs', 'refs', options) +} + +export function formatMcpImporters(results: EdgeResult[], options: FormatMcpEdgeResultsOptions = {}): string { + return formatMcpEdgeResults(results, 'no_importers', 'importers', options) +} + +export function formatMcpImplementers(results: EdgeResult[], options: FormatMcpEdgeResultsOptions = {}): string { + return formatMcpEdgeResults(results, 'no_implementers', 'implementers', options) +} + +export function formatMcpImpact(result: ImpactResult, options: FormatMcpImpactOptions = {}): string { + if (result.nodes.length === 0) { + return 'no_impact' + } + + const rendered = result.nodes.map((node) => formatMcpImpactNode(node)) + const notes = buildImpactNoteLines(result) + const output = joinSections([...rendered, ...notes]) + const maxTokens = options.maxTokens + if (maxTokens === undefined || output.length <= maxTokens * 4) { + return output + } + + return fitImpactOutputForBudget(rendered, notes, maxTokens * 4) +} + +function buildImpactNoteLines(result: ImpactResult): string[] { + const notes: string[] = [] + if (result.depthCapped) { + notes.push(`depth_capped=${result.depthLimit}`) + } + if (result.nodesCapped) { + notes.push(`nodes_capped=${result.maxNodes}`) + } + if (result.impactTruncated) { + notes.push('impact_truncated=true') + } + return notes +} + +function formatMcpImpactNode(node: ImpactResult['nodes'][number], maxChars?: number): string { + const context = node.srcSymbol ?? node.name + const resolution = node.resolution === 'name-only' ? 'approx:name-only' : node.resolution + const prefix = `${node.file}:${node.line} ${context} d${node.depth} ${node.edgeKind} ${resolution} | ` + const snippet = compactStatusValue(node.snippet, 240) + + if (maxChars === undefined || prefix.length + snippet.length <= maxChars) { + return `${prefix}${snippet}` + } + + if (prefix.length >= maxChars) { + return truncateWithEllipsis(prefix.trimEnd(), maxChars) + } + + return `${prefix}${truncateWithEllipsis(snippet, maxChars - prefix.length)}` +} + +function fitImpactOutputForBudget(rendered: string[], notes: string[], maxChars: number): string { + const kept: string[] = [] + const requiredNotes = notes.includes('impact_truncated=true') ? notes : [...notes, 'impact_truncated=true'] + const requiredNoteText = joinSections(requiredNotes) + + for (let index = 0; index < rendered.length; index += 1) { + const line = rendered[index]! + const candidate = joinSections([...kept, line, requiredNoteText]) + if (candidate.length <= maxChars) { + kept.push(line) + continue + } + + if (kept.length === 0) { + const lineBudget = Math.max(0, maxChars - requiredNoteText.length - 1) + const fittedLine = lineBudget > 0 ? formatMcpImpactNodeLineForBudget(line, lineBudget) : '' + return joinSections([fittedLine, requiredNoteText]) || bestEffortOmissionMarker(['impact_truncated=true'], maxChars) + } + + break + } + + return fitImpactNotesWithLeadingSections(kept, requiredNotes, maxChars) +} + +function fitImpactNotesWithLeadingSections(leadingSections: string[], notes: string[], maxChars: number): string { + const leading = joinSections(leadingSections) + for (let count = notes.length; count >= 1; count -= 1) { + const candidate = joinSections([...leadingSections, ...notes.slice(0, count)]) + if (candidate.length <= maxChars) { + return candidate + } + } + + const remaining = maxChars - (leading ? leading.length + 1 : 0) + if (remaining > 0) { + const note = bestEffortOmissionMarker(['impact_truncated=true'], remaining) + if (note) { + return joinSections([...leadingSections, note]) + } + } + + return leading ? truncateWithEllipsis(leading, maxChars) : bestEffortOmissionMarker(['impact_truncated=true'], maxChars) +} + +function formatMcpImpactNodeLineForBudget(line: string, maxChars: number): string { + const separator = ' | ' + const separatorIndex = line.indexOf(separator) + if (separatorIndex < 0 || line.length <= maxChars) { + return truncateWithEllipsis(line, maxChars) + } + + const prefix = line.slice(0, separatorIndex + separator.length) + const snippet = line.slice(separatorIndex + separator.length) + if (prefix.length >= maxChars) { + return truncateWithEllipsis(prefix.trimEnd(), maxChars) + } + + return `${prefix}${truncateWithEllipsis(snippet, maxChars - prefix.length)}` +} + +function formatMcpEdgeResults( + results: EdgeResult[], + emptyMarker: string, + omissionLabel: 'callers' | 'refs' | 'importers' | 'implementers', + options: FormatMcpEdgeResultsOptions +): string { + if (results.length === 0) { + return emptyMarker + } + + const rendered = results.map((result) => formatMcpEdgeResult(result)) + const output = rendered.join('\n') + const maxTokens = options.maxTokens + if (maxTokens === undefined || output.length <= maxTokens * 4) { + return output + } + + return fitEdgeResultOutputForBudget(rendered, omissionLabel, maxTokens * 4) +} + +function formatMcpEdgeResult(result: EdgeResult, maxChars?: number): string { + const context = result.srcSymbol ?? 'top-level' + const resolution = result.resolution === 'name-only' ? 'approx:name-only' : result.resolution + const prefix = `${result.file}:${result.line} ${context} ${result.edgeKind} ${resolution} | ` + const snippet = compactStatusValue(result.snippet, 240) + + if (maxChars === undefined || prefix.length + snippet.length <= maxChars) { + return `${prefix}${snippet}` + } + + if (prefix.length >= maxChars) { + return truncateWithEllipsis(prefix.trimEnd(), maxChars) + } + + return `${prefix}${truncateWithEllipsis(snippet, maxChars - prefix.length)}` +} + +function fitEdgeResultOutputForBudget( + rendered: string[], + omissionLabel: 'callers' | 'refs' | 'importers' | 'implementers', + maxChars: number +): string { + const kept: string[] = [] + + for (let index = 0; index < rendered.length; index += 1) { + const line = rendered[index]! + const omittedAfter = rendered.length - index - 1 + const marker = omittedAfter > 0 + ? edgeResultOmissionMarker(omissionLabel, omittedAfter, remainingCharsForTrailingLine([...kept, line], maxChars)) + : '' + const candidate = joinSections([...kept, line, marker]) + if (candidate.length <= maxChars) { + kept.push(line) + continue + } + + if (kept.length === 0) { + return formatSingleEdgeResultForBudget(line, omissionLabel, omittedAfter, maxChars) + } + + break + } + + const omitted = rendered.length - kept.length + if (omitted <= 0) { + return joinSections(kept) + } + + let marker = edgeResultOmissionMarker(omissionLabel, omitted, remainingCharsForTrailingLine(kept, maxChars)) + while (!marker && kept.length > 0) { + kept.pop() + marker = edgeResultOmissionMarker(omissionLabel, rendered.length - kept.length, remainingCharsForTrailingLine(kept, maxChars)) + } + + if (!marker) { + return bestEffortOmissionMarker(edgeResultOmissionMarkerVariants(omissionLabel, omitted), maxChars) + } + + return joinSections([...kept, marker]) +} + +function formatSingleEdgeResultForBudget( + line: string, + omissionLabel: 'callers' | 'refs' | 'importers' | 'implementers', + omittedAfter: number, + maxChars: number +): string { + if (omittedAfter <= 0) { + return formatMcpEdgeResultLineForBudget(line, maxChars) + } + + for (const marker of edgeResultOmissionMarkerVariants(omissionLabel, omittedAfter)) { + const lineBudget = maxChars - marker.length - 1 + if (lineBudget <= 0) { + continue + } + + const truncatedLine = formatMcpEdgeResultLineForBudget(line, lineBudget) + const candidate = joinSections([truncatedLine, marker]) + if (candidate.length <= maxChars) { + return candidate + } + } + + return bestEffortOmissionMarker(edgeResultOmissionMarkerVariants(omissionLabel, omittedAfter), maxChars) +} + +function formatMcpEdgeResultLineForBudget(line: string, maxChars: number): string { + const separator = ' | ' + const separatorIndex = line.indexOf(separator) + if (separatorIndex < 0 || line.length <= maxChars) { + return truncateWithEllipsis(line, maxChars) + } + + const prefix = line.slice(0, separatorIndex + separator.length) + const snippet = line.slice(separatorIndex + separator.length) + if (prefix.length >= maxChars) { + return truncateWithEllipsis(prefix.trimEnd(), maxChars) + } + + return `${prefix}${truncateWithEllipsis(snippet, maxChars - prefix.length)}` +} + export function formatMcpGrepHits(hits: GrepHit[], options: FormatMcpGrepHitsOptions = {}): string { if (hits.length === 0) { return 'no_matches' @@ -1681,6 +2309,24 @@ function usageOmissionMarkerVariants(omitted: number): string[] { ] } +function edgeResultOmissionMarker( + label: 'callers' | 'refs' | 'importers' | 'implementers', + omitted: number, + maxChars?: number +): string { + return firstMarkerThatFits(edgeResultOmissionMarkerVariants(label, omitted), maxChars) +} + +function edgeResultOmissionMarkerVariants( + label: 'callers' | 'refs' | 'importers' | 'implementers', + omitted: number +): string[] { + return [ + `${label}_omitted=${omitted}; raise max_tokens`, + `${label}_omitted=${omitted}` + ] +} + function firstMarkerThatFits(markers: string[], maxChars?: number): string { if (maxChars === undefined) { return markers[0] ?? '' diff --git a/packages/mcp/test/mcp.test.ts b/packages/mcp/test/mcp.test.ts index e371816..c57e33c 100644 --- a/packages/mcp/test/mcp.test.ts +++ b/packages/mcp/test/mcp.test.ts @@ -1,15 +1,16 @@ import { spawn, type ChildProcessWithoutNullStreams } from 'node:child_process' -import { mkdtemp, mkdir, rm, writeFile } from 'node:fs/promises' +import { cp, mkdtemp, mkdir, rm, writeFile } from 'node:fs/promises' import { tmpdir } from 'node:os' import { join } from 'node:path' import { afterEach, describe, expect, it } from 'vitest' -import { openRepo, registerEmbeddingProvider, type GrepHit, type RepoStatus, type SearchHit, type SymbolDefinition } from '@codesift/core' +import { openRepo, registerEmbeddingProvider, type EdgeResult, type GrepHit, type ImpactResult, type RepoStatus, type SearchHit, type SymbolDefinition } from '@codesift/core' import { DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS, DEFAULT_MCP_INDEX_STATUS_MAX_TOKENS, + DEFAULT_MCP_RELATION_MAX_TOKENS, DEFAULT_MCP_SEARCH_MAX_TOKENS, DEFAULT_MCP_READ_CHUNK_MAX_TOKENS, DEFAULT_SEARCH_K, @@ -19,9 +20,14 @@ import { createRouter, createStdioServer, callMcpTool, + formatMcpCallers, formatMcpGrepHits, + formatMcpImpact, + formatMcpImplementers, + formatMcpImporters, formatMcpIndexStatus, formatMcpReadChunk, + formatMcpReferences, formatMcpSearchHits, formatMcpSymbols, getToolDefinitions @@ -44,11 +50,24 @@ afterEach(async () => { ) }) +async function copyFixtureRepository(fixtureName: string, prefix: string): Promise { + const parentDirectory = await mkdtemp(join(tmpdir(), prefix)) + const repoRoot = join(parentDirectory, 'repo') + temporaryDirectories.push(parentDirectory) + await cp(join(process.cwd(), 'packages', 'eval', 'fixtures', fixtureName), repoRoot, { recursive: true }) + return repoRoot +} + describe('@codesift/mcp server', () => { it('exposes the planned tool surface with routing schemas', () => { expect(getToolDefinitions().map((tool) => tool.name)).toEqual([ 'search_code', 'find_symbol', + 'find_callers', + 'find_refs', + 'find_importers', + 'who_implements', + 'impact', 'grep_code', 'read_chunk', 'index_status' @@ -58,8 +77,17 @@ describe('@codesift/mcp server', () => { max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: 700 } }) expect(getToolDefinitions().find((tool) => tool.name === 'find_symbol')?.inputSchema.properties).toMatchObject({ + with_callers: { type: 'boolean' }, max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: DEFAULT_MCP_FIND_SYMBOL_MAX_TOKENS } }) + expect(getToolDefinitions().find((tool) => tool.name === 'find_callers')?.inputSchema.required).toEqual(['name']) + expect(getToolDefinitions().find((tool) => tool.name === 'find_callers')?.inputSchema.properties).toMatchObject({ + max_tokens: { type: 'integer', minimum: 1, maximum: 4000, default: DEFAULT_MCP_RELATION_MAX_TOKENS } + }) + expect(getToolDefinitions().find((tool) => tool.name === 'find_refs')?.inputSchema.required).toEqual(['name']) + expect(getToolDefinitions().find((tool) => tool.name === 'find_importers')?.inputSchema.required).toEqual(['file']) + expect(getToolDefinitions().find((tool) => tool.name === 'who_implements')?.inputSchema.required).toEqual(['name']) + expect(getToolDefinitions().find((tool) => tool.name === 'impact')?.inputSchema.required).toEqual(['name']) expect(getToolDefinitions().find((tool) => tool.name === 'read_chunk')?.inputSchema.properties).toMatchObject({ max_tokens: { type: 'integer', minimum: MIN_MCP_READ_CHUNK_MAX_TOKENS, maximum: 4000, default: DEFAULT_MCP_READ_CHUNK_MAX_TOKENS } }) @@ -71,29 +99,48 @@ describe('@codesift/mcp server', () => { expect(getToolDefinitions().find((tool) => tool.name === 'index_status')?.inputSchema.properties).toMatchObject({ max_tokens: { type: 'integer', minimum: 1, maximum: 1000, default: DEFAULT_MCP_INDEX_STATUS_MAX_TOKENS } }) - expect(MCP_SERVER_INSTRUCTIONS).toContain('literals/env/errors/operators/regex->grep_code') + expect(MCP_SERVER_INSTRUCTIONS).toContain('breakage/transitive callers->impact') + expect(MCP_SERVER_INSTRUCTIONS).toContain('approx:name-only') + expect(MCP_SERVER_INSTRUCTIONS).toContain('literals/regex/errors/env/operators->grep_code') }) it('asserts single-call sufficiency in instructions and tool descriptions', () => { const lowerInstructions = MCP_SERVER_INSTRUCTIONS.toLowerCase() - expect(lowerInstructions).toContain('top search_code body is inline') - expect(lowerInstructions).toContain('read_chunk only for non-top hits/wider context') + expect(lowerInstructions).toContain('top search_code body inline') + expect(lowerInstructions).toContain('read_chunk only for non-top/wider context') const tools = getToolDefinitions() const searchDescription = tools.find((tool) => tool.name === 'search_code')?.description ?? '' expect(searchDescription.toLowerCase()).toContain('top body inline') const findDescription = tools.find((tool) => tool.name === 'find_symbol')?.description ?? '' - expect(findDescription.toLowerCase()).toContain('top unambiguous body inline') + expect(findDescription.toLowerCase()).toContain('top body inline') + expect(findDescription.toLowerCase()).toContain('relations') + + const callersDescription = tools.find((tool) => tool.name === 'find_callers')?.description ?? '' + expect(callersDescription).toContain('approx:name-only') + + const refsDescription = tools.find((tool) => tool.name === 'find_refs')?.description ?? '' + expect(refsDescription).toContain('approx:name-only') + + const implementersDescription = tools.find((tool) => tool.name === 'who_implements')?.description ?? '' + expect(implementersDescription).toContain('approx:name-only') + + const impactDescription = tools.find((tool) => tool.name === 'impact')?.description ?? '' + expect(impactDescription).toContain('approx:name-only') const readDescription = tools.find((tool) => tool.name === 'read_chunk')?.description ?? '' expect(readDescription.toLowerCase()).toContain('not needed for top search_code/find_symbol hits') expect(readDescription.toLowerCase()).toContain('returned inline') - // The routing guidance for find_symbol/grep_code is preserved. - expect(MCP_SERVER_INSTRUCTIONS).toContain('identifiers/definitions->find_symbol') - expect(MCP_SERVER_INSTRUCTIONS).toContain('concepts/unknown names->search_code') - expect(MCP_SERVER_INSTRUCTIONS).toContain('Broad search_code: k=5-8') + // The routing guidance for find_symbol/graph/grep_code is preserved. + expect(MCP_SERVER_INSTRUCTIONS).toContain('identifiers->find_symbol') + expect(MCP_SERVER_INSTRUCTIONS).toContain('callers/uses->find_callers/find_refs') + expect(MCP_SERVER_INSTRUCTIONS).toContain('importers->find_importers') + expect(MCP_SERVER_INSTRUCTIONS).toContain('implements/extends->who_implements') + expect(MCP_SERVER_INSTRUCTIONS).toContain('breakage/transitive callers->impact') + expect(MCP_SERVER_INSTRUCTIONS).toContain('concepts/fuzzy names->search_code') + expect(MCP_SERVER_INSTRUCTIONS).toContain('Broad search k=5-8') expect(MCP_SERVER_INSTRUCTIONS).toContain('Check index_status') }) @@ -103,13 +150,18 @@ describe('@codesift/mcp server', () => { const toolDescriptionCeilings = new Map([ ['search_code', 190], ['find_symbol', 130], + ['find_callers', 120], + ['find_refs', 120], + ['find_importers', 90], + ['who_implements', 125], + ['impact', 110], ['grep_code', 120], ['read_chunk', 120], ['index_status', 90] ]) - expect(payload.length).toBeLessThanOrEqual(3500) - expect(Math.ceil(payload.length / 4)).toBeLessThanOrEqual(875) + expect(payload.length).toBeLessThanOrEqual(6000) + expect(Math.ceil(payload.length / 4)).toBeLessThanOrEqual(1500) expect(MCP_SERVER_INSTRUCTIONS.length).toBeLessThanOrEqual(430) for (const tool of tools) { expect(tool.description.length).toBeLessThanOrEqual(toolDescriptionCeilings.get(tool.name) ?? 120) @@ -179,6 +231,187 @@ describe('@codesift/mcp server', () => { expect(DEFAULT_SEARCH_K).toBe(8) }) + it('routes graph tool calls through the core repo contract', async () => { + const repoRoot = await copyFixtureRepository('usages-ts', 'codesift-mcp-graph-router-') + const repo = await openRepo(repoRoot) + + await repo.sync() + const router = createRouter(repo) + const callers = await router.findCallers({ name: 'parseToken' }) + const refs = await router.findReferences({ name: 'parseToken' }) + const importers = await router.findImporters({ file: 'src/token.ts' }) + + expect(callers.map((result) => `${result.file}:${result.line}:${result.srcSymbol}:${result.edgeKind}`)).toEqual([ + 'src/api.ts:4:readSubject:call', + 'src/worker.ts:4:enqueueToken:call' + ]) + expect(refs.map((result) => `${result.file}:${result.line}:${result.srcSymbol}:${result.edgeKind}`)).toEqual([ + 'src/api.ts:4:readSubject:call', + 'src/worker.ts:4:enqueueToken:call' + ]) + expect(importers.map((result) => `${result.file}:${result.line}:${result.edgeKind}`)).toEqual([ + 'src/api.ts:1:import', + 'src/worker.ts:1:import' + ]) + }) + + it('returns one-call find_symbol relations and bounded impact through the MCP contract', async () => { + const repoRoot = await copyFixtureRepository('usages-ts', 'codesift-mcp-relations-impact-') + await writeFile( + join(repoRoot, 'src', 'service.ts'), + `import { readSubject } from './api' + +export function handleToken(header: string): string { + return readSubject(header) +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'src', 'app.ts'), + `import { handleToken } from './service' + +export function runApp(header: string): string { + return handleToken(header) +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'src', 'cli.ts'), + `import { runApp } from './app' + +export function main(header: string): string { + return runApp(header) +} +`, + 'utf8' + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + const router = createRouter(repo) + + const symbols = await router.findSymbol({ name: 'parseToken', with_callers: true }) + expect(symbols[0]?.body).toContain('export function parseToken') + expect(symbols[0]?.relations?.sites.map((site) => `${site.file}:${site.line}:${site.srcSymbol}:${site.edgeKind}`)).toEqual([ + 'src/api.ts:4:readSubject:call', + 'src/worker.ts:4:enqueueToken:call' + ]) + + const impact = await router.impact({ name: 'parseToken', depth: 2 }) + expect(impact.nodes.map((node) => `${node.depth}:${node.file}:${node.line}:${node.srcSymbol}:${node.edgeKind}`)).toEqual([ + '0:src/api.ts:4:readSubject:call', + '0:src/worker.ts:4:enqueueToken:call', + '1:src/service.ts:4:handleToken:call', + '2:src/app.ts:4:runApp:call' + ]) + expect(impact.depthCapped).toBe(true) + }) + + it('routes implementer lookups through the core repo contract', async () => { + const repoRoot = await copyFixtureRepository('heritage-ts', 'codesift-mcp-heritage-router-') + const repo = await openRepo(repoRoot) + + await repo.sync() + const router = createRouter(repo) + const implementers = await router.findImplementers({ name: 'AuthStrategy' }) + + expect(implementers.map((result) => `${result.file}:${result.line}:${result.srcSymbol}:${result.edgeKind}`)).toEqual([ + 'src/impl.ts:3:JwtVerifier:implements', + 'src/impl.ts:9:StrictStrategy:extends' + ]) + }) + + it('formats graph MCP tools with honest resolution labels and definition-path disambiguation', async () => { + const usagesRepoRoot = await copyFixtureRepository('usages-ts', 'codesift-mcp-graph-tools-') + const usagesRepo = await openRepo(usagesRepoRoot) + await usagesRepo.sync() + + const callersOutput = await callMcpTool(usagesRepo, 'find_callers', { name: 'parseToken', max_tokens: 80 }) + expect(callersOutput).toContain('src/api.ts:4 readSubject call import-resolved |') + expect(callersOutput).toContain('src/worker.ts:4 enqueueToken call import-resolved |') + + const importersOutput = await callMcpTool(usagesRepo, 'find_importers', { file: 'src/token.ts', max_tokens: 80 }) + expect(importersOutput).toContain("src/api.ts:1 top-level import import-resolved | import { parseToken } from './token'") + expect(importersOutput).toContain("src/worker.ts:1 top-level import import-resolved | import { parseToken } from './token'") + + const collisionRepoRoot = await copyFixtureRepository('collision-ts', 'codesift-mcp-graph-collision-') + const collisionRepo = await openRepo(collisionRepoRoot) + await collisionRepo.sync() + + const refsOutput = await callMcpTool(collisionRepo, 'find_refs', { + name: 'validate', + kind: 'function', + path_glob: 'src/schema/**', + max_tokens: 80 + }) + expect(refsOutput).toContain('src/api/handler.ts:6 handleRequest call import-resolved |') + expect(refsOutput).not.toContain('src/auth/token.ts') + expect(refsOutput).not.toContain('src/forms/checkout.ts') + + const heritageRepoRoot = await copyFixtureRepository('heritage-ts', 'codesift-mcp-graph-heritage-') + const heritageRepo = await openRepo(heritageRepoRoot) + await heritageRepo.sync() + + const implementersOutput = await callMcpTool(heritageRepo, 'who_implements', { + name: 'AuthStrategy', + max_tokens: 80 + }) + expect(implementersOutput).toContain('src/impl.ts:3 JwtVerifier implements import-resolved | export class JwtVerifier extends BaseVerifier implements AuthStrategy {') + expect(implementersOutput).toContain('src/impl.ts:9 StrictStrategy extends import-resolved | export interface StrictStrategy extends AuthStrategy {}') + }) + + it('formats one-call find_symbol relations and bounded impact output', async () => { + const repoRoot = await copyFixtureRepository('usages-ts', 'codesift-mcp-relations-output-') + await writeFile( + join(repoRoot, 'src', 'service.ts'), + `import { readSubject } from './api' + +export function handleToken(header: string): string { + return readSubject(header) +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'src', 'app.ts'), + `import { handleToken } from './service' + +export function runApp(header: string): string { + return handleToken(header) +} +`, + 'utf8' + ) + await writeFile( + join(repoRoot, 'src', 'cli.ts'), + `import { runApp } from './app' + +export function main(header: string): string { + return runApp(header) +} +`, + 'utf8' + ) + + const repo = await openRepo(repoRoot) + await repo.sync() + + const symbolOutput = await callMcpTool(repo, 'find_symbol', { name: 'parseToken', with_callers: true, max_tokens: 120 }) + expect(symbolOutput).toContain('#1 function parseToken src/token.ts:6-9') + expect(symbolOutput).toContain('relations:') + expect(symbolOutput).toContain('- call src/api.ts:4 readSubject import-resolved') + expect(symbolOutput).toContain('- call src/worker.ts:4 enqueueToken import-resolved') + expect(symbolOutput).toContain('- neighbor interface ParsedToken 1-4') + + const impactOutput = await callMcpTool(repo, 'impact', { name: 'parseToken', depth: 2, max_tokens: 180 }) + expect(impactOutput).toContain('src/api.ts:4 readSubject d0 call import-resolved |') + expect(impactOutput).toContain('src/service.ts:4 handleToken d1 call import-resolved |') + expect(impactOutput).toContain('src/app.ts:4 runApp d2 call import-resolved |') + expect(impactOutput).toContain('depth_capped=2') + }) + it('budgets find_symbol output through the MCP call surface', async () => { const repoRoot = await mkdtemp(join(tmpdir(), 'codesift-mcp-symbol-budget-')) temporaryDirectories.push(repoRoot) @@ -391,6 +624,51 @@ function makeSymbol(overrides: Partial> & { body? return definition } +function makeEdgeResult(overrides: Partial> & { srcSymbol?: string | undefined } = {}): EdgeResult { + const { srcSymbol, ...rest } = overrides + const result: EdgeResult = { + file: 'src/auth.ts', + range: { startLine: 10, endLine: 10 }, + line: 10, + snippet: ' return verifyToken(token)', + srcSymbol: 'handleRequest', + edgeKind: 'call', + resolution: 'import-resolved', + ...rest + } + + if ('srcSymbol' in overrides) { + if (srcSymbol === undefined) { + delete result.srcSymbol + } else { + result.srcSymbol = srcSymbol + } + } + + return result +} + +function makeImpactResult(overrides: Partial = {}): ImpactResult { + return { + nodes: [ + { + name: 'handleRequest', + file: 'src/auth.ts', + range: { startLine: 10, endLine: 10 }, + line: 10, + snippet: ' return verifyToken(token)', + srcSymbol: 'handleRequest', + depth: 1, + edgeKind: 'call', + resolution: 'import-resolved' + } + ], + depthLimit: 2, + maxNodes: 50, + ...overrides + } +} + function makeStatus(overrides: Partial = {}): RepoStatus { return { root: '/tmp/repo', @@ -530,6 +808,107 @@ describe('formatMcpIndexStatus compact output', () => { }) }) +describe('formatMcpCallers/References/Importers/Implementers budgeted output', () => { + it('preserves empty markers and under-budget edge rows byte-for-byte', () => { + const results = [makeEdgeResult()] + const expected = 'src/auth.ts:10 handleRequest call import-resolved | return verifyToken(token)' + + expect(formatMcpCallers([])).toBe('no_callers') + expect(formatMcpReferences([])).toBe('no_refs') + expect(formatMcpImporters([])).toBe('no_importers') + expect(formatMcpImplementers([])).toBe('no_implementers') + expect(formatMcpCallers(results)).toBe(expected) + expect(formatMcpReferences(results)).toBe(expected) + expect(formatMcpImporters([makeEdgeResult({ srcSymbol: undefined, edgeKind: 'import', snippet: "import { verifyToken } from './auth'" })])).toBe( + "src/auth.ts:10 top-level import import-resolved | import { verifyToken } from './auth'" + ) + expect(formatMcpImplementers([makeEdgeResult({ srcSymbol: 'JwtVerifier', edgeKind: 'implements', snippet: 'class JwtVerifier extends BaseVerifier implements AuthStrategy {' })])).toBe( + 'src/auth.ts:10 JwtVerifier implements import-resolved | class JwtVerifier extends BaseVerifier implements AuthStrategy {' + ) + }) + + it('marks name-only rows as approximate and emits omission markers when truncated', () => { + const results = Array.from({ length: 5 }, (_, index) => makeEdgeResult({ + file: `src/callers/${index}.ts`, + line: index + 10, + range: { startLine: index + 10, endLine: index + 10 }, + resolution: index === 0 ? 'name-only' : 'import-resolved', + snippet: ` return verifyToken(token, '${'x'.repeat(90)}')` + })) + + const output = formatMcpCallers(results, { maxTokens: 35 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(35) + expect(output).toContain('approx:name-only') + expect(output).toContain('callers_omitted=') + expect(output).not.toContain('src/callers/4.ts') + }) + + it('uses the tool-specific omission marker for tiny budgets', () => { + const output = formatMcpImporters([ + makeEdgeResult({ srcSymbol: undefined, edgeKind: 'import', snippet: `import { verifyToken } from './${'deep/'.repeat(8)}auth'` }), + makeEdgeResult({ file: 'src/other.ts', srcSymbol: undefined, edgeKind: 'import' }) + ], { maxTokens: 8 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(8) + expect(output).toContain('importers_omitted=1') + expect(output).not.toContain('src/other.ts') + }) + + it('uses the implementers omission marker for tight budgets', () => { + const output = formatMcpImplementers([ + makeEdgeResult({ edgeKind: 'implements', snippet: `class JwtVerifier implements ${'AuthStrategy'.repeat(12)}` }), + makeEdgeResult({ file: 'src/other.ts', line: 11, range: { startLine: 11, endLine: 11 }, edgeKind: 'extends', srcSymbol: 'StrictStrategy' }) + ], { maxTokens: 10 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(10) + expect(output).toContain('implementers_omitted=1') + expect(output).not.toContain('src/other.ts') + }) +}) + +describe('formatMcpImpact budgeted output', () => { + it('renders bounded impact notes and keeps a truncation marker under tight budgets', () => { + const output = formatMcpImpact(makeImpactResult({ + nodes: [ + { + name: 'handleRequest', + file: 'src/auth.ts', + range: { startLine: 10, endLine: 10 }, + line: 10, + snippet: ` return verifyToken(token, '${'x'.repeat(120)}')`, + srcSymbol: 'handleRequest', + depth: 2, + edgeKind: 'call', + resolution: 'name-only' + }, + { + name: 'runApp', + file: 'src/app.ts', + range: { startLine: 4, endLine: 4 }, + line: 4, + snippet: ' return handleRequest(token)', + srcSymbol: 'runApp', + depth: 3, + edgeKind: 'call', + resolution: 'import-resolved' + } + ], + depthCapped: true, + nodesCapped: true, + impactTruncated: true, + depthLimit: 2, + maxNodes: 50 + }), { maxTokens: 28 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(28) + expect(output).toContain('approx:name-only') + expect(output).toContain('impact_truncated=true') + expect(output).toContain('depth_capped=2') + expect(output).toContain('nodes_capped=50') + }) +}) + describe('formatMcpSymbols budgeted output', () => { it('preserves no_symbols and under-budget output byte-for-byte', () => { const definitions = [makeSymbol()] @@ -585,6 +964,29 @@ describe('formatMcpSymbols budgeted output', () => { expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(5) expect(output).toContain('symbols_omitted=1') }) + + it('adds a bounded relations block and relation omission marker when requested output is tight', () => { + const output = formatMcpSymbols([ + makeSymbol({ + relations: { + sites: Array.from({ length: 5 }, (_, index) => makeEdgeResult({ + file: `src/callers/${index}.ts`, + line: index + 10, + range: { startLine: index + 10, endLine: index + 10 }, + srcSymbol: `caller${index}` + })), + neighbors: [ + { name: 'ParsedToken', file: 'src/auth.ts', range: { startLine: 1, endLine: 4 }, kind: 'interface' } + ], + omitted: 3 + } + }) + ], { maxTokens: 45 }) + + expect(Math.ceil(output.length / 4)).toBeLessThanOrEqual(45) + expect(output).toContain('relations:') + expect(output).toContain('relations_omitted=') + }) }) describe('formatMcpSearchHits structure-preserving output', () => { From 16cc515e076a76238dc43bf406fbe4ff87b2c9c8 Mon Sep 17 00:00:00 2001 From: Rutvik Chandla Date: Sat, 20 Jun 2026 14:33:08 +0530 Subject: [PATCH 13/13] fix(core): use posix paths for module edge resolution --- packages/core/src/repo.ts | 14 +++++++------- packages/eval/losses.json | 1 + 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/packages/core/src/repo.ts b/packages/core/src/repo.ts index c42867e..83a0f56 100644 --- a/packages/core/src/repo.ts +++ b/packages/core/src/repo.ts @@ -2,7 +2,7 @@ import { createHash } from 'node:crypto' import ts from 'typescript' import { existsSync, watch as watchFs, type FSWatcher } from 'node:fs' import { copyFile, mkdir, readFile, rename, rm, writeFile } from 'node:fs/promises' -import { dirname, isAbsolute, join, resolve } from 'node:path' +import { dirname, isAbsolute, join, posix as pathPosix, resolve } from 'node:path' import Database from 'better-sqlite3' import { minimatch } from 'minimatch' @@ -3608,8 +3608,8 @@ function resolveTypeScriptModuleTarget(file: ScannedFile, moduleSpecifier: strin } function buildTypeScriptModuleCandidates(importerFile: string, moduleSpecifier: string): string[] { - const importerDirectory = dirname(importerFile) - const base = normalizeRelativeRepoPath(resolve('/', importerDirectory, moduleSpecifier)) + const importerDirectory = pathPosix.dirname(normalizeRelativeRepoPath(importerFile)) + const base = normalizeRelativeRepoPath(pathPosix.resolve('/', importerDirectory, moduleSpecifier)) return [ base, `${base}.ts`, @@ -3638,17 +3638,17 @@ function resolvePythonModuleTarget(file: ScannedFile, moduleSpecifier: string): } function buildPythonModuleCandidates(importerFile: string, moduleSpecifier: string): string[] { - const importerDirectory = dirname(importerFile) + const importerDirectory = pathPosix.dirname(normalizeRelativeRepoPath(importerFile)) const relativeMatch = moduleSpecifier.match(/^(\.+)(.*)$/) let base: string if (relativeMatch) { const dots = relativeMatch[1]!.length const rest = relativeMatch[2] ?? '' - const parentDirectory = dots > 1 ? resolve('/', importerDirectory, ...Array.from({ length: dots - 1 }, () => '..')) : resolve('/', importerDirectory) - base = normalizeRelativeRepoPath(resolve(parentDirectory, rest.replace(/\./g, '/'))) + const parentDirectory = dots > 1 ? pathPosix.resolve('/', importerDirectory, ...Array.from({ length: dots - 1 }, () => '..')) : pathPosix.resolve('/', importerDirectory) + base = normalizeRelativeRepoPath(pathPosix.resolve(parentDirectory, rest.replace(/\./g, '/'))) } else { - base = normalizeRelativeRepoPath(resolve('/', moduleSpecifier.replace(/\./g, '/'))) + base = normalizeRelativeRepoPath(pathPosix.resolve('/', moduleSpecifier.replace(/\./g, '/'))) } return [base, `${base}.py`, `${base}/__init__.py`] diff --git a/packages/eval/losses.json b/packages/eval/losses.json index 0275982..5bd05ec 100644 --- a/packages/eval/losses.json +++ b/packages/eval/losses.json @@ -48,6 +48,7 @@ "queryId": "escape-regexp-symbol", "queryType": "symbol-def", "axes": [ + "latency.cold", "precision" ] },