fix: max_results can result in returning incomplete matches

Author: sebthom

.github/workflows/test.yml

@@ -63,7 +63,7 @@ jobs:
 
     - name: Run unit tests
       run: |
-        uv run pytest tests/test_main.py -v --cov=main --cov-report=term-missing
+        uv run pytest tests/test_unit.py -v --cov=main --cov-report=term-missing
 
     - name: Run integration tests
       run: |

README.md

@@ -132,9 +132,23 @@ Test ast-grep YAML rules against code snippets before applying them to larger co
 Search codebases using simple ast-grep patterns for straightforward structural matches.
 
 **Parameters:**
-- `max_results`: Limit number of results (default: unlimited)
+- `max_results`: Limit number of complete matches returned (default: unlimited)
 - `output_format`: Choose between `"text"` (default, ~75% fewer tokens) or `"json"` (full metadata)
 
+**Text Output Format:**
+```
+Found 2 matches:
+
+path/to/file.py:10-15
+def example_function():
+    # function body
+    return result
+
+path/to/file.py:20-22
+def another_function():
+    pass
+```
+
 **Use cases:**
 - Find function calls with specific patterns
 - Locate variable declarations
@@ -144,7 +158,7 @@ Search codebases using simple ast-grep patterns for straightforward structural m
 Advanced codebase search using complex YAML rules that can express sophisticated matching criteria.
 
 **Parameters:**
-- `max_results`: Limit number of results (default: unlimited)
+- `max_results`: Limit number of complete matches returned (default: unlimited)
 - `output_format`: Choose between `"text"` (default, ~75% fewer tokens) or `"json"` (full metadata)
 
 **Use cases:**

main.py

