Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
227 changes: 222 additions & 5 deletions packages/mcp-server/src/formatters/__tests__/formatters.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -173,14 +173,20 @@ describe('Formatters', () => {
expect(result.content).not.toContain('3.');
});

it('should estimate more tokens than compact', () => {
it('should estimate more tokens than compact when snippets disabled', () => {
// When snippets are disabled, verbose still has more metadata
const compactFormatter = new CompactFormatter();
const verboseFormatter = new VerboseFormatter();
const verboseFormatter = new VerboseFormatter({
includeSnippets: false,
includeImports: false,
});

const compactTokens = compactFormatter.estimateTokens(mockResults[0]);
const verboseTokens = verboseFormatter.estimateTokens(mockResults[0]);
// Use formatResult which includes all the metadata lines
const compactOutput = compactFormatter.formatResult(mockResults[0]);
const verboseOutput = verboseFormatter.formatResult(mockResults[0]);

expect(verboseTokens).toBeGreaterThan(compactTokens);
// Verbose output should be longer (has Location, Signature, Metadata lines)
expect(verboseOutput.length).toBeGreaterThan(compactOutput.length);
});

it('should handle missing metadata gracefully', () => {
Expand Down Expand Up @@ -250,4 +256,215 @@ describe('Formatters', () => {
expect(result.tokens).toBeGreaterThan(0);
});
});

describe('Snippet and Import Formatting', () => {
const resultWithSnippet: SearchResult = {
id: 'src/auth/handler.ts:handleAuth:45',
score: 0.85,
metadata: {
path: 'src/auth/handler.ts',
type: 'function',
language: 'typescript',
name: 'handleAuth',
startLine: 45,
endLine: 67,
exported: true,
snippet:
'export async function handleAuth(req: Request): Promise<Response> {\n const token = extractToken(req);\n return validateToken(token);\n}',
imports: ['./service', '../utils/jwt', 'express'],
},
};

const resultWithManyImports: SearchResult = {
id: 'src/index.ts:main:1',
score: 0.75,
metadata: {
path: 'src/index.ts',
type: 'function',
name: 'main',
startLine: 1,
endLine: 10,
imports: ['a', 'b', 'c', 'd', 'e', 'f', 'g'],
},
};

describe('CompactFormatter with snippets', () => {
it('should not include snippet by default', () => {
const formatter = new CompactFormatter();
const formatted = formatter.formatResult(resultWithSnippet);

expect(formatted).not.toContain('export async function');
expect(formatted).not.toContain('Imports:');
});

it('should include snippet when enabled', () => {
const formatter = new CompactFormatter({ includeSnippets: true });
const formatted = formatter.formatResult(resultWithSnippet);

expect(formatted).toContain('export async function handleAuth');
expect(formatted).toContain('extractToken');
});

it('should include imports when enabled', () => {
const formatter = new CompactFormatter({ includeImports: true });
const formatted = formatter.formatResult(resultWithSnippet);

expect(formatted).toContain('Imports:');
expect(formatted).toContain('./service');
expect(formatted).toContain('express');
});

it('should truncate long import lists', () => {
const formatter = new CompactFormatter({ includeImports: true });
const formatted = formatter.formatResult(resultWithManyImports);

expect(formatted).toContain('Imports:');
expect(formatted).toContain('a, b, c, d, e');
expect(formatted).toContain('...');
expect(formatted).not.toContain('f, g');
});

it('should truncate long snippets', () => {
const longSnippet = Array(20).fill('const x = 1;').join('\n');
const result: SearchResult = {
id: 'test',
score: 0.8,
metadata: {
path: 'test.ts',
type: 'function',
name: 'test',
snippet: longSnippet,
},
};

const formatter = new CompactFormatter({ includeSnippets: true, maxSnippetLines: 5 });
const formatted = formatter.formatResult(result);

expect(formatted).toContain('// ... 15 more lines');
});

it('should increase token estimate with snippets', () => {
const formatterWithout = new CompactFormatter();
const formatterWith = new CompactFormatter({ includeSnippets: true, includeImports: true });

const tokensWithout = formatterWithout.estimateTokens(resultWithSnippet);
const tokensWith = formatterWith.estimateTokens(resultWithSnippet);

expect(tokensWith).toBeGreaterThan(tokensWithout);
});
});

describe('VerboseFormatter with snippets', () => {
it('should include snippet by default', () => {
const formatter = new VerboseFormatter();
const formatted = formatter.formatResult(resultWithSnippet);

expect(formatted).toContain('Code:');
expect(formatted).toContain('export async function handleAuth');
});

it('should include imports by default', () => {
const formatter = new VerboseFormatter();
const formatted = formatter.formatResult(resultWithSnippet);

expect(formatted).toContain('Imports: ./service, ../utils/jwt, express');
});

it('should show location with line range', () => {
const formatter = new VerboseFormatter();
const formatted = formatter.formatResult(resultWithSnippet);

expect(formatted).toContain('Location: src/auth/handler.ts:45-67');
});

it('should not truncate imports in verbose mode', () => {
const formatter = new VerboseFormatter();
const formatted = formatter.formatResult(resultWithManyImports);

expect(formatted).toContain('Imports: a, b, c, d, e, f, g');
expect(formatted).not.toContain('...');
});

it('should respect maxSnippetLines option', () => {
const longSnippet = Array(30).fill('const x = 1;').join('\n');
const result: SearchResult = {
id: 'test',
score: 0.8,
metadata: {
path: 'test.ts',
type: 'function',
name: 'test',
snippet: longSnippet,
},
};

const formatter = new VerboseFormatter({ maxSnippetLines: 10 });
const formatted = formatter.formatResult(result);

expect(formatted).toContain('// ... 20 more lines');
});

it('should be able to disable snippets', () => {
const formatter = new VerboseFormatter({ includeSnippets: false });
const formatted = formatter.formatResult(resultWithSnippet);

expect(formatted).not.toContain('Code:');
expect(formatted).not.toContain('export async function');
});

it('should be able to disable imports', () => {
const formatter = new VerboseFormatter({ includeImports: false });
const formatted = formatter.formatResult(resultWithSnippet);

expect(formatted).not.toContain('Imports:');
});

it('should increase token estimate with snippets', () => {
const formatterWithout = new VerboseFormatter({
includeSnippets: false,
includeImports: false,
});
const formatterWith = new VerboseFormatter();

const tokensWithout = formatterWithout.estimateTokens(resultWithSnippet);
const tokensWith = formatterWith.estimateTokens(resultWithSnippet);

expect(tokensWith).toBeGreaterThan(tokensWithout);
});
});

describe('Empty snippets and imports', () => {
it('should handle missing snippet gracefully', () => {
const formatter = new VerboseFormatter();
const formatted = formatter.formatResult(mockResults[0]);

expect(formatted).not.toContain('Code:');
});

it('should handle missing imports gracefully', () => {
const formatter = new VerboseFormatter();
const formatted = formatter.formatResult(mockResults[0]);

expect(formatted).not.toContain('Imports:');
});

it('should handle empty imports array', () => {
const result: SearchResult = {
id: 'test',
score: 0.8,
metadata: {
path: 'test.ts',
type: 'function',
name: 'test',
imports: [],
},
};

const formatter = new VerboseFormatter();
const formatted = formatter.formatResult(result);

expect(formatted).not.toContain('Imports:');
});
});
});
});
73 changes: 69 additions & 4 deletions packages/mcp-server/src/formatters/compact-formatter.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,14 @@ import type { SearchResult } from '@lytics/dev-agent-core';
import type { 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;

/**
* Compact formatter - optimized for token efficiency
* Returns: path, type, name, score
* Returns: path, type, name, score, optional snippet and imports
*/
export class CompactFormatter implements ResultFormatter {
private options: Required<FormatterOptions>;
Expand All @@ -21,14 +26,45 @@ export class CompactFormatter implements ResultFormatter {
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
maxSnippetLines: options.maxSnippetLines ?? DEFAULT_MAX_SNIPPET_LINES,
tokenBudget: options.tokenBudget ?? 1000,
};
}

formatResult(result: SearchResult): string {
const lines: string[] = [];

// Line 1: Header with score, type, name, path
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));
}

