feat: improve Search tool truncation logic for long lines

extensions/cli/src/tools/searchCode.test.ts

@@ -1,54 +1,68 @@
-// Since we want to test just the interface and not the internals,
-// let's create a simplified version of the run function to test the truncation logic
-describe("searchCodeTool", () => {
-  // We'll test the functionality without mocking the dependencies
-  // Instead, we'll focus on the core truncation logic by directly checking
-  // if truncation happens with large outputs
-
-  it("should include truncation message when output exceeds limit", () => {
+import { smartTruncate, formatTruncationMessage } from './truncation.js';
+
+describe("searchCodeTool truncation logic", () => {
+  it("should include truncation message when output exceeds line limit", () => {
     // Create a large sample output (more than 100 lines)
     const largeOutput = Array.from(
       { length: 150 },
       (_, i) => `file${i}.ts:${i}:const foo = 'bar';`,
     ).join("\n");
 
-    // Check if the truncation logic is applied by the function
-    const truncatedOutput = `Search results for pattern "foo":\n\n${largeOutput.split("\n").slice(0, 100).join("\n")}\n\n[Results truncated: showing 100 of 150 matches]`;
-
-    // Verify the truncation message is included and only 100 lines are shown
-    const lines = truncatedOutput.split("\n");
-    const nonEmptyLines = lines.filter((line) => line.trim() !== "");
-
-    // Count the content lines (excluding header and truncation message)
-    const contentLines = nonEmptyLines.slice(1, -1);
+    // Apply the new smart truncation logic
+    const truncationResult = smartTruncate(largeOutput, { maxLines: 100 });
+    const truncationMessage = formatTruncationMessage(truncationResult);
 
-    // Check that we have exactly 100 content lines
-    expect(contentLines.length).toBe(100);
-
-    // Check that the truncation message is present
-    expect(truncatedOutput).toContain(
-      "[Results truncated: showing 100 of 150 matches]",
-    );
+    expect(truncationResult.truncated).toBe(true);
+    expect(truncationResult.truncatedLineCount).toBe(100);
+    expect(truncationResult.originalLineCount).toBe(150);
+    expect(truncationMessage).toContain("showing 100 of 150 matches");
   });
 
-  it("should not include truncation message when output is within limit", () => {
+  it("should not include truncation message when output is within limits", () => {
     // Create a sample output (less than 100 lines)
     const smallOutput = Array.from(
       { length: 50 },
       (_, i) => `file${i}.ts:${i}:const foo = 'bar';`,
     ).join("\n");
 
-    // Format the output as the function would
-    const output = `Search results for pattern "foo":\n\n${smallOutput}`;
+    // Apply the new smart truncation logic
+    const truncationResult = smartTruncate(smallOutput, { maxLines: 100 });
+    const truncationMessage = formatTruncationMessage(truncationResult);
+
+    expect(truncationResult.truncated).toBe(false);
+    expect(truncationMessage).toBe("");
+  });
+
+  it("should handle very long lines like base64 content", () => {
+    // Simulate a search result with base64 content
+    const base64Line = `data.js:1:const image = "data:image/png;base64,${'iVBORw0KGgoAAAANSUhEUgAA'.repeat(100)}";`;
+    const normalLines = Array.from({ length: 5 }, (_, i) => `file${i}.js:${i}:normal line`);
+    const content = [base64Line, ...normalLines].join('\n');
+
+    const truncationResult = smartTruncate(content, { 
+      maxLines: 100, 
+      maxLineLength: 200 
+    });
 
-    // Verify no truncation message is included
-    expect(output).not.toContain("[Results truncated:");
+    expect(truncationResult.truncated).toBe(true);
+    expect(truncationResult.content).toContain('... [line truncated]');
+    expect(truncationResult.content.split('\n')[0].length).toBeLessThan(base64Line.length);
+  });
+
+  it("should handle content that exceeds character limit", () => {
+    // Create content that exceeds character limit but not line limit
+    const longLines = Array.from({ length: 10 }, (_, i) => 
+      `file${i}.js:${i}:${'x'.repeat(600)} // long line content`
+    );
+    const content = longLines.join('\n');
 
-    // Count the lines
-    const lines = output.split("\n");
-    const nonEmptyLines = lines.filter((line) => line.trim() !== "");
+    const truncationResult = smartTruncate(content, { 
+      maxLines: 100, 
+      maxChars: 3000 
+    });
 
-    // Check that we have the expected number of lines (header + 50 content lines)
-    expect(nonEmptyLines.length).toBe(51);
+    expect(truncationResult.truncated).toBe(true);
+    expect(truncationResult.truncatedCharCount).toBeLessThanOrEqual(3000);
+    expect(truncationResult.truncatedLineCount).toBeLessThan(10);
   });
 });

extensions/cli/src/tools/searchCode.ts

@@ -3,11 +3,16 @@
 import * as util from "util";
 
 import { Tool } from "./types.js";
+import { smartTruncate, formatTruncationMessage } from "./truncation.js";
 
 const execPromise = util.promisify(child_process.exec);
 
-// Default maximum number of results to display
-const DEFAULT_MAX_RESULTS = 100;
+// Default truncation limits
+const DEFAULT_TRUNCATION_OPTIONS = {
+  maxLines: 100,
+  maxChars: 50000, // 50KB
+  maxLineLength: 1000, // Handle very long lines like base64
+};
 
 export const searchCodeTool: Tool = {
   name: "Search",
@@ -78,19 +83,14 @@
           }.`;
         }
 
-        // Split the results into lines and limit the number of results
-        const lines = stdout.split("\n");
-        const truncated = lines.length > DEFAULT_MAX_RESULTS;
-        const limitedLines = lines.slice(0, DEFAULT_MAX_RESULTS);
-        const resultText = limitedLines.join("\n");
-
-        const truncationMessage = truncated
-          ? `\n\n[Results truncated: showing ${DEFAULT_MAX_RESULTS} of ${lines.length} matches]`
-          : "";
+        // Apply smart truncation to handle both line count and character limits
+        const truncationResult = smartTruncate(stdout.trim(), DEFAULT_TRUNCATION_OPTIONS);
+        const truncationMessage = formatTruncationMessage(truncationResult);
+        const truncationSuffix = truncationMessage ? `\n\n${truncationMessage}` : "";
 
         return `Search results for pattern "${args.pattern}"${
           args.file_pattern ? ` in files matching "${args.file_pattern}"` : ""
-        }:\n\n${resultText}${truncationMessage}`;
+        }:\n\n${truncationResult.content}${truncationSuffix}`;
       } catch (error: any) {
         if (error.code === 1) {
           return `No matches found for pattern "${args.pattern}"${

extensions/cli/src/tools/truncation.test.ts

@@ -0,0 +1,196 @@
+import { smartTruncate, formatTruncationMessage } from './truncation.js';
+
+describe('smartTruncate', () => {
+  describe('line-based truncation', () => {
+    it('should not truncate when content is within line limit', () => {
+      const content = Array.from({ length: 50 }, (_, i) => `line ${i}`).join('\n');
+      const result = smartTruncate(content, { maxLines: 100 });
+
+      expect(result.truncated).toBe(false);
+      expect(result.content).toBe(content);
+      expect(result.originalLineCount).toBe(50);
+      expect(result.truncatedLineCount).toBe(50);
+    });
+
+    it('should truncate when content exceeds line limit', () => {
+      const lines = Array.from({ length: 150 }, (_, i) => `line ${i}`);
+      const content = lines.join('\n');
+      const result = smartTruncate(content, { maxLines: 100 });
+
+      expect(result.truncated).toBe(true);
+      expect(result.originalLineCount).toBe(150);
+      expect(result.truncatedLineCount).toBe(100);
+
+      const expectedContent = lines.slice(0, 100).join('\n');
+      expect(result.content).toBe(expectedContent);
+    });
+  });
+
+  describe('character-based truncation', () => {
+    it('should not truncate when content is within character limit', () => {
+      const content = 'short content';
+      const result = smartTruncate(content, { maxChars: 1000 });
+
+      expect(result.truncated).toBe(false);
+      expect(result.content).toBe(content);
+      expect(result.originalCharCount).toBe(13);
+      expect(result.truncatedCharCount).toBe(13);
+    });
+
+    it('should truncate when content exceeds character limit', () => {
+      // Create content that exceeds character limit but not line limit
+      const longLine = 'x'.repeat(600); // 600 chars
+      const lines = Array.from({ length: 10 }, () => longLine); // 10 lines, ~6000 chars total
+      const content = lines.join('\n');
+
+      const result = smartTruncate(content, { maxChars: 3000, maxLines: 100 });
+
+      expect(result.truncated).toBe(true);
+      expect(result.originalCharCount).toBeGreaterThan(3000);
+      expect(result.truncatedCharCount).toBeLessThanOrEqual(3000);
+    });
+  });
+
+  describe('long line truncation', () => {
+    it('should truncate individual lines that are too long', () => {
+      const veryLongLine = 'x'.repeat(2000); // Simulates base64 or other long content
+      const content = `normal line\n${veryLongLine}\nanother normal line`;
+
+      const result = smartTruncate(content, { maxLineLength: 1000 });
+
+      expect(result.truncated).toBe(true);
+      expect(result.content).toContain('... [line truncated]');
+      expect(result.content.split('\n')[1].length).toBeLessThan(veryLongLine.length);
+    });
+
+    it('should handle base64-like content', () => {
+      // Simulate a file with base64 content
+      const base64Content = 'data:image/png;base64,' + 'iVBORw0KGgoAAAANSUhEUgAA'.repeat(100);
+      const content = `file.js:1:const image = "${base64Content}";`;
+
+      const result = smartTruncate(content, { maxLineLength: 200 });
+
+      expect(result.truncated).toBe(true);
+      expect(result.content).toContain('... [line truncated]');
+      expect(result.content.length).toBeLessThan(content.length);
+    });
+  });
+
+  describe('combined limits', () => {
+    it('should respect whichever limit is hit first', () => {
+      // Create content that hits character limit before line limit
+      const longLines = Array.from({ length: 50 }, (_, i) => 'x'.repeat(200) + ` line ${i}`);
+      const content = longLines.join('\n');
+
+      const result = smartTruncate(content, { 
+        maxLines: 100, 
+        maxChars: 5000,
+        maxLineLength: 1000
+      });
+
+      expect(result.truncated).toBe(true);
+      expect(result.truncatedLineCount).toBeLessThan(50); // Hit char limit first
+      expect(result.truncatedCharCount).toBeLessThanOrEqual(5000);
+    });
+  });
+
+  describe('edge cases', () => {
+    it('should handle empty content', () => {
+      const result = smartTruncate('');
+
+      expect(result.truncated).toBe(false);
+      expect(result.content).toBe('');
+      expect(result.originalLineCount).toBe(1); // split('\n') on empty string gives ['']
+      expect(result.truncatedLineCount).toBe(1);
+    });
+
+    it('should handle single very long line', () => {
+      const veryLongContent = 'x'.repeat(100000);
+      const result = smartTruncate(veryLongContent, { maxChars: 1000 });
+
+      expect(result.truncated).toBe(true);
+      expect(result.truncatedCharCount).toBeLessThanOrEqual(1000);
+    });
+
+    it('should handle content with only newlines', () => {
+      const content = '\n'.repeat(200);
+      const result = smartTruncate(content, { maxLines: 100 });
+
+      expect(result.truncated).toBe(true);
+      expect(result.truncatedLineCount).toBe(100);
+    });
+  });
+});
+
+describe('formatTruncationMessage', () => {
+  it('should return empty string when not truncated', () => {
+    const result = {
+      content: 'test',
+      truncated: false,
+      originalLineCount: 5,
+      truncatedLineCount: 5,
+      originalCharCount: 100,
+      truncatedCharCount: 100,
+    };
+
+    expect(formatTruncationMessage(result)).toBe('');
+  });
+
+  it('should format line truncation message', () => {
+    const result = {
+      content: 'test',
+      truncated: true,
+      originalLineCount: 150,
+      truncatedLineCount: 100,
+      originalCharCount: 1000,
+      truncatedCharCount: 800,
+    };
+
+    const message = formatTruncationMessage(result);
+    expect(message).toContain('showing 100 of 150 matches');
+  });
+
+  it('should format character truncation message', () => {
+    const result = {
+      content: 'test',
+      truncated: true,
+      originalLineCount: 50,
+      truncatedLineCount: 50,
+      originalCharCount: 50000,
+      truncatedCharCount: 25000,
+    };
+
+    const message = formatTruncationMessage(result);
+    expect(message).toContain('25KB of 50KB total');
+  });
+
+  it('should format combined truncation message', () => {
+    const result = {
+      content: 'test',
+      truncated: true,
+      originalLineCount: 150,
+      truncatedLineCount: 100,
+      originalCharCount: 50000,
+      truncatedCharCount: 25000,
+    };
+
+    const message = formatTruncationMessage(result);
+    expect(message).toContain('showing 100 of 150 matches');
+    expect(message).toContain('25KB of 50KB total');
+  });
+
+  it('should not show KB message when difference is negligible', () => {
+    const result = {
+      content: 'test',
+      truncated: true,
+      originalLineCount: 150,
+      truncatedLineCount: 100,
+      originalCharCount: 1200,
+      truncatedCharCount: 1100,
+    };
+
+    const message = formatTruncationMessage(result);
+    expect(message).toContain('showing 100 of 150 matches');
+    expect(message).not.toContain('KB');
+  });
+});