Optimize tryToDoStringReplacementWithExtraIndentation with first-line fast bailout

packages/agent-runtime/src/__tests__/generate-diffs-prompt.test.ts

@@ -104,4 +104,36 @@ describe('tryToDoStringReplacementWithExtraIndentation', () => {
     expect(result!.searchContent).toBe(' const x = 1;\n')
     expect(result!.replaceContent).toBe(' const x = 2;\n')
   })
+
+  it('should return null when first line matches with indentation but subsequent lines mismatch', () => {
+    const oldFileContent = '  function foo() {\n    return 999;\n  }\n'
+    const searchContent = 'function foo() {\n  return 1;\n}\n'
+    const replaceContent = 'function foo() {\n  return 2;\n}\n'
+
+    const result = tryToDoStringReplacementWithExtraIndentation({
+      oldFileContent,
+      searchContent,
+      replaceContent,
+    })
+
+    expect(result).toBeNull()
+  })
+
+  it('should handle searchContent with leading empty lines and extra indentation', () => {
+    const oldFileContent = '\n    const a = 1;\n    const b = 2;\n'
+    const searchContent = '\nconst a = 1;\nconst b = 2;\n'
+    const replaceContent = '\nconst a = 10;\nconst b = 20;\n'
+
+    const result = tryToDoStringReplacementWithExtraIndentation({
+      oldFileContent,
+      searchContent,
+      replaceContent,
+    })
+
+    expect(result).not.toBeNull()
+    expect(result!.searchContent).toBe('\n    const a = 1;\n    const b = 2;\n')
+    expect(result!.replaceContent).toBe(
+      '\n    const a = 10;\n    const b = 20;\n',
+    )
+  })
 })

packages/agent-runtime/src/generate-diffs-prompt.ts

@@ -4,35 +4,45 @@ export const tryToDoStringReplacementWithExtraIndentation = (params: {
   replaceContent: string
 }) => {
   const { oldFileContent, searchContent, replaceContent } = params
+  const searchLines = searchContent.split('\n')
+  const firstNonEmptyLine = searchLines.find((line) => Boolean(line))
+
+  const indentLines = (lines: string[], prefix: string) =>
+    lines.map((line) => (line ? prefix + line : line)).join('\n')
+
   for (let i = 1; i <= 12; i++) {
-    const searchContentWithIndentation = searchContent
-      .split('\n')
-      .map((line) => (line ? ' '.repeat(i) + line : line))
-      .join('\n')
+    const prefix = ' '.repeat(i)
+    if (
+      firstNonEmptyLine !== undefined &&
+      !oldFileContent.includes(prefix + firstNonEmptyLine)
+    ) {
+      continue
+    }
+    const searchContentWithIndentation = indentLines(searchLines, prefix)
     if (oldFileContent.includes(searchContentWithIndentation)) {
       return {
         searchContent: searchContentWithIndentation,
-        replaceContent: replaceContent
-          .split('\n')
-          .map((line) => (line ? ' '.repeat(i) + line : line))
-          .join('\n'),
+        replaceContent: indentLines(replaceContent.split('\n'), prefix),
       }
     }
   }
+
   for (let i = 1; i <= 6; i++) {
-    const searchContentWithIndentation = searchContent
-      .split('\n')
-      .map((line) => (line ? '\t'.repeat(i) + line : line))
-      .join('\n')
+    const prefix = '\t'.repeat(i)
+    if (
+      firstNonEmptyLine !== undefined &&
+      !oldFileContent.includes(prefix + firstNonEmptyLine)
+    ) {
+      continue
+    }
+    const searchContentWithIndentation = indentLines(searchLines, prefix)
     if (oldFileContent.includes(searchContentWithIndentation)) {
       return {
         searchContent: searchContentWithIndentation,
-        replaceContent: replaceContent
-          .split('\n')
-          .map((line) => (line ? '\t'.repeat(i) + line : line))
-          .join('\n'),
+        replaceContent: indentLines(replaceContent.split('\n'), prefix),
       }
     }
   }
+
   return null
 }