diff --git a/packages/mcp-server/src/adapters/built-in/search-adapter.ts b/packages/mcp-server/src/adapters/built-in/search-adapter.ts index 68d2dff..c272b50 100644 --- a/packages/mcp-server/src/adapters/built-in/search-adapter.ts +++ b/packages/mcp-server/src/adapters/built-in/search-adapter.ts @@ -41,8 +41,6 @@ export class SearchAdapter extends ToolAdapter { }; private indexer: RepositoryIndexer; - private compactFormatter: CompactFormatter; - private verboseFormatter: VerboseFormatter; private config: Required; constructor(config: SearchAdapterConfig) { @@ -53,17 +51,6 @@ export class SearchAdapter extends ToolAdapter { defaultFormat: config.defaultFormat ?? 'compact', defaultLimit: config.defaultLimit ?? 10, }; - - // Initialize formatters - this.compactFormatter = new CompactFormatter({ - maxResults: this.config.defaultLimit, - tokenBudget: 1000, - }); - - this.verboseFormatter = new VerboseFormatter({ - maxResults: this.config.defaultLimit, - tokenBudget: 5000, - }); } async initialize(context: AdapterContext): Promise { @@ -107,6 +94,13 @@ export class SearchAdapter extends ToolAdapter { maximum: 1, default: 0, }, + tokenBudget: { + type: 'number', + description: + 'Maximum tokens for results. Uses progressive disclosure to fit within budget (default: 2000 compact, 5000 verbose)', + minimum: 500, + maximum: 10000, + }, }, required: ['query'], }, @@ -119,6 +113,7 @@ export class SearchAdapter extends ToolAdapter { format = this.config.defaultFormat, limit = this.config.defaultLimit, scoreThreshold = 0, + tokenBudget, } = args; // Validate query @@ -165,9 +160,29 @@ export class SearchAdapter extends ToolAdapter { }; } + // Validate tokenBudget if provided + if ( + tokenBudget !== undefined && + (typeof tokenBudget !== 'number' || tokenBudget < 500 || tokenBudget > 10000) + ) { + return { + success: false, + error: { + code: 'INVALID_TOKEN_BUDGET', + message: 'Token budget must be a number between 500 and 10000', + }, + }; + } + try { const startTime = Date.now(); - context.logger.debug('Executing search', { query, format, limit, scoreThreshold }); + context.logger.debug('Executing search', { + query, + format, + limit, + scoreThreshold, + tokenBudget, + }); // Perform search const results = await this.indexer.search(query as string, { @@ -175,8 +190,22 @@ export class SearchAdapter extends ToolAdapter { scoreThreshold: scoreThreshold as number, }); - // Format results - const formatter = format === 'verbose' ? this.verboseFormatter : this.compactFormatter; + // Create formatter with token budget if specified + const formatter = + format === 'verbose' + ? new VerboseFormatter({ + maxResults: limit as number, + tokenBudget: (tokenBudget as number | undefined) ?? 5000, + includeSnippets: true, + includeImports: true, + }) + : new CompactFormatter({ + maxResults: limit as number, + tokenBudget: (tokenBudget as number | undefined) ?? 2000, + includeSnippets: true, + includeImports: true, + }); + const formatted = formatter.formatResults(results); const duration_ms = Date.now() - startTime; diff --git a/packages/mcp-server/src/formatters/__tests__/formatters.test.ts b/packages/mcp-server/src/formatters/__tests__/formatters.test.ts index b674cb7..45e95d8 100644 --- a/packages/mcp-server/src/formatters/__tests__/formatters.test.ts +++ b/packages/mcp-server/src/formatters/__tests__/formatters.test.ts @@ -467,4 +467,183 @@ describe('Formatters', () => { }); }); }); + + describe('Token Budget Management', () => { + // Generate mock results with snippets for testing + const generateMockResults = (count: number): SearchResult[] => { + return Array.from({ length: count }, (_, i) => ({ + id: `test-${i}`, + score: 0.9 - i * 0.05, + metadata: { + path: `src/file${i}.ts`, + type: 'function', + language: 'typescript', + name: `function${i}`, + startLine: i * 10, + endLine: i * 10 + 20, + exported: true, + signature: `function function${i}(): void`, + snippet: `export function function${i}() {\n // Line 1\n // Line 2\n // Line 3\n return ${i};\n}`, + imports: ['./utils', '../lib'], + }, + })); + }; + + describe('CompactFormatter budget', () => { + it('should respect token budget', () => { + const formatter = new CompactFormatter({ + tokenBudget: 100, + maxResults: 10, + includeSnippets: true, + includeImports: true, + progressiveDisclosure: false, // Force all results to use full detail + }); + + const results = generateMockResults(10); + const output = formatter.formatResults(results); + + // Should be within budget (with some tolerance for truncation notice) + expect(output.tokens).toBeLessThanOrEqual(150); + // Should have truncation notice since budget is small + expect(output.content).toContain('more results'); + }); + + it('should use progressive disclosure', () => { + const formatter = new CompactFormatter({ + tokenBudget: 1000, + maxResults: 10, + includeSnippets: true, + includeImports: true, + fullDetailCount: 2, + signatureDetailCount: 2, + }); + + const results = generateMockResults(10); + const output = formatter.formatResults(results); + + // First results should have snippets + expect(output.content).toContain('export function function0'); + // Later results should be minimal (just name + path) + // Check that result 5+ doesn't have its snippet + const lines = output.content.split('\n'); + const result5Line = lines.find((l) => l.includes('function5')); + expect(result5Line).toBeDefined(); + }); + + it('should always include at least first result', () => { + const formatter = new CompactFormatter({ + tokenBudget: 10, // Very small budget + maxResults: 10, + includeSnippets: true, + }); + + const results = generateMockResults(5); + const output = formatter.formatResults(results); + + // Should have at least one result + expect(output.content).toContain('function0'); + }); + + it('should disable progressive disclosure when option is false', () => { + const formatter = new CompactFormatter({ + tokenBudget: 500, + maxResults: 5, + includeSnippets: true, + progressiveDisclosure: false, + }); + + const results = generateMockResults(5); + const output = formatter.formatResults(results); + + // All results should have full detail (until budget runs out) + // Check that early results have snippets + expect(output.content).toContain('export function function0'); + }); + }); + + describe('VerboseFormatter budget', () => { + it('should respect token budget', () => { + const formatter = new VerboseFormatter({ + tokenBudget: 500, + maxResults: 10, + includeSnippets: true, + includeImports: true, + }); + + const results = generateMockResults(10); + const output = formatter.formatResults(results); + + // Should be within budget (with some tolerance) + expect(output.tokens).toBeLessThanOrEqual(600); + }); + + it('should use progressive disclosure', () => { + const formatter = new VerboseFormatter({ + tokenBudget: 2000, + maxResults: 10, + includeSnippets: true, + includeImports: true, + fullDetailCount: 2, + signatureDetailCount: 3, + }); + + const results = generateMockResults(10); + const output = formatter.formatResults(results); + + // First results should have Code: section + expect(output.content).toContain('Code:'); + expect(output.content).toContain('export function function0'); + }); + + it('should show truncation notice when budget exceeded', () => { + const formatter = new VerboseFormatter({ + tokenBudget: 100, // Very small budget + maxResults: 10, + includeSnippets: true, + progressiveDisclosure: false, // Force all results to use full detail + }); + + const results = generateMockResults(10); + const output = formatter.formatResults(results); + + expect(output.content).toContain('more results (token budget reached)'); + }); + }); + + describe('Detail levels', () => { + it('should format with full detail', () => { + const formatter = new CompactFormatter({ includeSnippets: true, includeImports: true }); + const result = generateMockResults(1)[0]; + + const output = formatter.formatResultWithDetail(result, 'full'); + + expect(output).toContain('export function'); + expect(output).toContain('Imports:'); + }); + + it('should format with signature detail', () => { + const formatter = new CompactFormatter({ includeSnippets: true, includeImports: true }); + const result = generateMockResults(1)[0]; + + const output = formatter.formatResultWithDetail(result, 'signature'); + + expect(output).toContain('function function0(): void'); + expect(output).not.toContain('export function'); + expect(output).not.toContain('Imports:'); + }); + + it('should format with minimal detail', () => { + const formatter = new CompactFormatter({ includeSnippets: true, includeImports: true }); + const result = generateMockResults(1)[0]; + + const output = formatter.formatResultWithDetail(result, 'minimal'); + + expect(output).toContain('function0'); + expect(output).toContain('src/file0.ts'); + expect(output).not.toContain('export function'); + expect(output).not.toContain('Imports:'); + expect(output).not.toContain('function function0(): void'); + }); + }); + }); }); diff --git a/packages/mcp-server/src/formatters/compact-formatter.ts b/packages/mcp-server/src/formatters/compact-formatter.ts index d00ee41..ccd13c9 100644 --- a/packages/mcp-server/src/formatters/compact-formatter.ts +++ b/packages/mcp-server/src/formatters/compact-formatter.ts @@ -1,20 +1,26 @@ /** * Compact Formatter - * Token-efficient formatter that returns summaries only + * Token-efficient formatter with progressive disclosure support */ import type { SearchResult } from '@lytics/dev-agent-core'; -import type { FormattedResult, FormatterOptions, ResultFormatter } from './types'; +import type { DetailLevel, FormattedResult, FormatterOptions, ResultFormatter } from './types'; import { estimateTokensForText } from './utils'; /** Default max snippet lines for compact mode */ const DEFAULT_MAX_SNIPPET_LINES = 10; /** Max imports to show before truncating */ const MAX_IMPORTS_DISPLAY = 5; +/** Default token budget */ +const DEFAULT_TOKEN_BUDGET = 2000; +/** Default number of results with full detail */ +const DEFAULT_FULL_DETAIL_COUNT = 3; +/** Default number of results with signature detail */ +const DEFAULT_SIGNATURE_DETAIL_COUNT = 4; /** * Compact formatter - optimized for token efficiency - * Returns: path, type, name, score, optional snippet and imports + * Supports progressive disclosure to fit within token budgets */ export class CompactFormatter implements ResultFormatter { private options: Required; @@ -25,38 +31,55 @@ export class CompactFormatter implements ResultFormatter { includePaths: options.includePaths ?? true, includeLineNumbers: options.includeLineNumbers ?? true, includeTypes: options.includeTypes ?? true, - includeSignatures: options.includeSignatures ?? false, // Compact mode excludes signatures - includeSnippets: options.includeSnippets ?? false, // Off by default for compact - includeImports: options.includeImports ?? false, // Off by default for compact + includeSignatures: options.includeSignatures ?? false, + includeSnippets: options.includeSnippets ?? false, + includeImports: options.includeImports ?? false, maxSnippetLines: options.maxSnippetLines ?? DEFAULT_MAX_SNIPPET_LINES, - tokenBudget: options.tokenBudget ?? 1000, + tokenBudget: options.tokenBudget ?? DEFAULT_TOKEN_BUDGET, + progressiveDisclosure: options.progressiveDisclosure ?? true, + fullDetailCount: options.fullDetailCount ?? DEFAULT_FULL_DETAIL_COUNT, + signatureDetailCount: options.signatureDetailCount ?? DEFAULT_SIGNATURE_DETAIL_COUNT, }; } formatResult(result: SearchResult): string { + return this.formatResultWithDetail(result, 'full'); + } + + /** + * Format a result with a specific detail level + */ + formatResultWithDetail(result: SearchResult, level: DetailLevel): string { const lines: string[] = []; - // Line 1: Header with score, type, name, path + // Always include header lines.push(this.formatHeader(result)); - // Code snippet (if enabled) - if (this.options.includeSnippets && typeof result.metadata.snippet === 'string') { - const truncatedSnippet = this.truncateSnippet( - result.metadata.snippet, - this.options.maxSnippetLines - ); - lines.push(this.indentText(truncatedSnippet, 3)); - } + if (level === 'full') { + // Full detail: snippet + imports + if (this.options.includeSnippets && typeof result.metadata.snippet === 'string') { + const truncatedSnippet = this.truncateSnippet( + result.metadata.snippet, + this.options.maxSnippetLines + ); + lines.push(this.indentText(truncatedSnippet, 3)); + } - // Imports (if enabled) - if (this.options.includeImports && Array.isArray(result.metadata.imports)) { - const imports = result.metadata.imports as string[]; - if (imports.length > 0) { - const displayImports = imports.slice(0, MAX_IMPORTS_DISPLAY); - const suffix = imports.length > MAX_IMPORTS_DISPLAY ? ' ...' : ''; - lines.push(` Imports: ${displayImports.join(', ')}${suffix}`); + if (this.options.includeImports && Array.isArray(result.metadata.imports)) { + const imports = result.metadata.imports as string[]; + if (imports.length > 0) { + const displayImports = imports.slice(0, MAX_IMPORTS_DISPLAY); + const suffix = imports.length > MAX_IMPORTS_DISPLAY ? ' ...' : ''; + lines.push(` Imports: ${displayImports.join(', ')}${suffix}`); + } + } + } else if (level === 'signature') { + // Signature detail: just the signature + if (typeof result.metadata.signature === 'string') { + lines.push(` ${result.metadata.signature}`); } } + // 'minimal' level = header only return lines.join('\n'); } @@ -64,20 +87,16 @@ export class CompactFormatter implements ResultFormatter { private formatHeader(result: SearchResult): string { const parts: string[] = []; - // Score parts.push(`[${(result.score * 100).toFixed(0)}%]`); - // Type if (this.options.includeTypes && typeof result.metadata.type === 'string') { parts.push(`${result.metadata.type}:`); } - // Name if (typeof result.metadata.name === 'string') { parts.push(result.metadata.name); } - // Path with line numbers if (this.options.includePaths && typeof result.metadata.path === 'string') { const pathPart = this.options.includeLineNumbers && typeof result.metadata.startLine === 'number' @@ -107,8 +126,31 @@ export class CompactFormatter implements ResultFormatter { .join('\n'); } + /** + * Determine detail level based on result position and remaining budget + */ + private getDetailLevel(index: number, remainingBudget: number): DetailLevel { + if (!this.options.progressiveDisclosure) { + return 'full'; + } + + // Top N get full detail if budget allows + if (index < this.options.fullDetailCount && remainingBudget > 300) { + return 'full'; + } + + // Next M get signatures if budget allows + if ( + index < this.options.fullDetailCount + this.options.signatureDetailCount && + remainingBudget > 100 + ) { + return 'signature'; + } + + return 'minimal'; + } + formatResults(results: SearchResult[]): FormattedResult { - // Handle empty results if (results.length === 0) { const content = 'No results found'; return { @@ -117,21 +159,43 @@ export class CompactFormatter implements ResultFormatter { }; } - // Respect max results const limitedResults = results.slice(0, this.options.maxResults); + const budget = this.options.tokenBudget; + let usedTokens = 0; + const formatted: string[] = []; + let truncatedCount = 0; + + for (let i = 0; i < limitedResults.length; i++) { + const result = limitedResults[i]; + const remainingBudget = budget - usedTokens; + + // Determine detail level + const detailLevel = this.getDetailLevel(i, remainingBudget); + const formattedResult = `${i + 1}. ${this.formatResultWithDetail(result, detailLevel)}`; + const tokens = estimateTokensForText(formattedResult); + + // Check if we have budget (always include at least first result) + if (usedTokens + tokens > budget && i > 0) { + truncatedCount = limitedResults.length - i; + break; + } + + formatted.push(formattedResult); + usedTokens += tokens; + } - // Format each result - const formatted = limitedResults.map((result, index) => { - return `${index + 1}. ${this.formatResult(result)}`; - }); + // Add truncation notice if needed + if (truncatedCount > 0) { + const notice = `\n... ${truncatedCount} more results (token budget reached)`; + formatted.push(notice); + usedTokens += estimateTokensForText(notice); + } - // Calculate total tokens (content only, no footer) const content = formatted.join('\n'); - const tokens = estimateTokensForText(content); return { content, - tokens, + tokens: usedTokens, }; } @@ -143,7 +207,6 @@ export class CompactFormatter implements ResultFormatter { } if (this.options.includeImports && Array.isArray(result.metadata.imports)) { - // ~3 tokens per import path estimate += (result.metadata.imports as string[]).length * 3; } diff --git a/packages/mcp-server/src/formatters/types.ts b/packages/mcp-server/src/formatters/types.ts index 58aed77..a6a07ee 100644 --- a/packages/mcp-server/src/formatters/types.ts +++ b/packages/mcp-server/src/formatters/types.ts @@ -38,6 +38,14 @@ export interface ResultFormatter { estimateTokens(result: SearchResult): number; } +/** + * Detail level for progressive disclosure + * - full: snippet + imports + signature + * - signature: signature only (no snippet) + * - minimal: name + path only + */ +export type DetailLevel = 'full' | 'signature' | 'minimal'; + /** * Formatter options */ @@ -83,7 +91,24 @@ export interface FormatterOptions { maxSnippetLines?: number; /** - * Token budget (soft limit) + * Token budget (soft limit) - enables progressive disclosure when set */ tokenBudget?: number; + + /** + * Enable progressive disclosure based on token budget + * When enabled, top results get full detail, lower results get less + * (default: true when tokenBudget is set) + */ + progressiveDisclosure?: boolean; + + /** + * Number of top results to show with full detail (default: 3) + */ + fullDetailCount?: number; + + /** + * Number of results after fullDetailCount to show with signatures (default: 4) + */ + signatureDetailCount?: number; } diff --git a/packages/mcp-server/src/formatters/verbose-formatter.ts b/packages/mcp-server/src/formatters/verbose-formatter.ts index 977aab4..ec8b384 100644 --- a/packages/mcp-server/src/formatters/verbose-formatter.ts +++ b/packages/mcp-server/src/formatters/verbose-formatter.ts @@ -1,18 +1,24 @@ /** * Verbose Formatter - * Full-detail formatter that includes signatures, snippets, and metadata + * Full-detail formatter with progressive disclosure support */ import type { SearchResult } from '@lytics/dev-agent-core'; -import type { FormattedResult, FormatterOptions, ResultFormatter } from './types'; +import type { DetailLevel, FormattedResult, FormatterOptions, ResultFormatter } from './types'; import { estimateTokensForText } from './utils'; /** Default max snippet lines for verbose mode */ const DEFAULT_MAX_SNIPPET_LINES = 20; +/** Default token budget for verbose */ +const DEFAULT_TOKEN_BUDGET = 5000; +/** Default number of results with full detail */ +const DEFAULT_FULL_DETAIL_COUNT = 3; +/** Default number of results with signature detail */ +const DEFAULT_SIGNATURE_DETAIL_COUNT = 4; /** * Verbose formatter - includes all available information - * Returns: path, type, name, signature, imports, snippet, metadata, score + * Supports progressive disclosure to fit within token budgets */ export class VerboseFormatter implements ResultFormatter { private options: Required; @@ -23,54 +29,74 @@ export class VerboseFormatter implements ResultFormatter { includePaths: options.includePaths ?? true, includeLineNumbers: options.includeLineNumbers ?? true, includeTypes: options.includeTypes ?? true, - includeSignatures: options.includeSignatures ?? true, // Verbose mode includes signatures - includeSnippets: options.includeSnippets ?? true, // On by default for verbose - includeImports: options.includeImports ?? true, // On by default for verbose + includeSignatures: options.includeSignatures ?? true, + includeSnippets: options.includeSnippets ?? true, + includeImports: options.includeImports ?? true, maxSnippetLines: options.maxSnippetLines ?? DEFAULT_MAX_SNIPPET_LINES, - tokenBudget: options.tokenBudget ?? 5000, + tokenBudget: options.tokenBudget ?? DEFAULT_TOKEN_BUDGET, + progressiveDisclosure: options.progressiveDisclosure ?? true, + fullDetailCount: options.fullDetailCount ?? DEFAULT_FULL_DETAIL_COUNT, + signatureDetailCount: options.signatureDetailCount ?? DEFAULT_SIGNATURE_DETAIL_COUNT, }; } formatResult(result: SearchResult): string { + return this.formatResultWithDetail(result, 'full'); + } + + /** + * Format a result with a specific detail level + */ + formatResultWithDetail(result: SearchResult, level: DetailLevel): string { const lines: string[] = []; - // Header: score + type + name + // Always include header lines.push(this.formatHeader(result)); - // Path with line range + // Path with line range (always include in verbose) if (this.options.includePaths && typeof result.metadata.path === 'string') { const location = this.formatLocation(result); lines.push(` Location: ${location}`); } - // Signature (if available and enabled) - if (this.options.includeSignatures && typeof result.metadata.signature === 'string') { - lines.push(` Signature: ${result.metadata.signature}`); - } + if (level === 'full') { + // Full detail: signature + imports + metadata + snippet + if (this.options.includeSignatures && typeof result.metadata.signature === 'string') { + lines.push(` Signature: ${result.metadata.signature}`); + } - // Imports (if enabled) - if (this.options.includeImports && Array.isArray(result.metadata.imports)) { - const imports = result.metadata.imports as string[]; - if (imports.length > 0) { - lines.push(` Imports: ${imports.join(', ')}`); + if (this.options.includeImports && Array.isArray(result.metadata.imports)) { + const imports = result.metadata.imports as string[]; + if (imports.length > 0) { + lines.push(` Imports: ${imports.join(', ')}`); + } } - } - // Additional metadata - const metadata = this.formatMetadata(result); - if (metadata) { - lines.push(` Metadata: ${metadata}`); - } + const metadata = this.formatMetadata(result); + if (metadata) { + lines.push(` Metadata: ${metadata}`); + } - // Code snippet (if enabled) - if (this.options.includeSnippets && typeof result.metadata.snippet === 'string') { - lines.push(' Code:'); - const truncatedSnippet = this.truncateSnippet( - result.metadata.snippet, - this.options.maxSnippetLines - ); - lines.push(this.indentText(truncatedSnippet, 4)); + if (this.options.includeSnippets && typeof result.metadata.snippet === 'string') { + lines.push(' Code:'); + const truncatedSnippet = this.truncateSnippet( + result.metadata.snippet, + this.options.maxSnippetLines + ); + lines.push(this.indentText(truncatedSnippet, 4)); + } + } else if (level === 'signature') { + // Signature detail: signature + metadata (no snippet) + if (typeof result.metadata.signature === 'string') { + lines.push(` Signature: ${result.metadata.signature}`); + } + + const metadata = this.formatMetadata(result); + if (metadata) { + lines.push(` Metadata: ${metadata}`); + } } + // 'minimal' level = header + location only return lines.join('\n'); } @@ -145,8 +171,31 @@ export class VerboseFormatter implements ResultFormatter { .join('\n'); } + /** + * Determine detail level based on result position and remaining budget + */ + private getDetailLevel(index: number, remainingBudget: number): DetailLevel { + if (!this.options.progressiveDisclosure) { + return 'full'; + } + + // Top N get full detail if budget allows + if (index < this.options.fullDetailCount && remainingBudget > 500) { + return 'full'; + } + + // Next M get signatures if budget allows + if ( + index < this.options.fullDetailCount + this.options.signatureDetailCount && + remainingBudget > 150 + ) { + return 'signature'; + } + + return 'minimal'; + } + formatResults(results: SearchResult[]): FormattedResult { - // Handle empty results if (results.length === 0) { const content = 'No results found'; return { @@ -155,21 +204,43 @@ export class VerboseFormatter implements ResultFormatter { }; } - // Respect max results const limitedResults = results.slice(0, this.options.maxResults); + const budget = this.options.tokenBudget; + let usedTokens = 0; + const formatted: string[] = []; + let truncatedCount = 0; + + for (let i = 0; i < limitedResults.length; i++) { + const result = limitedResults[i]; + const remainingBudget = budget - usedTokens; + + // Determine detail level + const detailLevel = this.getDetailLevel(i, remainingBudget); + const formattedResult = `${i + 1}. ${this.formatResultWithDetail(result, detailLevel)}`; + const tokens = estimateTokensForText(formattedResult); + + // Check if we have budget (always include at least first result) + if (usedTokens + tokens > budget && i > 0) { + truncatedCount = limitedResults.length - i; + break; + } - // Format each result with separator - const formatted = limitedResults.map((result, index) => { - return `${index + 1}. ${this.formatResult(result)}`; - }); + formatted.push(formattedResult); + usedTokens += tokens; + } + + // Add truncation notice if needed + if (truncatedCount > 0) { + const notice = `\n... ${truncatedCount} more results (token budget reached)`; + formatted.push(notice); + usedTokens += estimateTokensForText(notice); + } - // Calculate total tokens (content only, no footer) const content = formatted.join('\n\n'); // Double newline for separation - const tokens = estimateTokensForText(content); return { content, - tokens, + tokens: usedTokens, }; } @@ -181,7 +252,6 @@ export class VerboseFormatter implements ResultFormatter { } if (this.options.includeImports && Array.isArray(result.metadata.imports)) { - // ~3 tokens per import path estimate += (result.metadata.imports as string[]).length * 3; }