@@ -111,8 +111,23 @@ def find_code(
     Internally calls: ast-grep run --pattern <pattern> [--json] <project_folder>
 
     Output formats:
-    - text (default): Simple file:line:content format, ~75% fewer tokens
-    - json: Full match objects with metadata
+    - text (default): Compact text format with file:line-range headers and complete match text
+      Example:
+        Found 2 matches:
+
+        path/to/file.py:10-15
+        def example_function():
+            # function body
+            return result
+
+        path/to/file.py:20-22
+        def another_function():
+            pass
+
+    - json: Full match objects with metadata including ranges, meta-variables, etc.
+
+    The max_results parameter limits the number of complete matches returned (not individual lines).
+    When limited, the header shows "Found X matches (showing first Y of Z)".
 
     Example usage:
       find_code(pattern="class $NAME", max_results=20)  # Returns text format
@@ -125,33 +140,24 @@ def find_code(
     if language:
         args.extend(["--lang", language])
 
-    if output_format == "json":
-        # JSON format - return structured data
-        result = run_ast_grep("run", args + ["--json", project_folder])
-        matches = json.loads(result.stdout.strip() or "[]")
-        # Limit results if max_results is specified
-        if max_results is not None and len(matches) > max_results:
-            matches = matches[:max_results]
-        return matches  # type: ignore[no-any-return]
-    else:
-        # Text format - return plain text output
-        result = run_ast_grep("run", args + [project_folder])
-        output = result.stdout.strip()
-        if not output:
-            output = "No matches found"
-        else:
-            # Apply max_results limit if specified
-            lines = output.split('\n')
-            non_empty_lines = [line for line in lines if line.strip()]
-            if max_results is not None and len(non_empty_lines) > max_results:
-                # Limit the results
-                non_empty_lines = non_empty_lines[:max_results]
-                output = '\n'.join(non_empty_lines)
-                header = f"Found {len(non_empty_lines)} matches (limited to {max_results}):\n"
-            else:
-                header = f"Found {len(non_empty_lines)} matches:\n"
-            output = header + output
-        return output  # type: ignore[no-any-return]
+    # Always get JSON internally for accurate match limiting
+    result = run_ast_grep("run", args + ["--json", project_folder])
+    matches = json.loads(result.stdout.strip() or "[]")
+
+    # Apply max_results limit to complete matches
+    total_matches = len(matches)
+    if max_results is not None and total_matches > max_results:
+        matches = matches[:max_results]
+
+    if output_format == "text":
+        if not matches:
+            return "No matches found"
+        text_output = format_matches_as_text(matches)
+        header = f"Found {len(matches)} matches"
+        if max_results is not None and total_matches > max_results:
+            header += f" (showing first {max_results} of {total_matches})"
+        return header + ":\n\n" + text_output
+    return matches  # type: ignore[no-any-return]
 
 @mcp.tool()
 def find_code_by_rule(
@@ -170,8 +176,23 @@ def find_code_by_rule(
     Internally calls: ast-grep scan --inline-rules <yaml> [--json] <project_folder>
 
     Output formats:
-    - text (default): Simple file:line:content format, ~75% fewer tokens
-    - json: Full match objects with metadata
+    - text (default): Compact text format with file:line-range headers and complete match text
+      Example:
+        Found 2 matches:
+
+        src/models.py:45-52
+        class UserModel:
+            def __init__(self):
+                self.id = None
+                self.name = None
+
+        src/views.py:12
+        class SimpleView: pass
+
+    - json: Full match objects with metadata including ranges, meta-variables, etc.
+
+    The max_results parameter limits the number of complete matches returned (not individual lines).
+    When limited, the header shows "Found X matches (showing first Y of Z)".
 
     Example usage:
       find_code_by_rule(yaml="id: x\\nlanguage: python\\nrule: {pattern: 'class $NAME'}", max_results=20)
@@ -182,33 +203,50 @@ def find_code_by_rule(
 
     args = ["--inline-rules", yaml]
 
-    if output_format == "json":
-        # JSON format - return structured data
-        result = run_ast_grep("scan", args + ["--json", project_folder])
-        matches = json.loads(result.stdout.strip() or "[]")
-        # Limit results if max_results is specified
-        if max_results is not None and len(matches) > max_results:
-            matches = matches[:max_results]
-        return matches  # type: ignore[no-any-return]
-    else:
-        # Text format - return plain text output
-        result = run_ast_grep("scan", args + [project_folder])
-        output = result.stdout.strip()
-        if not output:
-            output = "No matches found"
+    # Always get JSON internally for accurate match limiting
+    result = run_ast_grep("scan", args + ["--json", project_folder])
+    matches = json.loads(result.stdout.strip() or "[]")
+
+    # Apply max_results limit to complete matches
+    total_matches = len(matches)
+    if max_results is not None and total_matches > max_results:
+        matches = matches[:max_results]
+
+    if output_format == "text":
+        if not matches:
+            return "No matches found"
+        text_output = format_matches_as_text(matches)
+        header = f"Found {len(matches)} matches"
+        if max_results is not None and total_matches > max_results:
+            header += f" (showing first {max_results} of {total_matches})"
+        return header + ":\n\n" + text_output
+    return matches  # type: ignore[no-any-return]
+
+def format_matches_as_text(matches: List[dict]) -> str:
+    """Convert JSON matches to LLM-friendly text format.
+
+    Format: file:start-end followed by the complete match text.
+    Matches are separated by blank lines for clarity.
+    """
+    if not matches:
+        return ""
+
+    output_blocks = []
+    for m in matches:
+        file_path = m.get('file', '')
+        start_line = m.get('range', {}).get('start', {}).get('line', 0) + 1
+        end_line = m.get('range', {}).get('end', {}).get('line', 0) + 1
+        match_text = m.get('text', '').rstrip()
+
+        # Format: filepath:start-end (or just :line for single-line matches)
+        if start_line == end_line:
+            header = f"{file_path}:{start_line}"
         else:
-            # Apply max_results limit if specified
-            lines = output.split('\n')
-            non_empty_lines = [line for line in lines if line.strip()]
-            if max_results is not None and len(non_empty_lines) > max_results:
-                # Limit the results
-                non_empty_lines = non_empty_lines[:max_results]
-                output = '\n'.join(non_empty_lines)
-                header = f"Found {len(non_empty_lines)} matches (limited to {max_results}):\n"
-            else:
-                header = f"Found {len(non_empty_lines)} matches:\n"
-            output = header + output
-        return output  # type: ignore[no-any-return]
+            header = f"{file_path}:{start_line}-{end_line}"
+
+        output_blocks.append(f"{header}\n{match_text}")
+
+    return '\n\n'.join(output_blocks)
 
 def run_command(args: List[str], input_text: Optional[str] = None) -> subprocess.CompletedProcess:
     try:

tests/test_integration.py

@@ -1,5 +1,6 @@
 """Integration tests for ast-grep MCP server"""
 
+import json
 import os
 import sys
 from unittest.mock import Mock, patch
@@ -79,9 +80,14 @@ def test_find_code_json_format(self, fixtures_dir):
     @patch("main.run_ast_grep")
     def test_find_code_by_rule(self, mock_run, fixtures_dir):
         """Test find_code_by_rule with mocked ast-grep"""
-        # Mock the response
+        # Mock the response with JSON format (since we always use JSON internally)
         mock_result = Mock()
-        mock_result.stdout = "fixtures/example.py:7:class Calculator:"
+        mock_matches = [{
+            "text": "class Calculator:\n    pass",
+            "file": "fixtures/example.py",
+            "range": {"start": {"line": 6}, "end": {"line": 7}}
+        }]
+        mock_result.stdout = json.dumps(mock_matches)
         mock_run.return_value = mock_result
 
         yaml_rule = """id: test
@@ -95,10 +101,11 @@ def test_find_code_by_rule(self, mock_run, fixtures_dir):
 
         assert "Calculator" in result
         assert "Found 1 match" in result
+        assert "fixtures/example.py:7-8" in result
 
         # Verify the command was called correctly
         mock_run.assert_called_once_with(
-            "scan", ["--inline-rules", yaml_rule, fixtures_dir]
+            "scan", ["--inline-rules", yaml_rule, "--json", fixtures_dir]
         )
 
     def test_find_code_with_max_results(self, fixtures_dir):
@@ -111,7 +118,8 @@ def test_find_code_with_max_results(self, fixtures_dir):
             output_format="text",
         )
 
-        assert "limited to 1" in result
+        # The new format says "showing first X of Y" instead of "limited to X"
+        assert "showing first 1 of" in result or "Found 1 match" in result
         # Should only have one match in the output
         assert result.count("def ") == 1