fix: Return whole lines when looking up source code

packages/cli/src/rpc/explain/buildContext.ts

@@ -70,13 +70,13 @@ export default async function buildContext(
 
       codeSnippetLocations.add(event.location);
 
-      const snippets = await lookupSourceCode(result.directory, event.location);
-      if (snippets) {
+      const snippet = await lookupSourceCode(result.directory, event.location);
+      if (snippet) {
         codeSnippets.push({
           directory: result.directory,
           type: ContextV2.ContextItemType.CodeSnippet,
-          location: event.location,
-          content: snippets.join('\n'),
+          location: snippet.location,
+          content: snippet.content,
           score: event.score,
         });
       }

packages/cli/src/rpc/explain/lookupSourceCode.ts

@@ -29,7 +29,7 @@ const FILE_NAMES = new Set<string>();
 export default async function lookupSourceCode(
   directory: string,
   location: string
-): Promise<string[] | undefined> {
+): Promise<{ content: string; location: string } | undefined> {
   const parsedLocation = parseLocation(location);
   if (!parsedLocation) return;
 
@@ -77,9 +77,9 @@ export default async function lookupSourceCode(
   }
 
   const fileContent = await readFile(fileName, 'utf-8');
-  if (!lineNo) return [fileContent];
+  const path = fileName.slice(directory.length + 1);
 
-  if (lineNo <= 0) return [fileContent];
+  if (!lineNo || lineNo <= 0) return { content: fileContent, location: path };
 
   const fileExtension = fileName.slice(fileName.lastIndexOf('.'));
   const language = LANGUAGE_BY_FILE_EXTENSION[fileExtension];
@@ -100,6 +100,30 @@ export default async function lookupSourceCode(
   const matches = chunks.filter(
     (chunk) => chunk.metadata.loc.lines.from <= lineNo && chunk.metadata.loc.lines.to >= lineNo
   );
-  if (verbose()) warn(chalk.gray(`Obtained ${matches.length} source code chunks for ${location}`));
-  return matches.map((match) => match.pageContent);
+
+  // determine the extent of the snippets and return a single snippet that contains all the complete lines
+  if (matches.length === 0) {
+    warn(chalk.gray(`No source code found for ${location}`));
+    return;
+  }
+
+  const extent = {
+    from: Math.min(...matches.map((match) => match.metadata.loc.lines.from)),
+    to: Math.max(...matches.map((match) => match.metadata.loc.lines.to)),
+  };
+
+  if (verbose())
+    warn(
+      chalk.gray(
+        `Found ${matches.length} matches for ${location} (extent: ${extent.from}-${extent.to})`
+      )
+    );
+
+  return {
+    content: fileContent
+      .split('\n')
+      .slice(extent.from - 1, extent.to)
+      .join('\n'),
+    location: `${path}:${extent.from}-${extent.to}`,
+  };
 }

packages/cli/tests/unit/rpc/explain/buildContext.spec.ts

@@ -36,7 +36,9 @@ describe('buildContext', () => {
 
   describe('function call', () => {
     it('is emitted with source code', async () => {
-      jest.mocked(lookupSourceCode).mockResolvedValue(['the code']);
+      jest
+        .mocked(lookupSourceCode)
+        .mockResolvedValue({ content: 'the code', location: 'app/models/user.rb:42-47' });
 
       const context = await buildContext([
         {
@@ -53,7 +55,7 @@ describe('buildContext', () => {
       ).toEqual([
         {
           directory: 'a',
-          location: 'app/models/user.rb',
+          location: 'app/models/user.rb:42-47',
           type: 'code-snippet',
           score: 1,
           content: 'the code',

packages/cli/tests/unit/rpc/explain/lookupSourceCode.spec.ts

@@ -0,0 +1,52 @@
+import path from 'node:path';
+
+import { RecursiveCharacterTextSplitter } from 'langchain/text_splitter';
+import { fs, vol } from 'memfs';
+
+import lookupSourceCode from '../../../../src/rpc/explain/lookupSourceCode';
+
+describe('lookupSourceCode', () => {
+  const directory = '/project';
+  const fileName = 'example.rb';
+  const sourceCodeSnippet = 'module A\n  class B\n    def foo\n      pass\n    end\n  end\nend';
+  const filePath = path.join(directory, fileName);
+
+  it('returns the correct snippet object', async () => {
+    const snippet = await lookupSourceCode(directory, `${fileName}:3`);
+    expect(snippet).toEqual({
+      content: sourceCodeSnippet,
+      location: 'example.rb:1-7',
+    });
+  });
+
+  it('preserves indents even on a chunk boundary', async () => {
+    const smallSplitter = RecursiveCharacterTextSplitter.fromLanguage('ruby', {
+      chunkSize: 10,
+      chunkOverlap: 0,
+    });
+    jest.spyOn(RecursiveCharacterTextSplitter, 'fromLanguage').mockReturnValue(smallSplitter);
+    const snippet = await lookupSourceCode(directory, `${fileName}:4`);
+    expect(snippet).toEqual({
+      content: '      pass',
+      location: 'example.rb:4-4',
+    });
+  });
+
+  it('returns undefined for non-existent line', async () => {
+    const snippet = await lookupSourceCode(directory, `${fileName}:10`);
+    expect(snippet).toBeUndefined();
+  });
+
+  beforeEach(() => {
+    vol.mkdirSync(directory, { recursive: true });
+    vol.writeFileSync(filePath, sourceCodeSnippet);
+  });
+
+  afterEach(() => {
+    vol.reset();
+    jest.restoreAllMocks();
+  });
+});
+
+jest.mock('fs', () => fs);
+jest.mock('fs/promises', () => fs.promises);