[AI Assistance] Use unified diff to patching

Bug: 399561340
Change-Id: I895a443f08e2c3e5eeccdae55c3cd3b4d208558b
Reviewed-on: https://chromium-review.googlesource.com/c/devtools/devtools-frontend/+/6434494
Reviewed-by: Alex Rudenko <[email protected]>
Commit-Queue: Nikolay Vitkov <[email protected]>

Author: Lightning00Blade

front_end/core/host/AidaClient.ts

@@ -228,7 +228,7 @@ export enum CitationSourceType {
   TRAINING_DATA = 'TRAINING_DATA',
   WORLD_FACTS = 'WORLD_FACTS',
   LOCAL_FACTS = 'LOCAL_FACTS',
-  INDIRECT = 'INDERECT',
+  INDIRECT = 'INDIRECT',
 }
 
 export interface Citation {

front_end/models/ai_assistance/AgentProject.test.ts

@@ -65,10 +65,198 @@ describeWithEnvironment('AgentProject', () => {
     assert.deepEqual(project.getProcessedFiles(), ['index.html']);
   });
 
-  it('can write files files', async () => {
-    const {project} = await mockProject();
-    project.writeFile('index.html', 'updated');
-    assert.deepEqual(project.readFile('index.html'), 'updated');
+  describe('write file', () => {
+    describe('full', () => {
+      it('can write files files', async () => {
+        const {project} = await mockProject();
+        project.writeFile('index.html', 'updated');
+        assert.deepEqual(project.readFile('index.html'), 'updated');
+      });
+    });
+
+    describe('unified', () => {
+      it('can write files', async () => {
+        const {project} = await mockProject();
+        const unifiedDiff = `\`\`\`\`\`
+diff
+--- a/index.html
++++ b/index.html
+@@ -817,5 +817,5 @@
+-content
++updated
+\`\`\`\`\``;
+
+        project.writeFile('index.html', unifiedDiff, AiAssistanceModel.ReplaceStrategy.UNIFIED_DIFF);
+        assert.deepEqual(project.readFile('index.html'), 'updated');
+      });
+
+      it('can write files with multiple changes', async () => {
+        const {project} = await mockProject(
+            [
+              {
+                path: 'index.css',
+                content: `Line:1
+Line:2
+Line:3
+Line:4
+Line:5`,
+              },
+            ],
+        );
+        const unifiedDiff = `\`\`\`\`\`
+diff
+--- a/index.css
++++ b/index.css
+@@ -817,1 +817,1 @@
+-Line:1
++LineUpdated:1
+@@ -856,7 +857,7 @@
+-Line:4
++LineUpdated:4
+\`\`\`\`\``;
+
+        project.writeFile('index.css', unifiedDiff, AiAssistanceModel.ReplaceStrategy.UNIFIED_DIFF);
+        assert.deepEqual(project.readFile('index.css'), `LineUpdated:1
+Line:2
+Line:3
+LineUpdated:4
+Line:5`);
+      });
+
+      it('can write files with only addition', async () => {
+        const {project} = await mockProject(
+            [
+              {
+                path: 'index.css',
+                content: '',
+              },
+            ],
+        );
+        const unifiedDiff = `\`\`\`\`\`
+diff
+--- a/index.css
++++ b/index.css
+@@ -817,1 +817,1 @@
++Line:1
++Line:4
+\`\`\`\`\``;
+
+        project.writeFile('index.css', unifiedDiff, AiAssistanceModel.ReplaceStrategy.UNIFIED_DIFF);
+        assert.deepEqual(project.readFile('index.css'), `Line:1
+Line:4`);
+      });
+
+      it('can write files with multiple additions', async () => {
+        const {project} = await mockProject(
+            [
+              {
+                path: 'index.css',
+                content: `Line:1
+Line:2
+Line:3
+Line:4
+Line:5`,
+              },
+            ],
+        );
+        const unifiedDiff = `\`\`\`\`\`
+diff
+--- a/index.css
++++ b/index.css
+@@ -817,1 +817,1 @@
+-Line:1
++LineUpdated:1
++LineUpdated:1.5
+\`\`\`\`\``;
+
+        project.writeFile('index.css', unifiedDiff, AiAssistanceModel.ReplaceStrategy.UNIFIED_DIFF);
+        assert.deepEqual(project.readFile('index.css'), `LineUpdated:1
+LineUpdated:1.5
+Line:2
+Line:3
+Line:4
+Line:5`);
+      });
+
+      it('can write files with only deletion', async () => {
+        const {project} = await mockProject(
+            [
+              {
+                path: 'index.css',
+                content: `Line:1
+Line:2
+Line:3
+Line:4
+Line:5`,
+              },
+            ],
+        );
+        const unifiedDiff = `\`\`\`\`\`
+diff
+--- a/index.css
++++ b/index.css
+@@ -817,1 +817,1 @@
+ Line:1
+-Line:2
+ Line:3
+\`\`\`\`\``;
+
+        project.writeFile('index.css', unifiedDiff, AiAssistanceModel.ReplaceStrategy.UNIFIED_DIFF);
+        assert.deepEqual(project.readFile('index.css'), `Line:1
+Line:3
+Line:4
+Line:5`);
+      });
+
+      it('can write files with only deletion no search lines', async () => {
+        const {project} = await mockProject(
+            [
+              {
+                path: 'index.css',
+                content: 'Line:1',
+              },
+            ],
+        );
+        const unifiedDiff = `\`\`\`\`\`
+diff
+--- a/index.css
++++ b/index.css
+@@ -817,1 +817,1 @@
+-Line:1
+\`\`\`\`\``;
+
+        project.writeFile('index.css', unifiedDiff, AiAssistanceModel.ReplaceStrategy.UNIFIED_DIFF);
+        assert.deepEqual(project.readFile('index.css'), '');
+      });
+
+      it('can write files with first line next to @@', async () => {
+        const {project} = await mockProject(
+            [
+              {
+                path: 'index.css',
+                content: `Line:1
+Line:2
+Line:3
+Line:4
+Line:5`,
+              },
+            ],
+        );
+        const unifiedDiff = `\`\`\`\`\`
+diff
+--- a/index.css
++++ b/index.css
+@@ -817,1 +817,1 @@-Line:1
+-Line:2
+ Line:3
+\`\`\`\`\``;
+
+        project.writeFile('index.css', unifiedDiff, AiAssistanceModel.ReplaceStrategy.UNIFIED_DIFF);
+        assert.deepEqual(project.readFile('index.css'), `Line:3
+Line:4
+Line:5`);
+      });
+    });
   });
 
   describe('limits', () => {

front_end/models/ai_assistance/AgentProject.ts

@@ -9,6 +9,13 @@ import type * as Workspace from '../workspace/workspace.js';
 
 import {debugLog} from './debug.js';
 
+const LINE_END_RE = /\r\n?|\n/;
+
+export const enum ReplaceStrategy {
+  FULL_FILE = 'full',
+  UNIFIED_DIFF = 'unified'
+}
+
 /**
  * AgentProject wraps around a Workspace.Workspace.Project and
  * implements AI Assistance-specific logic for accessing workspace files
@@ -18,7 +25,7 @@ export class AgentProject {
   #project: Workspace.Workspace.Project;
   #ignoredFolderNames = new Set(['node_modules']);
   #filesChanged = new Set<string>();
-  #linesChanged = 0;
+  #totalLinesChanged = 0;
 
   readonly #maxFilesChanged: number;
   readonly #maxLinesChanged: number;
@@ -75,27 +82,26 @@ export class AgentProject {
    * This method updates the file content in the working copy of the
    * UiSourceCode identified by the filepath.
    */
-  writeFile(filepath: string, content: string): void {
+  writeFile(filepath: string, update: string, mode = ReplaceStrategy.FULL_FILE): void {
     const {map} = this.#indexFiles();
     const uiSourceCode = map.get(filepath);
     if (!uiSourceCode) {
       throw new Error(`UISourceCode ${filepath} not found`);
     }
     const currentContent = this.readFile(filepath);
-    const lineEndRe = /\r\n?|\n/;
-    let linesChanged = 0;
-    if (currentContent) {
-      const diff = Diff.Diff.DiffWrapper.lineDiff(currentContent.split(lineEndRe), content.split(lineEndRe));
-      for (const item of diff) {
-        if (item[0] !== Diff.Diff.Operation.Equal) {
-          linesChanged++;
-        }
-      }
-    } else {
-      linesChanged += content.split(lineEndRe).length;
+    let content: string;
+    switch (mode) {
+      case ReplaceStrategy.FULL_FILE:
+        content = update;
+        break;
+      case ReplaceStrategy.UNIFIED_DIFF:
+        content = this.#writeWithUnifiedDiff(update, currentContent);
+        break;
     }
 
-    if (this.#linesChanged + linesChanged > this.#maxLinesChanged) {
+    const linesChanged = this.getLinesChanged(currentContent, content);
+
+    if (this.#totalLinesChanged + linesChanged > this.#maxLinesChanged) {
       throw new Error('Too many lines changed');
     }
 
@@ -104,11 +110,92 @@ export class AgentProject {
       this.#filesChanged.delete(filepath);
       throw new Error('Too many files changed');
     }
-    this.#linesChanged += linesChanged;
+    this.#totalLinesChanged += linesChanged;
     uiSourceCode.setWorkingCopy(content);
     uiSourceCode.setContainsAiChanges(true);
   }
 
+  #writeWithUnifiedDiff(llmDiff: string, content = ''): string {
+    let updatedContent = content;
+    const diffChunk = llmDiff.trim();
+    const normalizedDiffLines = diffChunk.split(LINE_END_RE);
+
+    const lineAfterSeparatorRegEx = /^@@.*@@([- +].*)/;
+    const changeChunk: string[][] = [];
+    let currentChunk: string[] = [];
+    for (const line of normalizedDiffLines) {
+      if (line.startsWith('```')) {
+        continue;
+      }
+
+      // The ending is not always @@
+      if (line.startsWith('@@')) {
+        line.search('@@');
+        currentChunk = [];
+        changeChunk.push(currentChunk);
+        if (!line.endsWith('@@')) {
+          const match = line.match(lineAfterSeparatorRegEx);
+          if (match?.[1]) {
+            currentChunk.push(match[1]);
+          }
+        }
+      } else {
+        currentChunk.push(line);
+      }
+    }
+
+    for (const chunk of changeChunk) {
+      const search = [];
+      const replace = [];
+      for (const changeLine of chunk) {
+        // Unified diff first char is ' ', '-', '+'
+        // to represent what happened to the line
+        const line = changeLine.slice(1);
+
+        if (changeLine.startsWith('-')) {
+          search.push(line);
+        } else if (changeLine.startsWith('+')) {
+          replace.push(line);
+        } else {
+          search.push(line);
+          replace.push(line);
+        }
+      }
+      if (replace.length === 0) {
+        const searchString = search.join('\n');
+        // If we remove we want to
+        if (updatedContent.search(searchString + '\n') !== -1) {
+          updatedContent = updatedContent.replace(searchString + '\n', '');
+        } else {
+          updatedContent = updatedContent.replace(searchString, '');
+        }
+      } else if (search.length === 0) {
+        // This just adds it to the beginning of the file
+        updatedContent = updatedContent.replace('', replace.join('\n'));
+      } else {
+        updatedContent = updatedContent.replace(search.join('\n'), replace.join('\n'));
+      }
+    }
+
+    return updatedContent;
+  }
+
+  getLinesChanged(currentContent: string|undefined, updatedContent: string): number {
+    let linesChanged = 0;
+    if (currentContent) {
+      const diff = Diff.Diff.DiffWrapper.lineDiff(updatedContent.split(LINE_END_RE), currentContent.split(LINE_END_RE));
+      for (const item of diff) {
+        if (item[0] !== Diff.Diff.Operation.Equal) {
+          linesChanged++;
+        }
+      }
+    } else {
+      linesChanged += updatedContent.split(LINE_END_RE).length;
+    }
+
+    return linesChanged;
+  }
+
   /**
    * This method searches in files for the agent and provides the
    * matches to the agent.
@@ -159,11 +246,11 @@ export class AgentProject {
     const map = new Map();
     // TODO: this could be optimized and cached.
     for (const uiSourceCode of this.#project.uiSourceCodes()) {
-      const pathParths = Persistence.FileSystemWorkspaceBinding.FileSystemWorkspaceBinding.relativePath(uiSourceCode);
-      if (this.#shouldSkipPath(pathParths)) {
+      const pathParts = Persistence.FileSystemWorkspaceBinding.FileSystemWorkspaceBinding.relativePath(uiSourceCode);
+      if (this.#shouldSkipPath(pathParts)) {
         continue;
       }
-      const path = pathParths.join('/');
+      const path = pathParts.join('/');
       files.push(path);
       map.set(path, uiSourceCode);
     }

front_end/models/ai_assistance/agents/AiAgent.ts

@@ -313,12 +313,12 @@ export abstract class AiAgent<T> {
     }
     const enableAidaFunctionCalling = declarations.length && !this.functionCallEmulationEnabled;
     const userTier = Host.AidaClient.convertToUserTierEnum(this.userTier);
-    const premable = userTier === Host.AidaClient.UserTier.TESTERS ? this.preamble : undefined;
+    const preamble = userTier === Host.AidaClient.UserTier.TESTERS ? this.preamble : undefined;
     const facts = Array.from(this.#facts);
     const request: Host.AidaClient.AidaRequest = {
       client: Host.AidaClient.CLIENT_NAME,
       current_message: currentMessage,
-      preamble: premable,
+      preamble,
 
       historical_contexts: history.length ? history : undefined,
       facts: facts.length ? facts : undefined,