diff --git a/agent/tools/edit/edit.py b/agent/tools/edit/edit.py index b6a2b57c..5671b79a 100644 --- a/agent/tools/edit/edit.py +++ b/agent/tools/edit/edit.py @@ -13,7 +13,7 @@ from agent.tools.utils.diff import ( detect_line_ending, normalize_to_lf, restore_line_endings, - normalize_for_fuzzy_match, + count_matches, fuzzy_find_text, generate_diff_string ) @@ -110,10 +110,10 @@ class Edit(BaseTool): "The old text must match exactly including all whitespace and newlines." ) - # Calculate occurrence count (use fuzzy normalized content for consistency) - fuzzy_content = normalize_for_fuzzy_match(normalized_content) - fuzzy_old_text = normalize_for_fuzzy_match(normalized_old_text) - occurrences = fuzzy_content.count(fuzzy_old_text) + # Count occurrences with the same matcher used to locate and + # replace (fuzzy_find_text), so the uniqueness guard cannot + # disagree with what actually gets replaced. + occurrences = count_matches(normalized_content, normalized_old_text) if occurrences > 1: return ToolResult.fail( diff --git a/agent/tools/utils/__init__.py b/agent/tools/utils/__init__.py index 698f75ef..e70cc46c 100644 --- a/agent/tools/utils/__init__.py +++ b/agent/tools/utils/__init__.py @@ -15,6 +15,7 @@ from .diff import ( normalize_to_lf, restore_line_endings, normalize_for_fuzzy_match, + count_matches, fuzzy_find_text, generate_diff_string, FuzzyMatchResult @@ -39,6 +40,7 @@ __all__ = [ 'normalize_to_lf', 'restore_line_endings', 'normalize_for_fuzzy_match', + 'count_matches', 'fuzzy_find_text', 'generate_diff_string', 'FuzzyMatchResult', diff --git a/agent/tools/utils/diff.py b/agent/tools/utils/diff.py index bcb5c694..7002c046 100644 --- a/agent/tools/utils/diff.py +++ b/agent/tools/utils/diff.py @@ -93,6 +93,40 @@ class FuzzyMatchResult: self.content_for_replacement = content_for_replacement +def _build_fuzzy_pattern(old_text: str) -> Optional[str]: + """ + Build the whitespace-flexible regex used to locate ``old_text`` fuzzily. + + Returns ``None`` when ``old_text`` has no non-whitespace content to match. + This is the single source of truth for fuzzy matching, so that *finding* a + match (:func:`fuzzy_find_text`) and *counting* occurrences + (:func:`count_matches`) always use the exact same rules. + """ + stripped = old_text.strip('\n') + if not stripped.strip(): + return None + + source_lines = stripped.split('\n') + line_patterns = [] + for i, line in enumerate(source_lines): + tokens = line.split() + if not tokens: + line_patterns.append(r'[ \t]*') + continue + # Tolerate any run of blanks between tokens. + core = r'[ \t]+'.join(re.escape(tok) for tok in tokens) + # First-line leading whitespace is folded into the match only when + # old_text itself was indented here; otherwise it stays OUTSIDE the + # match so a no-indent old_text preserves (does not swallow and drop) + # the file's existing indentation -- mirroring an exact substring + # match. Inner lines always tolerate indentation: it sits inside the + # matched region and is re-supplied by new_text. + if i > 0 or line[:1] in (' ', '\t'): + core = r'[ \t]*' + core + line_patterns.append(core + r'[ \t]*') + return '\n'.join(line_patterns) + + def fuzzy_find_text(content: str, old_text: str) -> FuzzyMatchResult: """ Find text in content, try exact match first, then fuzzy match @@ -110,7 +144,7 @@ def fuzzy_find_text(content: str, old_text: str) -> FuzzyMatchResult: match_length=len(old_text), content_for_replacement=content ) - + # Fuzzy match: the exact substring was not found, most likely because the # whitespace differs (indentation, spaces around operators, trailing # spaces). Locate the region in the ORIGINAL content using a @@ -121,27 +155,8 @@ def fuzzy_find_text(content: str, old_text: str) -> FuzzyMatchResult: # doing so previously returned the normalized copy as # content_for_replacement, which caused the whole file to be rewritten # with collapsed indentation (every untouched line got reformatted). - stripped = old_text.strip('\n') - if stripped.strip(): - source_lines = stripped.split('\n') - line_patterns = [] - for i, line in enumerate(source_lines): - tokens = line.split() - if not tokens: - line_patterns.append(r'[ \t]*') - continue - # Tolerate any run of blanks between tokens. - core = r'[ \t]+'.join(re.escape(tok) for tok in tokens) - # First-line leading whitespace is folded into the match only when - # old_text itself was indented here; otherwise it stays OUTSIDE the - # match so a no-indent old_text preserves (does not swallow and drop) - # the file's existing indentation -- mirroring an exact substring - # match. Inner lines always tolerate indentation: it sits inside the - # matched region and is re-supplied by new_text. - if i > 0 or line[:1] in (' ', '\t'): - core = r'[ \t]*' + core - line_patterns.append(core + r'[ \t]*') - pattern = '\n'.join(line_patterns) + pattern = _build_fuzzy_pattern(old_text) + if pattern is not None: match = re.search(pattern, content) if match: return FuzzyMatchResult( @@ -155,6 +170,28 @@ def fuzzy_find_text(content: str, old_text: str) -> FuzzyMatchResult: return FuzzyMatchResult(found=False) +def count_matches(content: str, old_text: str) -> int: + """ + Count occurrences of ``old_text`` using the SAME strategy as + :func:`fuzzy_find_text`: an exact substring when one is present, otherwise + the whitespace-flexible fuzzy regex. + + The edit tool's uniqueness guard must agree with the matcher that actually + performs the replacement. Counting through a separate normalization pass + (the previous approach) could disagree with the regex used to locate and + replace, so both paths now share :func:`_build_fuzzy_pattern`. + """ + if not old_text: + return 0 + # Mirror fuzzy_find_text: prefer exact matching when it applies. + if content.find(old_text) != -1: + return content.count(old_text) + pattern = _build_fuzzy_pattern(old_text) + if pattern is None: + return 0 + return len(re.findall(pattern, content)) + + def generate_diff_string(old_content: str, new_content: str) -> dict: """ Generate unified diff string diff --git a/tests/test_edit_tool.py b/tests/test_edit_tool.py index 24b77759..f558929a 100644 --- a/tests/test_edit_tool.py +++ b/tests/test_edit_tool.py @@ -98,6 +98,41 @@ class TestEditFuzzyPreservesWhitespace(unittest.TestCase): "def foo():\n x = 100\n y = 2\n return x + y\n", ) + def test_exact_match_rejects_multiple_occurrences(self): + # Two byte-identical statements; the exact-match path applies and the + # uniqueness guard counts exact occurrences, so the ambiguous edit is + # rejected instead of silently editing only the first. + with open(self.path, "w", encoding="utf-8") as f: + f.write("a = 1\nb = 2\na = 1\n") + result = self.tool.execute({ + "path": self.path, + "oldText": "a = 1", + "newText": "a = 9", + }) + self.assertEqual(result.status, "error", result.result) + self.assertIn("occurrences", result.result) + # An ambiguous match must leave the file untouched. + self.assertEqual(self._read(), "a = 1\nb = 2\na = 1\n") + + def test_fuzzy_match_rejects_multiple_occurrences(self): + # oldText uses loose spacing, so the exact match fails and the fuzzy + # path runs. The uniqueness guard now counts with the SAME regex used + # to match/replace, so an ambiguous fuzzy match (two hits) is rejected + # rather than silently editing the first one. + with open(self.path, "w", encoding="utf-8") as f: + f.write("def foo():\n x = 1\n y = 2\n x = 1\n") + result = self.tool.execute({ + "path": self.path, + "oldText": "x = 1", + "newText": "x = 99", + }) + self.assertEqual(result.status, "error", result.result) + self.assertIn("occurrences", result.result) + self.assertEqual( + self._read(), + "def foo():\n x = 1\n y = 2\n x = 1\n", + ) + if __name__ == "__main__": unittest.main()