From 7b632f2ca43138c4c102e1222bccb62bf434f13b Mon Sep 17 00:00:00 2001 From: sebasidian sunday Date: Sat, 3 Oct 2026 19:42:38 -0400 Subject: [PATCH] fix(mcp): rank symbol searches and report truncated pages --- README.md | 9 ++ codeprism/core/storage.py | 60 ++++++-- codeprism/mcp/server.py | 18 ++- codeprism/mcp/tools.py | 10 +- codeprism/query/engine.py | 49 +++++- docs/changelog.md | 4 + tests/test_search_pagination.py | 256 ++++++++++++++++++++++++++++++++ 7 files changed, 376 insertions(+), 30 deletions(-) create mode 100644 tests/test_search_pagination.py diff --git a/README.md b/README.md index cf9262c..76e0be8 100644 --- a/README.md +++ b/README.md @@ -277,6 +277,15 @@ codeprism serve --no-auto-index codeprism serve /path/to/project --transport sse --port 8765 ``` +The MCP `search_symbol` tool ranks substring matches by exact name, then prefix, then other +matches, with stable name/file ordering. Its `limit` defaults to 50 and accepts 1-500. +`kind` and `project_path` filters apply before this limit; the response includes `count` +(returned matches), `total` (all filtered matches) and `truncated` (whether any were cut). +For example, `search_symbol("run", kind="function", limit=100)` can return more matches; +narrow the query or filters when `truncated` is true. With semantic embeddings enabled, +search keeps its existing 20-neighbor cap and ignores `limit`; that response has +`count`/`matches` only because a complete `total` and `truncated` are unavailable. + The SSE server listens on `127.0.0.1` only and has no built-in authentication; there is no `--host` option yet ([#45](https://github.com/knight22-21/CodePrism/issues/45)). To reach it from another machine, use an SSH tunnel or a reverse proxy on the same host (see [INTEGRATIONS.md](INTEGRATIONS.md)). diff --git a/codeprism/core/storage.py b/codeprism/core/storage.py index 2aa6c38..60806db 100644 --- a/codeprism/core/storage.py +++ b/codeprism/core/storage.py @@ -26,6 +26,9 @@ # 2: Python calls record call style/receiver; function-local imports are indexed INDEX_FORMAT_VERSION = 2 +DEFAULT_SEARCH_LIMIT = 50 +MAX_SEARCH_LIMIT = 500 + _SCHEMA = """ PRAGMA journal_mode=WAL; PRAGMA foreign_keys=ON; @@ -331,20 +334,49 @@ async def find_symbols(self, name: str, kind: str | None = None) -> list[SymbolR return [_row_to_symbol(r) for r in rows] async def search_symbols(self, query: str, kind: str | None = None) -> list[SymbolRecord]: - pattern = f"%{query}%" - kind_param = kind.lower() if kind else None - if kind_param: - async with self.db.execute( - "SELECT * FROM symbols WHERE name LIKE ? AND kind = ? LIMIT 50", - (pattern, kind_param), - ) as cur: - rows = await cur.fetchall() - else: - async with self.db.execute( - "SELECT * FROM symbols WHERE name LIKE ? LIMIT 50", (pattern,) - ) as cur: - rows = await cur.fetchall() - return [_row_to_symbol(r) for r in rows] + """Return the default ranked substring page, preserving the list API.""" + symbols, _ = await self.search_symbols_page(query, kind) + return symbols + + async def search_symbols_page( + self, + query: str, + kind: str | None = None, + *, + limit: int = DEFAULT_SEARCH_LIMIT, + file_prefix: str | None = None, + ) -> tuple[list[SymbolRecord], int]: + """Return ranked substring matches and the total before applying the limit.""" + if ( + isinstance(limit, bool) + or not isinstance(limit, int) + or not 1 <= limit <= MAX_SEARCH_LIMIT + ): + raise ValueError(f"limit must be an integer between 1 and {MAX_SEARCH_LIMIT}") + predicates = ["s.name LIKE ?"] + parameters: list[str | int] = [f"%{query}%"] + if kind: + predicates.append("s.kind = ?") + parameters.append(kind.lower()) + if file_prefix is not None: + # Match the existing literal startswith filter, including '%' and '_' + # in a directory name, before pagination and counting. + predicates.append("substr(f.path, 1, length(?)) = ?") + parameters.extend([file_prefix, file_prefix]) + sql = f""" + SELECT s.*, COUNT(*) OVER () AS total_matches + FROM symbols AS s JOIN files AS f ON f.id = s.file_id + WHERE {" AND ".join(predicates)} + ORDER BY s.name = ? DESC, s.name LIKE ? DESC, s.name, f.path, s.id + LIMIT ? + """ + parameters.extend([query, f"{query}%", limit]) + # A window count keeps the page and total in the same SQLite snapshot, + # even if a watcher updates the index while the query runs. + async with self.db.execute(sql, parameters) as cur: + rows = list(await cur.fetchall()) + total = int(rows[0]["total_matches"]) if rows else 0 + return [_row_to_symbol(row) for row in rows], total async def get_all_symbols(self) -> list[SymbolRecord]: async with self.db.execute("SELECT * FROM symbols") as cur: diff --git a/codeprism/mcp/server.py b/codeprism/mcp/server.py index 7c9a229..e90d054 100644 --- a/codeprism/mcp/server.py +++ b/codeprism/mcp/server.py @@ -424,26 +424,32 @@ async def get_data_flow(file: str, symbol: str) -> dict[str, Any]: @mcp.tool( description="Find symbols by name (substring match).\n\n" "project_path: optional directory prefix to restrict results to a sub-project.\n" - "kind: case-insensitive symbol kind; one of " + ", ".join(SYMBOL_KINDS) + "kind: case-insensitive symbol kind; one of " + ", ".join(SYMBOL_KINDS) + ".\n" + "limit: substring result limit, 1-500 (default 50). Substring results rank exact names, " + "then prefixes, then other matches, and include total/truncated.\n" + "Semantic search retains its 20-neighbor cap, ignores limit, and omits total/truncated." ) async def search_symbol( query: str, project_path: str | None = None, kind: str | None = None, + limit: int = 50, ) -> dict[str, Any]: """Find symbols by name (substring match). project_path: optional directory prefix to restrict results to a sub-project. kind: case-insensitive symbol kind from SYMBOL_KINDS. + limit: 1-500 (default 50), applied after filtering in substring mode. + Substring responses include total and truncated. Semantic search retains + its 20-neighbor cap and ignores limit; its total/truncated are unavailable. """ try: - matches = await _get().search_symbols(query, kind) + page = await _get().search_symbols_page( + query, kind, limit=limit, file_prefix=_abs(project_path) if project_path else None + ) except ValueError as error: return {"error": str(error)} - if project_path: - prefix = _abs(project_path) - matches = [m for m in matches if m.file_path.startswith(prefix)] - return search_matches_to_dict(matches) + return search_matches_to_dict(page.matches, total=page.total) @mcp.tool() diff --git a/codeprism/mcp/tools.py b/codeprism/mcp/tools.py index 56772bc..a7e3b4a 100644 --- a/codeprism/mcp/tools.py +++ b/codeprism/mcp/tools.py @@ -94,8 +94,11 @@ def file_map_to_dict(r: FileMap) -> dict[str, Any]: } -def search_matches_to_dict(matches: list[SearchMatch]) -> dict[str, Any]: - return { +def search_matches_to_dict( + matches: list[SearchMatch], *, total: int | None = None +) -> dict[str, Any]: + """Serialize matches, adding completeness metadata only when the total is known.""" + result = { "count": len(matches), "matches": [ { @@ -110,6 +113,9 @@ def search_matches_to_dict(matches: list[SearchMatch]) -> dict[str, Any]: for m in matches ], } + if total is not None: + result.update(total=total, truncated=total > len(matches)) + return result def deps_to_dict(r: DependencyResult) -> dict[str, Any]: diff --git a/codeprism/query/engine.py b/codeprism/query/engine.py index 7fdfbbc..4157e7c 100644 --- a/codeprism/query/engine.py +++ b/codeprism/query/engine.py @@ -8,7 +8,7 @@ from ..core.graph import GraphEngine from ..core.models import SYMBOL_KINDS, FileRecord, NodeKind, SymbolRecord -from ..core.storage import StorageManager +from ..core.storage import DEFAULT_SEARCH_LIMIT, MAX_SEARCH_LIMIT, StorageManager from . import context as _ctx_mod from . import impact as _imp_mod from . import summary as _sum_mod @@ -27,6 +27,14 @@ class SearchMatch: docstring_excerpt: str | None = None +@dataclass +class SearchPage: + """Search matches with a complete substring total, or None for semantic search.""" + + matches: list[SearchMatch] + total: int | None = None + + @dataclass class FileMapEntry: path: str @@ -128,21 +136,45 @@ async def get_callees(self, file_path: str, symbol_name: str) -> list[SymbolReco # ── Search ──────────────────────────────────────────────────────────────── async def search_symbols(self, query: str, kind: str | None = None) -> list[SearchMatch]: - """Search symbols, rejecting unknown kinds before querying either backend.""" + """Return the default search page while preserving the existing list API.""" + return (await self.search_symbols_page(query, kind)).matches + + async def search_symbols_page( + self, + query: str, + kind: str | None = None, + *, + limit: int = DEFAULT_SEARCH_LIMIT, + file_prefix: str | None = None, + ) -> SearchPage: + """Search with filtered substring pagination; semantic results keep their cap. + + file_prefix is an already-resolved literal path prefix. limit and total + apply only to substring mode; the semantic backend still requests 20 + neighbors and supplies no complete total. + """ + if ( + isinstance(limit, bool) + or not isinstance(limit, int) + or not 1 <= limit <= MAX_SEARCH_LIMIT + ): + raise ValueError(f"limit must be an integer between 1 and {MAX_SEARCH_LIMIT}") if kind is not None: normalized = kind.lower() if normalized not in SYMBOL_KINDS: raise ValueError(f"Unknown kind {kind!r}. Valid kinds: {', '.join(SYMBOL_KINDS)}") kind = normalized if self._embedder is not None and self._embed_store is not None: - return await self._semantic_search(query, kind) - return await self._substring_search(query, kind) - - async def _substring_search(self, query: str, kind: str | None = None) -> list[SearchMatch]: - raw = await self._storage.search_symbols(query, kind) + matches = await self._semantic_search(query, kind) + if file_prefix is not None: + matches = [match for match in matches if match.file_path.startswith(file_prefix)] + return SearchPage(matches) + raw, total = await self._storage.search_symbols_page( + query, kind, limit=limit, file_prefix=file_prefix + ) all_files = await self._storage.get_all_files() id_to_path = {f.id: f.path for f in all_files} - return [ + matches = [ SearchMatch( symbol=sym, file_path=id_to_path.get(sym.file_id, ""), @@ -151,6 +183,7 @@ async def _substring_search(self, query: str, kind: str | None = None) -> list[S ) for sym in raw ] + return SearchPage(matches, total) async def _semantic_search(self, query: str, kind: str | None = None) -> list[SearchMatch]: import asyncio as _asyncio diff --git a/docs/changelog.md b/docs/changelog.md index 206f452..7b612b9 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -8,6 +8,10 @@ Versioning follows [Semantic Versioning](https://semver.org/). ## [Unreleased] ### Fixed +- Substring symbol search ranks exact names and prefixes before limiting results. MCP + `search_symbol` now accepts `limit` (1-500, default 50) and reports filtered `total` and + `truncated`; kind/project filters apply before pagination. Semantic search keeps its + existing cap and response shape. (#38) - `index --languages` now rejects unknown names with a list of supported languages. Names remain case-insensitive, and supported aliases such as `py` and `C++` still work. (#55) - `codeprism --version` now exists, and `codeprism.__version__` is read from the installed diff --git a/tests/test_search_pagination.py b/tests/test_search_pagination.py new file mode 100644 index 0000000..3b7ba72 --- /dev/null +++ b/tests/test_search_pagination.py @@ -0,0 +1,256 @@ +"""Ranked substring search and truthful MCP pagination (issue #38).""" + +from pathlib import Path +from typing import cast + +import pytest + +import codeprism.mcp.server as server +from codeprism.core.config import CodePrismConfig +from codeprism.core.graph import GraphEngine +from codeprism.core.models import FileRecord, NodeKind, SymbolRecord +from codeprism.core.storage import StorageManager +from codeprism.embeddings.store import SearchResult +from codeprism.indexer.project_indexer import ProjectIndexer +from codeprism.mcp.tools import search_matches_to_dict +from codeprism.query.engine import QueryEngine + + +async def _index_sources( + root: Path, graph: GraphEngine, storage: StorageManager, sources: dict[str, str] +) -> None: + """Index real fixture files without inheriting the surrounding Git checkout.""" + for name, source in sources.items(): + path = root / name + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(source) + config = CodePrismConfig(languages=["python"], respect_gitignore=False) + await ProjectIndexer(graph, storage, config).index(str(root)) + + +@pytest.fixture +async def search_index( + storage: StorageManager, graph: GraphEngine, tmp_path: Path +) -> tuple[QueryEngine, Path]: + """Index 60 partial matches before the exact match in real SQLite storage.""" + root = tmp_path / "project" + await _index_sources( + root, + graph, + storage, + { + "a_helpers.py": "".join( + f"def prerun_step_{i}():\n return {i}\n\n" for i in range(60) + ), + "z_core.py": "def run():\n return 1\n", + }, + ) + return QueryEngine(graph, storage), root.resolve() + + +@pytest.fixture +def mcp_search(search_index: tuple[QueryEngine, Path], monkeypatch: pytest.MonkeyPatch) -> None: + """Inject the real indexed query engine and served root into the MCP tool.""" + engine, root = search_index + monkeypatch.setattr(server, "_engine", engine) + monkeypatch.setattr(server, "_project_path", str(root)) + + +async def test_existing_list_api_keeps_exact_match_before_limit( + search_index: tuple[QueryEngine, Path], storage: StorageManager +) -> None: + """The original reproduction must succeed through both existing list APIs.""" + engine, _ = search_index + matches = await engine.search_symbols("run") + assert isinstance(matches, list) + assert len(matches) == 50 + assert matches[0].symbol.name == "run" + symbols = await storage.search_symbols("run") + assert isinstance(symbols, list) + assert len(symbols) == 50 + assert symbols[0].name == "run" + + +async def test_existing_list_api_ranks_prefix_before_substrings( + storage: StorageManager, graph: GraphEngine, tmp_path: Path +) -> None: + """SQL ranking precedes the cap and uses stable names within each tier.""" + await _index_sources( + tmp_path / "rank", + graph, + storage, + { + "rank.py": "\n".join( + f"def {name}(): pass" for name in ["prerun_z", "run_z", "prerun_a", "run_a", "run"] + ) + }, + ) + names = [match.symbol.name for match in await QueryEngine(graph, storage).search_symbols("run")] + assert names == ["run", "run_a", "run_z", "prerun_a", "prerun_z"] + + +async def test_existing_list_api_orders_tied_names_by_file( + storage: StorageManager, graph: GraphEngine +) -> None: + """Equal names must not depend on SQLite insertion order.""" + for path in ["/project/z.py", "/project/a.py"]: + file = FileRecord.create(path) + await storage.upsert_file(file) + await storage.upsert_symbol(SymbolRecord.create(path, file.id, "run", NodeKind.FUNCTION)) + engine = QueryEngine(graph, storage) + assert [match.file_path for match in await engine.search_symbols("run")] == [ + "/project/a.py", + "/project/z.py", + ] + + +@pytest.mark.parametrize("count,limit", [(0, 5), (4, 5), (5, 5), (6, 5), (50, 50), (51, 50)]) +async def test_substring_page_reports_actual_total_and_truncation( + storage: StorageManager, graph: GraphEngine, tmp_path: Path, count: int, limit: int +) -> None: + """Empty, below-limit, exact-limit and cut pages have accurate metadata.""" + await _index_sources( + tmp_path / "boundary", + graph, + storage, + {"symbols.py": "".join(f"def item_{i}(): pass\n" for i in range(count))}, + ) + page = await QueryEngine(graph, storage).search_symbols_page("item", limit=limit) + result = search_matches_to_dict(page.matches, total=page.total) + assert result["count"] == min(count, limit) + assert result["total"] == count + assert result["truncated"] is (count > limit) + + +@pytest.mark.parametrize("limit", [1, 50, 500]) +async def test_mcp_limit_and_default_ranked_page(mcp_search: None, limit: int) -> None: + """MCP exposes the default and both valid bound values without losing run.""" + result = await server.search_symbol("run", limit=limit) + assert result["count"] == min(61, limit) + assert result["total"] == 61 + assert result["truncated"] is (limit < 61) + assert result["matches"][0]["name"] == "run" + if limit == 50: + assert await server.search_symbol("run") == result + + +async def test_mcp_filters_before_pagination_and_total( + search_index: tuple[QueryEngine, Path], mcp_search: None +) -> None: + """Relative project and kind filters select/count rows before the cap.""" + engine, root = search_index + await _index_sources( + root, + engine._graph, + engine._storage, + { + "sub%_project/target.py": "def run(): pass\ndef run_again(): pass\nclass run_type: pass\n", + "subXXproject/other.py": "def run_other(): pass\n", + }, + ) + result = await server.search_symbol("run", "sub%_project", "FUNCTION", limit=1) + assert result["count"] == 1 + assert result["total"] == 2 + assert result["truncated"] is True + assert result["matches"][0]["name"] == "run" + assert result["matches"][0]["file_path"] == str(root / "sub%_project/target.py") + assert ( + await server.search_symbol("run", str(root / "sub%_project"), "FUNCTION", limit=1) == result + ) + uncut = await server.search_symbol("run", "sub%_project", "FUNCTION", limit=2) + assert uncut["total"] == 2 + assert uncut["truncated"] is False + empty = await server.search_symbol("run", "absent", limit=1) + assert empty == {"count": 0, "matches": [], "total": 0, "truncated": False} + + +@pytest.mark.parametrize("kind", ["function", "Function", "FUNCTION"]) +async def test_page_kind_filter_counts_only_selected_kind( + search_index: tuple[QueryEngine, Path], kind: str +) -> None: + """Case-insensitive kinds retain the complete filtered total.""" + engine, root = search_index + await _index_sources(root, engine._graph, engine._storage, {"types.py": "class run: pass\n"}) + page = await engine.search_symbols_page("run", kind, limit=1) + assert page.total == 61 + assert len(page.matches) == 1 + assert page.matches[0].symbol.kind == NodeKind.FUNCTION + + +@pytest.mark.parametrize("limit", [0, -1, 501, True, 1.5, "5", None]) +async def test_invalid_limits_return_mcp_error_and_query_value_error( + search_index: tuple[QueryEngine, Path], mcp_search: None, limit: object +) -> None: + """Out-of-range input is rejected instead of silently clamped.""" + engine, _ = search_index + with pytest.raises(ValueError, match="limit"): + await engine.search_symbols_page("run", limit=cast(int, limit)) + with pytest.raises(ValueError, match="limit"): + await engine._storage.search_symbols_page("run", limit=cast(int, limit)) + result = await server.search_symbol("run", limit=cast(int, limit)) + assert set(result) == {"error"} + assert "1" in result["error"] and "500" in result["error"] + + +async def test_invalid_kind_keeps_mcp_error_contract(mcp_search: None) -> None: + """Unknown kinds still explain accepted choices through the MCP wrapper.""" + result = await server.search_symbol("run", kind="method") + assert "Unknown kind 'method'" in result["error"] + + +async def test_serializer_without_total_keeps_original_shape( + search_index: tuple[QueryEngine, Path], +) -> None: + """Existing serializer callers retain count/matches without invented metadata.""" + engine, _ = search_index + matches = await engine.search_symbols("run") + result = search_matches_to_dict(matches) + assert set(result) == {"count", "matches"} + assert result["count"] == len(matches) + + +class _Embedder: + """Lightweight semantic encoder; no optional model is downloaded.""" + + def encode_one(self, query: str) -> list[float]: + """Return a deterministic vector for the semantic compatibility test.""" + return [1.0] + + +class _Store: + """Return actual indexed symbol IDs in a fixed semantic order.""" + + def __init__(self, results: list[SearchResult]) -> None: + self.results = results + self.requested_top_k: int | None = None + + def search(self, vector: list[float], top_k: int) -> list[SearchResult]: + """Mimic the existing vector store's capped nearest-neighbor API.""" + self.requested_top_k = top_k + return self.results[:top_k] + + +async def test_semantic_backend_keeps_cap_order_and_original_metadata_shape( + search_index: tuple[QueryEngine, Path], mcp_search: None +) -> None: + """Substring limit/total must not pretend to describe the semantic result set.""" + engine, root = search_index + symbols = await engine._storage.get_all_symbols() + results = [ + SearchResult(symbol.id, str(root / "a_helpers.py"), symbol.name, 0.1, {"kind": "function"}) + for symbol in symbols[:30] + ] + store = _Store(results) + engine.set_embeddings(_Embedder(), store) + page = await engine.search_symbols_page("run", limit=1) + assert store.requested_top_k == 20 + assert page.total is None + assert len(page.matches) == 20 + assert [match.symbol.id for match in page.matches] == [ + result.symbol_id for result in results[:20] + ] + assert page.matches == await engine.search_symbols("run") + response = await server.search_symbol("run", kind="FUNCTION", limit=1) + assert set(response) == {"count", "matches"} + assert response["count"] == 20 + assert await server.search_symbol("run", "absent", limit=1) == {"count": 0, "matches": []}