fix(search): paginate tool responses (#48)

Author: johnhuang316

src/code_index_mcp/server.py

@@ -15,7 +15,7 @@
 import logging
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
-from typing import AsyncIterator, Dict, Any, List
+from typing import AsyncIterator, Dict, Any, List, Optional
 
 # Third-party imports
 from mcp.server.fastmcp import FastMCP, Context
@@ -164,10 +164,12 @@ def search_code_advanced(
     context_lines: int = 0,
     file_pattern: str = None,
     fuzzy: bool = False,
-    regex: bool = None
+    regex: bool = None,
+    start_index: int = 0,
+        max_results: Optional[int] = 10
 ) -> Dict[str, Any]:
     """
-    Search for a code pattern in the project using an advanced, fast tool.
+    Search for a code pattern in the project using an advanced, fast tool with pagination support.
 
     This tool automatically selects the best available command-line search tool
     (like ugrep, ripgrep, ag, or grep) for maximum performance.
@@ -195,9 +197,15 @@ def search_code_advanced(
                - If False, forces literal string search
                - If None (default), automatically detects regex patterns and enables regex for patterns like "ERROR|WARN"
                The pattern will always be validated for safety to prevent ReDoS attacks.
+        start_index: Zero-based offset into the flattened match list. Use to fetch subsequent pages.
+        max_results: Maximum number of matches to return (default 10). Pass None to retrieve all matches.
 
     Returns:
-        A dictionary containing the search results or an error message.
+        A dictionary containing:
+        - results: List of matches with file, line, and text keys.
+        - pagination: Metadata with total_matches, returned, start_index, end_index, has_more,
+                      and optionally max_results.
+        If an error occurs, an error message is returned instead.
 
     """
     return SearchService(ctx).search_code(
@@ -206,7 +214,9 @@ def search_code_advanced(
         context_lines=context_lines,
         file_pattern=file_pattern,
         fuzzy=fuzzy,
-        regex=regex
+        regex=regex,
+        start_index=start_index,
+        max_results=max_results
     )
 
 @mcp.tool()

src/code_index_mcp/services/search_service.py

@@ -6,7 +6,7 @@
 """
 
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 from .base_service import BaseService
 from ..utils import FileFilter, ResponseFormatter, ValidationHelper
@@ -20,14 +20,16 @@ def __init__(self, ctx):
         super().__init__(ctx)
         self.file_filter = self._create_file_filter()
 
-    def search_code(  # pylint: disable=too-many-arguments
+    def search_code(  # pylint: disable=too-many-arguments, too-many-locals
         self,
         pattern: str,
         case_sensitive: bool = True,
         context_lines: int = 0,
         file_pattern: Optional[str] = None,
         fuzzy: bool = False,
-        regex: Optional[bool] = None
+        regex: Optional[bool] = None,
+        start_index: int = 0,
+        max_results: Optional[int] = 10
     ) -> Dict[str, Any]:
         """Search for code patterns in the project."""
         self._require_project_setup()
@@ -44,6 +46,10 @@ def search_code(  # pylint: disable=too-many-arguments
             if error:
                 raise ValueError(f"Invalid file pattern: {error}")
 
+        pagination_error = ValidationHelper.validate_pagination(start_index, max_results)
+        if pagination_error:
+            raise ValueError(pagination_error)
+
         if not self.settings:
             raise ValueError("Settings not available")
 
@@ -64,7 +70,15 @@ def search_code(  # pylint: disable=too-many-arguments
                 regex=regex
             )
             filtered = self._filter_results(results)
-            return ResponseFormatter.search_results_response(filtered)
+            formatted_results, pagination = self._paginate_results(
+                filtered,
+                start_index=start_index,
+                max_results=max_results
+            )
+            return ResponseFormatter.search_results_response(
+                formatted_results,
+                pagination
+            )
         except Exception as exc:
             raise ValueError(f"Search failed using '{strategy.name}': {exc}") from exc
 
@@ -168,3 +182,88 @@ def _filter_results(self, results: Dict[str, Any]) -> Dict[str, Any]:
                 continue
 
         return filtered
+
+    def _paginate_results(
+        self,
+        results: Dict[str, Any],
+        start_index: int,
+        max_results: Optional[int]
+    ) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
+        """Apply pagination to search results and format them for responses."""
+        total_matches = 0
+        for matches in results.values():
+            if isinstance(matches, (list, tuple)):
+                total_matches += len(matches)
+
+        effective_start = min(max(start_index, 0), total_matches)
+
+        if total_matches == 0 or effective_start >= total_matches:
+            pagination = self._build_pagination_metadata(
+                total_matches=total_matches,
+                returned=0,
+                start_index=effective_start,
+                max_results=max_results
+            )
+            return [], pagination
+
+        collected: List[Dict[str, Any]] = []
+        current_index = 0
+
+        sorted_items = sorted(
+            (
+                (path, matches)
+                for path, matches in results.items()
+                if isinstance(path, str) and isinstance(matches, (list, tuple))
+            ),
+            key=lambda item: item[0]
+        )
+
+        for path, matches in sorted_items:
+            sorted_matches = sorted(
+                (match for match in matches if isinstance(match, (list, tuple)) and len(match) >= 2),
+                key=lambda pair: pair[0]
+            )
+
+            for line_number, content, *_ in sorted_matches:
+                if current_index >= effective_start:
+                    if max_results is None or len(collected) < max_results:
+                        collected.append({
+                            "file": path,
+                            "line": line_number,
+                            "text": content
+                        })
+                    else:
+                        break
+                current_index += 1
+            if max_results is not None and len(collected) >= max_results:
+                break
+
+        pagination = self._build_pagination_metadata(
+            total_matches=total_matches,
+            returned=len(collected),
+            start_index=effective_start,
+            max_results=max_results
+        )
+        return collected, pagination
+
+    @staticmethod
+    def _build_pagination_metadata(
+        total_matches: int,
+        returned: int,
+        start_index: int,
+        max_results: Optional[int]
+    ) -> Dict[str, Any]:
+        """Construct pagination metadata for search responses."""
+        end_index = start_index + returned
+        metadata: Dict[str, Any] = {
+            "total_matches": total_matches,
+            "returned": returned,
+            "start_index": start_index,
+            "has_more": end_index < total_matches
+        }
+
+        if max_results is not None:
+            metadata["max_results"] = max_results
+
+        metadata["end_index"] = end_index
+        return metadata

src/code_index_mcp/utils/response_formatter.py

@@ -132,7 +132,10 @@ def file_list_response(files: List[str], status_message: str) -> Dict[str, Any]:
         }
 
     @staticmethod
-    def search_results_response(results: List[Dict[str, Any]]) -> Dict[str, Any]:
+    def search_results_response(
+        results: List[Dict[str, Any]],
+        pagination: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
         """
         Format search results response.
         
@@ -142,9 +145,14 @@ def search_results_response(results: List[Dict[str, Any]]) -> Dict[str, Any]:
         Returns:
             Formatted search results response
         """
-        return {
+        response = {
             "results": results
         }
+
+        if pagination is not None:
+            response["pagination"] = pagination
+
+        return response
 
     @staticmethod
     def config_response(config_data: Dict[str, Any]) -> str:
@@ -361,4 +369,4 @@ def settings_info_response(
         if message:
             response["message"] = message
 
-        return response
+        return response

src/code_index_mcp/utils/validation.py

@@ -158,6 +158,35 @@ def validate_search_pattern(pattern: str, regex: bool = False) -> Optional[str]:
 
         return None
 
+    @staticmethod
+    def validate_pagination(start_index: int, max_results: Optional[int]) -> Optional[str]:
+        """
+        Validate pagination parameters for search queries.
+
+        Args:
+            start_index: The index of the first result to include.
+            max_results: The maximum number of results to return.
+
+        Returns:
+            Error message if validation fails, None if valid.
+        """
+        if not isinstance(start_index, int):
+            return "start_index must be an integer"
+
+        if start_index < 0:
+            return "start_index cannot be negative"
+
+        if max_results is None:
+            return None
+
+        if not isinstance(max_results, int):
+            return "max_results must be an integer when provided"
+
+        if max_results <= 0:
+            return "max_results must be greater than zero when provided"
+
+        return None
+
     @staticmethod
     def validate_file_extensions(extensions: List[str]) -> Optional[str]:
         """
@@ -204,4 +233,4 @@ def sanitize_file_path(file_path: str) -> str:
         # Remove any leading slashes to ensure relative path
         sanitized = sanitized.lstrip('/')
 
-        return sanitized
+        return sanitized

tests/search/test_search_service_pagination.py

@@ -0,0 +1,98 @@
+"""Tests for search result pagination formatting."""
+from pathlib import Path as _TestPath
+from types import SimpleNamespace
+import sys
+
+ROOT = _TestPath(__file__).resolve().parents[2]
+SRC_PATH = ROOT / 'src'
+if str(SRC_PATH) not in sys.path:
+    sys.path.insert(0, str(SRC_PATH))
+
+from code_index_mcp.services.search_service import SearchService
+
+
+def _create_service() -> SearchService:
+    ctx = SimpleNamespace(
+        request_context=SimpleNamespace(
+            lifespan_context=SimpleNamespace(base_path="", settings=None)
+        )
+    )
+    return SearchService(ctx)
+
+
+def test_paginate_results_default_ordering():
+    service = _create_service()
+
+    raw_results = {
+        "b/file.py": [(12, "second match"), (3, "first match")],
+        "a/file.py": [(8, "another file")],
+    }
+
+    formatted, pagination = service._paginate_results(
+        raw_results,
+        start_index=0,
+        max_results=None,
+    )
+
+    assert pagination == {
+        "total_matches": 3,
+        "returned": 3,
+        "start_index": 0,
+        "has_more": False,
+        "end_index": 3,
+    }
+
+    assert formatted == [
+        {"file": "a/file.py", "line": 8, "text": "another file"},
+        {"file": "b/file.py", "line": 3, "text": "first match"},
+        {"file": "b/file.py", "line": 12, "text": "second match"},
+    ]
+
+
+def test_paginate_results_with_start_and_limit():
+    service = _create_service()
+
+    raw_results = {
+        "b/file.py": [(5, "line five"), (6, "line six")],
+        "a/file.py": [(1, "line one"), (2, "line two")],
+    }
+
+    formatted, pagination = service._paginate_results(
+        raw_results,
+        start_index=1,
+        max_results=2,
+    )
+
+    assert pagination == {
+        "total_matches": 4,
+        "returned": 2,
+        "start_index": 1,
+        "has_more": True,
+        "max_results": 2,
+        "end_index": 3,
+    }
+
+    assert formatted == [
+        {"file": "a/file.py", "line": 2, "text": "line two"},
+        {"file": "b/file.py", "line": 5, "text": "line five"},
+    ]
+
+
+def test_paginate_results_when_start_beyond_total():
+    service = _create_service()
+
+    formatted, pagination = service._paginate_results(
+        {"only/file.py": [(1, "match")]},
+        start_index=10,
+        max_results=5,
+    )
+
+    assert formatted == []
+    assert pagination == {
+        "total_matches": 1,
+        "returned": 0,
+        "start_index": 1,
+        "has_more": False,
+        "max_results": 5,
+        "end_index": 1,
+    }

tests/utils/test_validation_pagination.py

@@ -0,0 +1,22 @@
+"""Tests for pagination validation helper."""
+from pathlib import Path as _TestPath
+import sys
+
+ROOT = _TestPath(__file__).resolve().parents[2]
+SRC_PATH = ROOT / 'src'
+if str(SRC_PATH) not in sys.path:
+    sys.path.insert(0, str(SRC_PATH))
+
+from code_index_mcp.utils.validation import ValidationHelper
+
+
+def test_validate_pagination_accepts_valid_values():
+    assert ValidationHelper.validate_pagination(0, None) is None
+    assert ValidationHelper.validate_pagination(5, 10) is None
+
+
+def test_validate_pagination_rejects_invalid_values():
+    assert ValidationHelper.validate_pagination(-1, None) == "start_index cannot be negative"
+    assert ValidationHelper.validate_pagination(0, 0) == "max_results must be greater than zero when provided"
+    assert ValidationHelper.validate_pagination(0, "a") == "max_results must be an integer when provided"
+    assert ValidationHelper.validate_pagination("a", None) == "start_index must be an integer"