From a167f28eadd5b447906c2e9ae82efd80a8309e59 Mon Sep 17 00:00:00 2001
From: Shardul Natu <43422294+kiranani@users.noreply.github.com>
Date: Sat, 30 Aug 2025 13:56:10 -0700
Subject: [PATCH] fix(diffstats): Fix diff stats to correctly capture the edits
 (#7446)

Co-authored-by: Shnatu <snatu@google.com>
Co-authored-by: Gaurav <39389231+gsquared94@users.noreply.github.com>
---
 packages/core/src/tools/edit.test.ts | 24 +++++++++++++++++++-----
 packages/core/src/tools/edit.ts      | 25 +++++++++++--------------
 2 files changed, 30 insertions(+), 19 deletions(-)

diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts
index bf003dc210..a25f1e8419 100644
--- a/packages/core/src/tools/edit.test.ts
+++ b/packages/core/src/tools/edit.test.ts
@@ -507,7 +507,7 @@ describe('EditTool', () => {
     });
 
     it('should successfully replace multiple occurrences when expected_replacements specified', async () => {
-      fs.writeFileSync(filePath, 'old text old text old text', 'utf8');
+      fs.writeFileSync(filePath, 'old text\nold text\nold text', 'utf8');
       const params: EditToolParams = {
         file_path: filePath,
         old_string: 'old',
@@ -520,12 +520,19 @@ describe('EditTool', () => {
 
       expect(result.llmContent).toMatch(/Successfully modified file/);
       expect(fs.readFileSync(filePath, 'utf8')).toBe(
-        'new text new text new text',
+        'new text\nnew text\nnew text',
       );
       const display = result.returnDisplay as FileDiff;
-      expect(display.fileDiff).toMatch(/old text old text old text/);
-      expect(display.fileDiff).toMatch(/new text new text new text/);
+
+      expect(display.fileDiff).toMatch(/-old text\n-old text\n-old text/);
+      expect(display.fileDiff).toMatch(/\+new text\n\+new text\n\+new text/);
       expect(display.fileName).toBe(testFile);
+      expect((result.returnDisplay as FileDiff).diffStat).toStrictEqual({
+        ai_added_lines: 3,
+        ai_removed_lines: 3,
+        user_added_lines: 0,
+        user_removed_lines: 0,
+      });
     });
 
     it('should return error if expected_replacements does not match actual occurrences', async () => {
@@ -562,13 +569,14 @@ describe('EditTool', () => {
     });
 
     it('should include modification message when proposed content is modified', async () => {
-      const initialContent = 'This is some old text.';
+      const initialContent = 'Line 1\nold line\nLine 3\nLine 4\nLine 5\n';
       fs.writeFileSync(filePath, initialContent, 'utf8');
       const params: EditToolParams = {
         file_path: filePath,
         old_string: 'old',
         new_string: 'new',
         modified_by_user: true,
+        ai_proposed_content: 'Line 1\nAI line\nLine 3\nLine 4\nLine 5\n',
       };
 
       (mockConfig.getApprovalMode as Mock).mockReturnValueOnce(
@@ -580,6 +588,12 @@ describe('EditTool', () => {
       expect(result.llmContent).toMatch(
         /User modified the `new_string` content/,
       );
+      expect((result.returnDisplay as FileDiff).diffStat).toStrictEqual({
+        ai_added_lines: 1,
+        ai_removed_lines: 1,
+        user_added_lines: 1,
+        user_removed_lines: 1,
+      });
     });
 
     it('should not include modification message when proposed content is not modified', async () => {
diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts
index cd98ed0070..b1b51b7450 100644
--- a/packages/core/src/tools/edit.ts
+++ b/packages/core/src/tools/edit.ts
@@ -86,9 +86,9 @@ export interface EditToolParams {
   modified_by_user?: boolean;
 
   /**
-   * Initially proposed string.
+   * Initially proposed content.
    */
-  ai_proposed_string?: string;
+  ai_proposed_content?: string;
 }
 
 interface CalculatedEdit {
@@ -365,12 +365,12 @@ class EditToolInvocation implements ToolInvocation<EditToolParams, ToolResult> {
       let displayResult: ToolResultDisplay;
       const fileName = path.basename(this.params.file_path);
       const originallyProposedContent =
-        this.params.ai_proposed_string || this.params.new_string;
+        this.params.ai_proposed_content || editData.newContent;
       const diffStat = getDiffStat(
         fileName,
         editData.currentContent ?? '',
         originallyProposedContent,
-        this.params.new_string,
+        editData.newContent,
       );
 
       if (editData.isNewFile) {
@@ -572,16 +572,13 @@ Expectation for required parameters:
         oldContent: string,
         modifiedProposedContent: string,
         originalParams: EditToolParams,
-      ): EditToolParams => {
-        const content = originalParams.new_string;
-        return {
-          ...originalParams,
-          ai_proposed_string: content,
-          old_string: oldContent,
-          new_string: modifiedProposedContent,
-          modified_by_user: true,
-        };
-      },
+      ): EditToolParams => ({
+        ...originalParams,
+        ai_proposed_content: oldContent,
+        old_string: oldContent,
+        new_string: modifiedProposedContent,
+        modified_by_user: true,
+      }),
     };
   }
 }