// 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}`);
}
}

return lines.join('\n');
}

private formatHeader(result: SearchResult): string {
const parts: string[] = [];

// Score (2 decimals)
// Score
parts.push(`[${(result.score * 100).toFixed(0)}%]`);

// Type
Expand All @@ -41,7 +77,7 @@ export class CompactFormatter implements ResultFormatter {
parts.push(result.metadata.name);
}

// Path
// Path with line numbers
if (this.options.includePaths && typeof result.metadata.path === 'string') {
const pathPart =
this.options.includeLineNumbers && typeof result.metadata.startLine === 'number'
Expand All @@ -53,6 +89,24 @@ export class CompactFormatter implements ResultFormatter {
return parts.join(' ');
}

private truncateSnippet(snippet: string, maxLines: number): string {
const lines = snippet.split('\n');
if (lines.length <= maxLines) {
return snippet;
}
const truncated = lines.slice(0, maxLines).join('\n');
const remaining = lines.length - maxLines;
return `${truncated}\n// ... ${remaining} more lines`;
}

private indentText(text: string, spaces: number): string {
const indent = ' '.repeat(spaces);
return text
.split('\n')
.map((line) => indent + line)
.join('\n');
}

formatResults(results: SearchResult[]): FormattedResult {
// Handle empty results
if (results.length === 0) {
Expand Down Expand Up @@ -82,6 +136,17 @@ export class CompactFormatter implements ResultFormatter {
}

estimateTokens(result: SearchResult): number {
return estimateTokensForText(this.formatResult(result));
let estimate = estimateTokensForText(this.formatHeader(result));

if (this.options.includeSnippets && typeof result.metadata.snippet === 'string') {
estimate += estimateTokensForText(result.metadata.snippet);
}

if (this.options.includeImports && Array.isArray(result.metadata.imports)) {
// ~3 tokens per import path
estimate += (result.metadata.imports as string[]).length * 3;
}

return estimate;
}
}
15 changes: 15 additions & 0 deletions packages/mcp-server/src/formatters/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,21 @@ export interface FormatterOptions {
*/
includeSignatures?: boolean;

/**
* Include code snippets in output
*/
includeSnippets?: boolean;

/**
* Include import lists in output
*/
includeImports?: boolean;

/**
* Maximum lines to show in snippets (default: 10 compact, 20 verbose)
*/
maxSnippetLines?: number;

/**
* Token budget (soft limit)
*/
Expand Down
Loading